Batched-decode engine runtime (attention-mask path) + benchmark

swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift

@@ -93,6 +93,9 @@ public struct LanguageConfig: Codable, Sendable, Equatable {
     public let vocabSize: Int
     public let maxContextLength: Int
 
+    /// Largest batch the graph accepts. 1 for single-sequence graphs and for bundles without `max_batch_size`.
+    package let maxBatchSize: Int
+
     /// `true` if the bundle ships its own tokenizer directory; `false` to
     /// load via HuggingFace at runtime. Defaults to `true` when omitted.
     public let embeddedTokenizer: Bool
@@ -140,6 +143,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable {
         self.prefillChunkSize = prefillChunkSize
         self.prefillChunkThreshold = prefillChunkThreshold
         self.overrides = overrides
+        self.maxBatchSize = 1
     }
 
     enum CodingKeys: String, CodingKey {
@@ -153,6 +157,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable {
         case prefillChunkSize = "prefill_chunk_size"
         case prefillChunkThreshold = "prefill_chunk_threshold"
         case overrides
+        case maxBatchSize = "max_batch_size"
     }
 
     public init(from decoder: Swift.Decoder) throws {
@@ -167,6 +172,7 @@ public struct LanguageConfig: Codable, Sendable, Equatable {
         self.prefillChunkSize = try c.decodeIfPresent(Int.self, forKey: .prefillChunkSize)
         self.prefillChunkThreshold = try c.decodeIfPresent(Int.self, forKey: .prefillChunkThreshold)
         self.overrides = try c.decodeIfPresent(LanguageOverrides.self, forKey: .overrides)
+        self.maxBatchSize = try c.decodeIfPresent(Int.self, forKey: .maxBatchSize) ?? 1
     }
 
     // MARK: - Additional Stop Tokens

swift/Sources/CoreAILanguageModels/Bundle/LanguageModelBundle.swift

@@ -60,6 +60,8 @@ public struct LanguageModelBundle: Sendable {
     public var tokenizer: String { language.tokenizer }
     public var vocabSize: Int { language.vocabSize }
     public var maxContextLength: Int { language.maxContextLength }
+    /// Largest batch the graph accepts. 1 for single-sequence graphs.
+    package var maxBatchSize: Int { language.maxBatchSize }
 
     /// Model-specific runtime settings (`language.overrides`); nil for most models.
     public var overrides: LanguageOverrides? { language.overrides }

swift/Sources/CoreAILanguageModels/Handlers/BatchedInputBuilder.swift

@@ -0,0 +1,96 @@
+// 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
+
+/// Model-free input builders for an N-row batched step. Each returns a flat row-major buffer that
+/// `decodeBatch` copies into its tensor with `fillNDArray`.
+///
+/// The uniform builders take one `startPos` for all rows, and each of their rows matches the
+/// single-row inputs exactly. The `raggedDecode*` builders take a cursor per row and match the
+/// uniform builders when all cursors are equal.
+enum BatchedInputBuilder {
+    /// Flat `[N, queryLen]` token ids. Every row must have `queryLen` tokens.
+    static func tokenIDs(rows: [[Int32]]) -> [Int32] {
+        rows.flatMap { $0 }
+    }
+
+    /// Flat `[N, keyLen]` position ids with `keyLen = startPos + queryLen`. Each row is `0..<keyLen`,
+    /// covering the whole key axis (cached prefix plus new tokens).
+    static func positionIDs(rowCount: Int, startPos: Int, queryLen: Int) -> [Int32] {
+        let keyLen = startPos + queryLen
+        var out = [Int32](repeating: 0, count: rowCount * keyLen)
+        for b in 0..<rowCount {
+            for c in 0..<keyLen {
+                out[b * keyLen + c] = Int32(c)
+            }
+        }
+        return out
+    }
+
+    /// Flat `[N, 1, queryLen, keyLen]` additive causal mask with `keyLen = startPos + queryLen`.
+    /// Query `i` gets 0 for keys `0...startPos + i` and `sentinel` after that.
+    static func additiveCausalMask(
+        rowCount: Int, startPos: Int, queryLen: Int,
+        sentinel: Float16 = causalMaskSentinel
+    ) -> [Float16] {
+        let keyLen = startPos + queryLen
+        var out = [Float16](repeating: sentinel, count: rowCount * queryLen * keyLen)
+        for b in 0..<rowCount {
+            for query in 0..<queryLen {
+                let attendUpTo = startPos + query  // absolute position of this query
+                let base = (b * queryLen + query) * keyLen
+                for key in 0...attendUpTo { out[base + key] = 0 }
+            }
+        }
+        return out
+    }
+
+    // MARK: - Ragged decode (per-row cursor)
+
+    /// Position ids for a ragged decode step (`queryLen = 1`, one cursor per row). Flat `[N, keyLen]`
+    /// with `keyLen = max(cursors) + 1`. Every row writes its new token at slot `max(cursors)`, and
+    /// that slot holds the row's own cursor, which the model reads as the query's RoPE position.
+    ///
+    /// For cursors `[2, 5]`:
+    /// ```
+    /// row 0: [0, 1, 0, 0, 0, 2]   // slots 0..<2 cached, gap 2..<5 set to 0, slot 5 holds cursor 2
+    /// row 1: [0, 1, 2, 3, 4, 5]
+    /// ```
+    static func raggedDecodePositionIDs(cursors: [Int]) -> [Int32] {
+        let maxCursor = cursors.max() ?? 0
+        let keyLen = maxCursor + 1
+        var out = [Int32](repeating: 0, count: cursors.count * keyLen)
+        for (b, cursor) in cursors.enumerated() {
+            let base = b * keyLen
+            for key in 0..<cursor { out[base + key] = Int32(key) }
+            out[base + maxCursor] = Int32(cursor)
+        }
+        return out
+    }
+
+    /// Additive mask for a ragged decode step, flat `[N, 1, 1, keyLen]` with
+    /// `keyLen = max(cursors) + 1`. Each row attends its cached slots and the write slot
+    /// `max(cursors)`, and masks the gap between them.
+    ///
+    /// For cursors `[2, 5]` (`s` = `sentinel`):
+    /// ```
+    /// row 0: [0, 0, s, s, s, 0]
+    /// row 1: [0, 0, 0, 0, 0, 0]
+    /// ```
+    static func raggedDecodeMask(
+        cursors: [Int], sentinel: Float16 = causalMaskSentinel
+    ) -> [Float16] {
+        let maxCursor = cursors.max() ?? 0
+        let keyLen = maxCursor + 1
+        var out = [Float16](repeating: sentinel, count: cursors.count * keyLen)
+        for (b, cursor) in cursors.enumerated() {
+            let base = b * keyLen
+            for key in 0..<cursor { out[base + key] = 0 }
+            out[base + maxCursor] = 0
+        }
+        return out
+    }
+}

swift/Sources/CoreAILanguageModels/Handlers/InputHandler+AttentionMask.swift

@@ -0,0 +1,37 @@
+// Additive causal attention mask for dynamic (GPU) engines.
+//
+// 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 CoreAI
+import CoreAIShared
+
+/// Fills a `Float16` `[1, 1, queryLen, keyLen]` additive causal mask in place, reading both lengths
+/// from the array's shape. Query `i` gets 0 for keys `0...processedTokenCount + i` and
+/// ``causalMaskSentinel`` after that, which matches the implicit `is_causal` path.
+func fillAdditiveCausalMask(_ mask: inout NDArray, processedTokenCount: Int) {
+    mask.mutableView(as: Float16.self).withUnsafeMutablePointer { ptr, shape, strides in
+        // shape == [1, 1, query_len, key_len]
+        let queryLen = shape[2]
+        let keyLen = shape[3]
+        for query in 0..<queryLen {
+            let attendUpTo = processedTokenCount + query  // absolute position of this query
+            for key in 0..<keyLen {
+                let offset = query &* strides[2] &+ key &* strides[3]
+                ptr[offset] = key <= attendUpTo ? 0 : causalMaskSentinel
+            }
+        }
+    }
+}
+
+/// Allocates and fills a new single-row causal mask for one dynamic-engine step, with
+/// `queryLen = tokens.count` and `keyLen = processedTokenCount + queryLen`.
+func makeAdditiveCausalMask(context: InputContext, descriptor: NDArrayDescriptor) -> NDArray {
+    let queryLen = context.tokens.count
+    let keyLen = context.processedTokenCount + queryLen
+    var mask = NDArray(descriptor: descriptor.resolvingDynamicDimensions([1, 1, queryLen, keyLen]))
+    fillAdditiveCausalMask(&mask, processedTokenCount: context.processedTokenCount)
+    return mask
+}

swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift

@@ -37,7 +37,10 @@ public struct TokenInputHandler: SyncInputHandler {
         self.useCompactPositionIds = useCompactPositionIds
         self.inputNames = [inputIdsName, positionIdsName]
 
-        let initDesc = inputIdsDescriptor.resolvingDynamicDimensions([1, 1])
+        // Keep a static batch dim as is, because `resolvingDynamicDimensions` rejects a mismatch on
+        // a pinned dim. A dynamic batch dim resolves to 1.
+        let batch = (inputIdsDescriptor.shape.first ?? 1) > 0 ? inputIdsDescriptor.shape[0] : 1
+        let initDesc = inputIdsDescriptor.resolvingDynamicDimensions([batch, 1])
         self.inputIdsArray = NDArray(descriptor: initDesc)
         self.cachedBatchSize = 1
     }

swift/Sources/CoreAILanguageModels/Handlers/InputLayout.swift

@@ -11,6 +11,9 @@ public struct InputLayout: Sendable {
     public let positionIdsName: String
     public let logitsName: String
 
+    /// Name of the additive attention-mask input, or `nil` when the graph applies its own causal mask.
+    public let attnMaskName: String?
+
     public let queryPolicy: QueryPolicy
     public let logitsPolicy: LogitsPolicy
     public let positionPolicy: PositionPolicy
@@ -111,6 +114,8 @@ extension InputLayout {
             from: inputNames, candidates: knownPositionIdNames, label: "position_ids")
         let logitsName = outputNames[0]
 
+        let attnMaskName = knownAttnMaskNames.first(where: inputNames.contains)
+
         let queryPolicy: QueryPolicy =
             if let s = inputIdsSeqLen, s > 0 { .fixed(s) } else { .dynamic }
 
@@ -132,6 +137,7 @@ extension InputLayout {
             inputIdsName: inputIdsName,
             positionIdsName: positionIdsName,
             logitsName: logitsName,
+            attnMaskName: attnMaskName,
             queryPolicy: queryPolicy,
             logitsPolicy: logitsPolicy,
             positionPolicy: positionPolicy,
@@ -140,6 +146,7 @@ extension InputLayout {
 
     static let knownInputIdNames = ["input_ids", "in_new_token_ids"]
     static let knownPositionIdNames = ["position_ids", "pos_ids"]
+    static let knownAttnMaskNames = ["attn_mask"]
 
     static func resolveRequired(
         from names: [String], candidates: [String], label: String