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.

- PipelinedConstrainedDecodingStrategy: AsyncSequence-based strategy that
  drives the GPU constrained loop. Single-use iteration guard, proper
  error propagation, Task.checkCancellation, early-stop on consumer drop.

- Session cache (checkout/checkin): ConstrainedGenerationSession reused
  across calls with the same schema. Engine Task defer block owns the
  return-to-cache, ensuring the session 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: exercises session cache
  reuse, schema invalidation, stop sequences, error propagation, and
  maxTokens boundary without requiring Metal hardware.

Author: stikves

swift/Sources/CoreAILanguageModels/DecodingStrategies/PipelinedConstrainedDecodingStrategy.swift

@@ -0,0 +1,232 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import CoreAIShared
+import Foundation
+import Synchronization
+import Tokenizers
+
+private let defaultMaxConstrainedTokens = 512
+
+private final class AtomicFlag: Sendable {
+    private let value = Atomic<Bool>(false)
+    func testAndSet() -> Bool {
+        value.exchange(true, ordering: .acquiring)
+    }
+}
+
+/// Grammar-constrained decoding strategy using the GPU pipelined engine.
+///
+/// Unlike `ConstrainedDecodingStrategy` which forces the sequential engine path
+/// (CPU logits -> mask -> CPU sample), this strategy applies the xgrammar bitmask
+/// directly inside the MPSGraph sampler on GPU -- eliminating the ~296KB logit
+/// transfer and CPU-side softmax/sampling per token.
+///
+/// Only used when the engine conforms to `ConstrainedGenerationCapable`. The routing logic in
+/// `CoreAILanguageModel` selects this strategy based on engine type.
+public struct PipelinedConstrainedDecodingStrategy: DecodingStrategy {
+    private let jsonSchema: String
+    private let vocabSizeOverride: Int?
+
+    public init(jsonSchema: String, vocabSize: Int? = nil) {
+        self.jsonSchema = jsonSchema
+        self.vocabSizeOverride = vocabSize
+    }
+
+    // MARK: - DecodingStrategy conformance
+
+    public func decode(
+        from input: Input,
+        tokenizer: any Tokenizer,
+        inferenceEngine: any InferenceEngine,
+        samplingConfiguration: SamplingConfiguration,
+        options: InferenceOptions,
+        stopSequences: StopSequences
+    ) async throws -> PipelinedConstrainedSequence {
+        guard let constrainedEngine = inferenceEngine as? any ConstrainedGenerationCapable else {
+            throw InferenceRuntimeError.invalidArgument(
+                "PipelinedConstrainedDecodingStrategy requires a ConstrainedGenerationCapable engine")
+        }
+
+        let vocabSize = vocabSizeOverride ?? ConstrainedDecodingStrategy.deriveVocabSize(from: tokenizer)
+        guard let vocabSize else {
+            throw InferenceRuntimeError.invalidArgument(
+                "Cannot determine vocabulary size from tokenizer. "
+                    + "Pass vocabSize explicitly via CoreAIRunner or LLMAsset metadata."
+            )
+        }
+
+        let singleTokenStops = stopSequences.sequences.filter { $0.count == 1 }.map { $0[0] }
+        if stopSequences.sequences.contains(where: { $0.count > 1 }) {
+            CLILogger.log(
+                "Warning: Multi-token stop sequences not supported by xgrammar, using single-token stops only",
+                component: "PipelinedConstrained")
+        }
+        let stopTokenIds: [Int32]? = singleTokenStops.isEmpty ? nil : singleTokenStops
+
+        let sessionBox = try constrainedEngine.getOrCreateConstrainedSessionBox(
+            jsonSchema: jsonSchema,
+            tokenizer: tokenizer,
+            vocabSize: vocabSize,
+            stopTokenIds: stopTokenIds
+        )
+
+        let inputTokens = try PromptUtils.maybeApplyTokenizerChatTemplate(input, tokenizer: tokenizer)
+            .map(Int32.init)
+        let maxTokens = options.maxTokens ?? defaultMaxConstrainedTokens
+
+        CLILogger.log(
+            "Starting GPU pipelined constrained decoding (vocabSize=\(vocabSize), maxTokens=\(maxTokens))",
+            component: "PipelinedConstrained")
+
+        return PipelinedConstrainedSequence(
+            sessionBox: sessionBox,
+            inputTokens: inputTokens,
+            maxTokens: maxTokens,
+            tokenizer: tokenizer,
+            engine: constrainedEngine,
+            samplingConfiguration: samplingConfiguration,
+            stopSequences: stopSequences
+        )
+    }
+}
+
+// MARK: - PipelinedConstrainedSequence
+
+/// Single-use async sequence of `GenerationResult` produced by pipelined constrained decoding.
+///
+/// The engine's background Task owns session lifetime: it returns the session to
+/// cache in its `defer` block after generation completes (normal, error, or cancel).
+/// The iterator does not manage session lifecycle.
+public struct PipelinedConstrainedSequence: AsyncSequence {
+    public typealias Element = GenerationResult
+    public typealias Failure = Error
+
+    fileprivate let sessionBox: ConstrainedSessionBox
+    fileprivate let inputTokens: [Int32]
+    fileprivate let maxTokens: Int
+    let tokenizer: any Tokenizer
+    let engine: any ConstrainedGenerationCapable
+    let samplingConfiguration: SamplingConfiguration
+    let stopSequences: StopSequences
+
+    /// Guards against multiple `makeAsyncIterator()` calls sharing the same session box.
+    private let consumed = AtomicFlag()
+
+    public func makeAsyncIterator() -> Iterator {
+        precondition(!consumed.testAndSet(), "PipelinedConstrainedSequence may only be iterated once")
+
+        return Iterator(
+            sessionBox: sessionBox,
+            inputTokens: inputTokens,
+            maxTokens: maxTokens,
+            tokenizer: tokenizer,
+            engine: engine,
+            samplingConfiguration: samplingConfiguration,
+            stopSequences: stopSequences
+        )
+    }
+}
+
+extension PipelinedConstrainedSequence {
+    public final class Iterator: AsyncIteratorProtocol {
+        public typealias Element = GenerationResult
+        public typealias Failure = Error
+
+        private let tokenizer: any Tokenizer
+        private let engine: any ConstrainedGenerationCapable
+        private let samplingConfiguration: SamplingConfiguration
+        private let stopSequences: StopSequences
+        private let sessionBox: ConstrainedSessionBox
+
+        private let inputTokens: [Int32]
+        private let maxTokens: Int
+
+        private var innerIterator: AsyncThrowingStream<Int32, Error>.AsyncIterator?
+        private var generatedTokens: [Int32] = []
+        private var previousDecodedText: String = ""
+        private var recentTokens: [Int32] = []
+        private var finished: Bool = false
+
+        fileprivate init(
+            sessionBox: ConstrainedSessionBox,
+            inputTokens: [Int32],
+            maxTokens: Int,
+            tokenizer: any Tokenizer,
+            engine: any ConstrainedGenerationCapable,
+            samplingConfiguration: SamplingConfiguration,
+            stopSequences: StopSequences
+        ) {
+            self.sessionBox = sessionBox
+            self.inputTokens = inputTokens
+            self.maxTokens = maxTokens
+            self.tokenizer = tokenizer
+            self.engine = engine
+            self.samplingConfiguration = samplingConfiguration
+            self.stopSequences = stopSequences
+        }
+
+        public func next() async throws -> GenerationResult? {
+            if finished { return nil }
+
+            // Lazily start the engine stream on first iteration
+            if innerIterator == nil {
+                let stream = try engine.generateConstrained(
+                    with: inputTokens,
+                    samplingConfiguration: samplingConfiguration,
+                    maxTokens: maxTokens,
+                    session: sessionBox
+                )
+                self.innerIterator = stream.makeAsyncIterator()
+            }
+
+            do {
+                while true {
+                    try Task.checkCancellation()
+
+                    guard var iterator = self.innerIterator else {
+                        finished = true
+                        return nil
+                    }
+
+                    guard let tokenId = try await iterator.next() else {
+                        finished = true
+                        return nil
+                    }
+                    self.innerIterator = iterator
+
+                    // Check stop sequences
+                    recentTokens.append(tokenId)
+                    if recentTokens.count > stopSequences.maxLength {
+                        recentTokens.removeFirst()
+                    }
+                    if stopSequences.matches(recentTokens: recentTokens) {
+                        finished = true
+                        return nil
+                    }
+
+                    // Decode text incrementally
+                    generatedTokens.append(tokenId)
+                    let fullDecode = tokenizer.decode(tokens: generatedTokens.map { Int($0) })
+                    let common = fullDecode.commonPrefix(with: previousDecodedText)
+                    let delta = String(fullDecode.dropFirst(common.count))
+
+                    if delta.unicodeScalars.contains(where: { $0 == "\u{FFFD}" }) {
+                        continue
+                    }
+
+                    previousDecodedText = fullDecode
+
+                    if !delta.isEmpty {
+                        return GenerationResult(text: delta, tokenId: tokenId, rawLogits: nil)
+                    }
+                }
+            } catch {
+                finished = true
+                throw error
+            }
+        }
+    }
+}

swift/Sources/CoreAILanguageModels/GuidedGeneration/ConstrainedGenerationSession.swift

@@ -193,6 +193,26 @@ public struct ConstrainedGenerationSession: ~Copyable {
         matcher.reset()
         allTokensBlocked = false
     }
+
+    /// Fill the bitmask directly into a caller-provided buffer (e.g., a GPU-visible MTLBuffer).
+    ///
+    /// The caller must ensure the pointer has room for at least `(vocabularySize + 31) / 32`
+    /// Int32 words. Returns `false` if the grammar is terminated (buffer untouched).
+    public mutating func fillBitmask(into pointer: UnsafeMutablePointer<Int32>) -> Bool {
+        if isTerminated { return false }
+
+        let hasConstraints = matcher.fillNextTokenBitmask(pointer)
+        if !hasConstraints {
+            allTokensBlocked = true
+            return false
+        }
+        // Check for all-zeros (no tokens allowed)
+        for i in 0..<bitmaskSize {
+            if pointer[i] != 0 { return true }
+        }
+        allTokensBlocked = true
+        return false
+    }
 }
 
 // MARK: - Float16 Masking

swift/Sources/CoreAILanguageModels/InferenceEngines/ConstrainedGenerationCapable.swift

@@ -0,0 +1,43 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import Tokenizers
+
+/// An inference engine that supports GPU-accelerated grammar-constrained generation.
+///
+/// Conforming engines can apply xgrammar bitmasks directly in the GPU sampling path,
+/// eliminating the need to transfer logits to CPU for masking. This protocol is used as
+/// a capability signal for routing (e.g., choosing pipelined vs. sequential constrained
+/// strategies) and for testability (mock conformers in unit tests).
+///
+/// The session lifecycle follows a checkout/checkin pattern:
+/// 1. Call `getOrCreateConstrainedSessionBox` to obtain a session (cached or fresh)
+/// 2. Pass the session box to `generateConstrained` which drives the GPU loop
+/// 3. The engine returns the session to its internal cache after generation completes
+protocol ConstrainedGenerationCapable: InferenceEngine {
+    /// Obtain a constrained generation session, reusing a cached one if the schema matches.
+    ///
+    /// The internal cache slot is emptied on checkout — concurrent calls create independent sessions.
+    func getOrCreateConstrainedSessionBox(
+        jsonSchema: String,
+        tokenizer: any Tokenizer,
+        vocabSize: Int,
+        stopTokenIds: [Int32]?
+    ) throws -> ConstrainedSessionBox
+
+    /// Stream constrained token generation using GPU-side bitmask application.
+    ///
+    /// The engine resets state, prefills the prompt (unconstrained), then enters
+    /// a semi-pipelined loop: inference overlaps bitmask computation, but sampling
+    /// waits per token (grammar state is inherently sequential).
+    ///
+    /// The session is returned to the engine's cache automatically when the Task completes.
+    func generateConstrained(
+        with input: [TokenId],
+        samplingConfiguration: SamplingConfiguration,
+        maxTokens: Int,
+        session: ConstrainedSessionBox
+    ) throws -> AsyncThrowingStream<TokenId, Error>
+}