From 9a307a66953d655448ce82cc0e0c584e430efac7 Mon Sep 17 00:00:00 2001
From: Sukru Tikves <sukrutikves@gmail.com>
Date: Thu, 8 Oct 2026 18:06:22 -0700
Subject: [PATCH] Add batched-decode engine runtime (attention-mask path) +
 benchmark

Add the engine runtime for batched LLM bundles that take an additive attention
mask, with llm-benchmark as the first consumer.

- CoreAISequentialEngine: decodeBatch runs one forward step over N rows, and
  also handles rows at different cache positions during decode when the graph
  has an attn_mask input. lockstepBatchedGenerate generates from equal-length
  prompts together. prefillRow prefills one prompt and copies its KV prefix
  into one row of a batched cache. The engine accepts an optional attn_mask
  input and fills a causal mask for each step.
- Per-request state (in-flight token, prefix-hit count, logits buffer, input
  handler) moves to GenerationSessionState, so separate sessions no longer
  cancel each other or share buffers.
- Handlers: builders for batched token IDs, position IDs and masks; a KV cache
  that holds N rows, with copyRegion and copyRowPrefix for per-row copies; and
  the BatchedDecodeStep protocol. The per-row copies move raw bytes, so 8-bit
  KV caches work as well.
- BatchedSampler samples one token per row with per-row configurations. The
  single-row iterator now goes through it with one row.
- Bundle: LanguageConfig and LanguageModelBundle read max_batch_size, which
  defaults to 1. ModelBundle accepts metadata_version 0.3 as well as 0.2.
- llm-benchmark: new --batch-size and --inference-engine-variant flags.
  Batches larger than 1 need --inference-engine-variant coreai-sequential.

Existing single-sequence bundles run as before. prefillRow, maxBatchCapacity
and the ragged decode path are there for a future scheduler and have no
caller yet.
---
 .../Bundle/LanguageConfig.swift               |   6 +
 .../Bundle/LanguageModelBundle.swift          |   2 +
 .../Handlers/BatchedInputBuilder.swift        |  96 +++++
 .../Handlers/InputHandler+AttentionMask.swift |  37 ++
 .../Handlers/InputHandler+Token.swift         |   5 +-
 .../Handlers/InputLayout.swift                |   7 +
 .../Handlers/StateHandler+NDArray.swift       | 181 ++++++++-
 .../Handlers/StateHandlerFactory.swift        |  10 +-
 .../InferenceEngines/BatchedDecodeStep.swift  |  48 +++
 .../CoreAISequentialEngine.swift              | 357 +++++++++++++++---
 .../GenerationSessionState.swift              |  35 +-
 .../InferenceEngines/IdempotentEngine.swift   |   7 +-
 .../InferenceEngines/SequentialIterator.swift |   8 +-
 .../Samplers/BatchedSampler.swift             |  46 +++
 .../CoreAIShared/Bundle/ModelBundle.swift     |  10 +-
 .../Tools/benchmark/BenchmarkMain.swift       |  67 +++-
 .../CoreAISharedTests/ModelBundleTests.swift  |  41 ++
 .../BatchedInputBuilderTests.swift            |  94 +++++
 .../BatchedSamplerTests.swift                 |  58 +++
 .../IdempotentEngineTests.swift               |  35 ++
 .../InputHandlerHarnessTests.swift            |  23 ++
 .../LanguageModelBundleTests.swift            |  22 ++
 .../StateHandlerTests.swift                   | 137 +++++++
 23 files changed, 1250 insertions(+), 82 deletions(-)
 create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/BatchedInputBuilder.swift
 create mode 100644 swift/Sources/CoreAILanguageModels/Handlers/InputHandler+AttentionMask.swift
 create mode 100644 swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift
 create mode 100644 swift/Sources/CoreAILanguageModels/Samplers/BatchedSampler.swift
 create mode 100644 swift/Tests/LanguageModelsTests/BatchedInputBuilderTests.swift
 create mode 100644 swift/Tests/LanguageModelsTests/BatchedSamplerTests.swift

diff --git a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift b/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift
index 5203a05f..38d5f65a 100644
--- a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift
+++ b/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
diff --git a/swift/Sources/CoreAILanguageModels/Bundle/LanguageModelBundle.swift b/swift/Sources/CoreAILanguageModels/Bundle/LanguageModelBundle.swift
index 6c1cbe05..f30c03a2 100644
--- a/swift/Sources/CoreAILanguageModels/Bundle/LanguageModelBundle.swift
+++ b/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 }
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/BatchedInputBuilder.swift b/swift/Sources/CoreAILanguageModels/Handlers/BatchedInputBuilder.swift
new file mode 100644
index 00000000..6334a4b8
--- /dev/null
+++ b/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
+    }
+}
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/InputHandler+AttentionMask.swift b/swift/Sources/CoreAILanguageModels/Handlers/InputHandler+AttentionMask.swift
new file mode 100644
index 00000000..501672c8
--- /dev/null
+++ b/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
+}
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift b/swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift
index a1d1165b..d6cbb09c 100644
--- a/swift/Sources/CoreAILanguageModels/Handlers/InputHandler+Token.swift
+++ b/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
     }
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/InputLayout.swift b/swift/Sources/CoreAILanguageModels/Handlers/InputLayout.swift
index 9aa2426d..fe3fe7b3 100644
--- a/swift/Sources/CoreAILanguageModels/Handlers/InputLayout.swift
+++ b/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
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift
index 2d96dd6b..5a77d0a9 100644
--- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift
+++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift
@@ -60,8 +60,25 @@ public final class FixedNDArrayState: SyncStateHandler {
 
 // MARK: - Growing NDArray State
 
+/// A KV cache that can move one batch row's entry between sequence slots.
+///
+/// A ragged decode step writes every row's new token at slot `max(cursors)`, then moves each
+/// shorter row down to its own cursor. A byte copy is enough, because K was already rotated by
+/// RoPE with the row's real position.
+///
+/// ```swift
+/// for b in 0..<n where cursors[b] < writeSlot {
+///     cache.copyRegion(row: b, fromSeq: writeSlot, toSeq: cursors[b])
+/// }
+/// ```
+protocol RegionCopyable {
+    /// Copies row `row`'s K and V at slot `fromSeq` to slot `toSeq` in every state array. Does
+    /// nothing when the slots are equal.
+    func copyRegion(row: Int, fromSeq: Int, toSeq: Int)
+}
+
 /// Dynamically-growing KV cache state. Starts small and doubles capacity.
-public final class GrowingNDArrayState: SyncStateHandler {
+public final class GrowingNDArrayState: SyncStateHandler, RegionCopyable {
     public let stateNames: [String]
     public let supportsTruncation: Bool = true
     public private(set) var currentCapacity: Int
@@ -71,21 +88,34 @@ public final class GrowingNDArrayState: SyncStateHandler {
     private let descriptors: [NDArrayDescriptor]
     private let maxCapacity: Int
     private let sequenceDimIndex: Int
+    private let batchDimIndex: Int
+    private let batchSize: Int
 
     public init(
         states: [(name: String, descriptor: NDArrayDescriptor)],
         initialCapacity: Int,
-        maxCapacity: Int
+        maxCapacity: Int,
+        batchSize: Int = 1
     ) {
         self.descriptors = states.map(\.descriptor)
         self.stateNames = states.map(\.name)
+        self.batchSize = batchSize
 
         let firstDesc = states[0].descriptor
-        let dynamicDimIndex = firstDesc.shape.firstIndex(where: { $0 < 0 })
-        self.sequenceDimIndex = dynamicDimIndex ?? max(0, firstDesc.shape.count - 2)
-
-        // A static-shape cache (exported at its full context) has nowhere to grow.
-        let staticCapacity = dynamicDimIndex == nil ? firstDesc.shape[sequenceDimIndex] : nil
+        // The layout is [..., seq, head_dim]. The first dynamic dim would be wrong here, because it
+        // is the batch dim when batch is also dynamic.
+        self.sequenceDimIndex = max(0, firstDesc.shape.count - 2)
+        // Batch is the first dynamic dim other than seq, for example dim 1 in [G, B, H, seq, D].
+        // A static batch dim falls back to 0, so a static [G, B, H, seq, D] cache is not supported.
+        let seqDim = self.sequenceDimIndex
+        self.batchDimIndex =
+            firstDesc.shape.enumerated()
+            .first { $0.offset != seqDim && $0.element < 0 }?.offset ?? 0
+
+        // A static seq dim means the cache was exported at full context and cannot grow. A dynamic
+        // batch dim does not make it growable.
+        let sequenceExtent = firstDesc.shape[sequenceDimIndex]
+        let staticCapacity = sequenceExtent >= 0 ? sequenceExtent : nil
         self.maxCapacity = staticCapacity ?? maxCapacity
 
         let capacity = staticCapacity ?? min(initialCapacity, maxCapacity)
@@ -94,12 +124,24 @@ public final class GrowingNDArrayState: SyncStateHandler {
         var arrays: [String: NDArray] = [:]
         for (name, desc) in states {
             let resolved = desc.resolvingDynamicDimensions(
-                desc.shape.map { $0 < 0 ? capacity : $0 })
+                Self.resolveKVShape(
+                    desc.shape, sequenceDimIndex: sequenceDimIndex, sequenceLength: capacity,
+                    batchSize: batchSize))
             arrays[name] = NDArray(descriptor: resolved)
         }
         self.arrays = arrays
     }
 
+    /// Replaces the dynamic seq dim with `sequenceLength` and any other dynamic dim with `batchSize`.
+    static func resolveKVShape(
+        _ shape: [Int], sequenceDimIndex: Int, sequenceLength: Int, batchSize: Int
+    ) -> [Int] {
+        shape.enumerated().map { index, dim in
+            guard dim < 0 else { return dim }
+            return index == sequenceDimIndex ? sequenceLength : batchSize
+        }
+    }
+
     public func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
         guard contextLength > currentCapacity else { return false }
         guard contextLength <= maxCapacity else {
@@ -115,7 +157,9 @@ public final class GrowingNDArrayState: SyncStateHandler {
         let previousCapacity = currentCapacity
         for (i, name) in stateNames.enumerated() {
             let desc = descriptors[i]
-            let newShape = desc.shape.map { $0 < 0 ? newCapacity : $0 }
+            let newShape = Self.resolveKVShape(
+                desc.shape, sequenceDimIndex: sequenceDimIndex, sequenceLength: newCapacity,
+                batchSize: batchSize)
             let resolvedDesc = desc.resolvingDynamicDimensions(newShape)
             var newArray = NDArray(descriptor: resolvedDesc)
             _ = newArray.mutableRawView()
@@ -166,6 +210,125 @@ public final class GrowingNDArrayState: SyncStateHandler {
             from: source, to: &destination, blockCount: numBlocks, sourceBlockStride: srcBlockStride,
             destinationBlockStride: dstBlockStride, runElements: copyElements)
     }
+
+    // MARK: - Per-row copy
+
+    func copyRegion(row: Int, fromSeq: Int, toSeq: Int) {
+        guard fromSeq != toSeq else { return }
+        for name in stateNames {
+            Self.copyRegion(
+                in: &arrays[name]!, sequenceDimIndex: sequenceDimIndex, batchDimIndex: batchDimIndex,
+                row: row, fromSeq: fromSeq, toSeq: toSeq)
+        }
+    }
+
+    /// Array-level form of `copyRegion(row:fromSeq:toSeq:)`. The batch axis can be at any index, for
+    /// example 1 in `[G, B, H, seq, D]`.
+    static func copyRegion(
+        in array: inout NDArray, sequenceDimIndex: Int, batchDimIndex: Int,
+        row: Int, fromSeq: Int, toSeq: Int
+    ) {
+        guard fromSeq != toSeq else { return }
+        let shape = array.shape
+        guard let headDim = shape.last else { return }
+        let bases = blockBases(
+            shape: shape, batchDimIndex: batchDimIndex, sequenceDimIndex: sequenceDimIndex, row: row)
+
+        // A raw byte copy, so every byte-addressable scalar type works, 8-bit KV caches included.
+        guard let width = stateElementByteWidth(array.scalarType) else {
+            preconditionFailure("Packed sub-byte scalar type for region copy: \(array.scalarType)")
+        }
+        let runBytes = headDim * width
+        array.mutableRawView().withUnsafeMutableBytes { raw, _, _ in
+            let ptr = raw.assumingMemoryBound(to: UInt8.self)
+            for base in bases {
+                ptr.advanced(by: (base + toSeq * headDim) * width).update(
+                    from: ptr.advanced(by: (base + fromSeq * headDim) * width), count: runBytes)
+            }
+        }
+    }
+
+    /// Flat offsets of the `[seq, head_dim]` blocks of batch row `row`, one per index combination of
+    /// the other axes. For `[G, B, H, seq, D] = [2, 3, 2, 5, 2]` and row 1: `[20, 30, 80, 90]`.
+    private static func blockBases(
+        shape: [Int], batchDimIndex: Int, sequenceDimIndex: Int, row: Int
+    ) -> [Int] {
+        let outerCount = shape[0..<batchDimIndex].reduce(1, *)
+        let innerCount = shape[(batchDimIndex + 1)..<sequenceDimIndex].reduce(1, *)
+        let outerBlock = shape[batchDimIndex...].reduce(1, *)  // elements per outer index
+        let batchStride = shape[(batchDimIndex + 1)...].reduce(1, *)  // elements per batch index
+        let blockStride = shape[sequenceDimIndex...].reduce(1, *)  // seq * head_dim
+        var bases: [Int] = []
+        bases.reserveCapacity(outerCount * innerCount)
+        for o in 0..<outerCount {
+            let outerBase = o * outerBlock + row * batchStride
+            for n in 0..<innerCount {
+                bases.append(outerBase + n * blockStride)
+            }
+        }
+        return bases
+    }
+
+    /// Copies slots `[0, count)` of row 0 of the batch-1 cache `source` into row `toRow`, leaving
+    /// the other rows untouched. `source` can have a different sequence capacity, but this cache
+    /// must already hold `count` slots.
+    ///
+    /// Admitting a request into a running batch:
+    ///
+    /// ```swift
+    /// // Prefill the prompt alone into a batch-1 scratch cache first.
+    /// _ = try batchCache.ensureCapacity(forContextLength: promptLen)
+    /// batchCache.copyRowPrefix(from: scratchCache, toRow: row, count: promptLen)
+    /// ```
+    func copyRowPrefix(from source: GrowingNDArrayState, toRow: Int, count: Int) {
+        guard count > 0 else { return }
+        for name in stateNames {
+            guard let src = source.arrays[name] else {
+                preconditionFailure("copyRowPrefix: source missing state \(name)")
+            }
+            Self.copyRowPrefix(
+                from: src, sourceRow: 0, to: &arrays[name]!, destRow: toRow,
+                sequenceDimIndex: sequenceDimIndex, batchDimIndex: batchDimIndex, count: count)
+        }
+    }
+
+    /// Array-level form of `copyRowPrefix(from:toRow:count:)`. The arrays can differ in sequence
+    /// capacity and batch size, but must match in scalar type and every other dim.
+    static func copyRowPrefix(
+        from source: NDArray, sourceRow: Int, to destination: inout NDArray, destRow: Int,
+        sequenceDimIndex: Int, batchDimIndex: Int, count: Int
+    ) {
+        guard count > 0 else { return }
+        let srcShape = source.shape
+        let dstShape = destination.shape
+        guard let headDim = srcShape.last else { return }
+        precondition(
+            source.scalarType == destination.scalarType,
+            "copyRowPrefix: scalar type mismatch \(source.scalarType) vs \(destination.scalarType)")
+        let srcBases = blockBases(
+            shape: srcShape, batchDimIndex: batchDimIndex, sequenceDimIndex: sequenceDimIndex,
+            row: sourceRow)
+        let dstBases = blockBases(
+            shape: dstShape, batchDimIndex: batchDimIndex, sequenceDimIndex: sequenceDimIndex,
+            row: destRow)
+        precondition(srcBases.count == dstBases.count, "copyRowPrefix: block-count mismatch")
+        let copyElements = count * headDim  // slots [0,count) are contiguous from each block's start
+
+        guard let width = stateElementByteWidth(destination.scalarType) else {
+            preconditionFailure("Packed sub-byte scalar type for row-prefix copy: \(destination.scalarType)")
+        }
+        let copyBytes = copyElements * width
+        source.rawView().withUnsafeBytes { srcRaw, _, _ in
+            let srcPtr = srcRaw.assumingMemoryBound(to: UInt8.self)
+            destination.mutableRawView().withUnsafeMutableBytes { dstRaw, _, _ in
+                let dstPtr = dstRaw.assumingMemoryBound(to: UInt8.self)
+                for i in srcBases.indices {
+                    dstPtr.advanced(by: dstBases[i] * width).update(
+                        from: srcPtr.advanced(by: srcBases[i] * width), count: copyBytes)
+                }
+            }
+        }
+    }
 }
 
 // MARK: - Shared Utilities
diff --git a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift
index f1dbe61f..8b6aceff 100644
--- a/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift
+++ b/swift/Sources/CoreAILanguageModels/Handlers/StateHandlerFactory.swift
@@ -149,6 +149,7 @@ enum StateHandlerFactory {
         maxContextLength: Int,
         stateKinds: [String: StateKind]? = nil,
         options: EngineOptions = EngineOptions(),
+        batchSize: Int = 1,
         verbose: Bool = false
     ) throws -> SyncStateHandlerSet {
         guard !descriptor.stateNames.isEmpty else {
@@ -187,8 +188,12 @@ enum StateHandlerFactory {
         if !growingPairs.isEmpty {
             if options.kvCacheStrategy == .fixedSize {
                 let resolved = growingPairs.map { (name, desc) -> (name: String, descriptor: NDArrayDescriptor) in
+                    // Size the seq dim to the full context and a dynamic batch dim to `batchSize`.
+                    let seqDim = max(0, desc.shape.count - 2)
                     let resolvedDesc = desc.resolvingDynamicDimensions(
-                        desc.shape.map { $0 < 0 ? maxContextLength : $0 })
+                        GrowingNDArrayState.resolveKVShape(
+                            desc.shape, sequenceDimIndex: seqDim, sequenceLength: maxContextLength,
+                            batchSize: batchSize))
                     return (name, resolvedDesc)
                 }
                 kvCache = FixedNDArrayState(states: resolved)
@@ -197,7 +202,8 @@ enum StateHandlerFactory {
                 kvCache = GrowingNDArrayState(
                     states: growingPairs,
                     initialCapacity: initial,
-                    maxCapacity: maxContextLength
+                    maxCapacity: maxContextLength,
+                    batchSize: batchSize
                 )
             }
         } else {
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift
new file mode 100644
index 00000000..ff5805e4
--- /dev/null
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/BatchedDecodeStep.swift
@@ -0,0 +1,48 @@
+// 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
+
+// MARK: - Batched Decode Step
+
+/// Runs one forward pass over the `N` rows of a batch KV cache and returns each row's last-token
+/// logits. Sampling, stop handling and admission stay with the caller.
+///
+/// The caller tracks each row's cursor. Neither method advances `session.processedTokenCount`.
+///
+/// ```swift
+/// let session = try engine.makeSessionState(batchSize: 2)
+/// let firstA = try await engine.prefillRow(row: 0, promptTokens: promptA, session: session)
+/// let firstB = try await engine.prefillRow(row: 1, promptTokens: promptB, session: session)
+/// // Sample nextA and nextB, then decode with each row at its own cursor.
+/// let logits = try await engine.decodeBatch(
+///     tokensPerRow: [[nextA], [nextB]],
+///     startPositions: [promptA.count, promptB.count],
+///     session: session)
+/// ```
+protocol BatchedDecodeStep {
+    /// Most rows one step can take: `N` for a static `batch=N` graph, 1 for a single-batch graph, and
+    /// `Int.max` for a dynamic-batch graph, whose traced maximum is only enforced at run time.
+    var maxBatchCapacity: Int { get }
+
+    /// Runs `tokensPerRow` (`[N][q]`) and returns each row's last-token logits as `[N][vocab]`.
+    ///
+    /// `startPositions[i]` is row `i`'s cache cursor. A decode step (`q == 1`) may give each row its
+    /// own cursor, which needs an `attn_mask` input and a `RegionCopyable` KV cache. Every row then
+    /// writes at slot `max(startPositions)`, and shorter rows are copied down to their own cursor
+    /// afterwards. A prefill step (`q > 1`) needs the same cursor for every row.
+    func decodeBatch(
+        tokensPerRow: [[Int32]],
+        startPositions: [Int],
+        session: GenerationSessionState
+    ) async throws -> [[LogitsScalarType]]
+
+    /// Prefills `promptTokens` into row `row` of `session`'s KV cache, at positions
+    /// `[0, promptTokens.count)`, and returns the last prompt token's logits as `[vocab]`. The other
+    /// rows are left untouched, so a request can join a running batch.
+    func prefillRow(
+        row: Int,
+        promptTokens: [Int32],
+        session: GenerationSessionState
+    ) async throws -> [LogitsScalarType]
+}
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift
index 5b015648..69f557c1 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift
@@ -29,7 +29,7 @@ enum PrefillStrategy {
 ///
 /// KV cache NDArrays start small (256 tokens) and grow dynamically with 2× expansion.
 /// Passed as `states` on every forward pass; the model graph updates them in-place.
-public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @unchecked Sendable {
+public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, BatchedDecodeStep, @unchecked Sendable {
     public typealias ConfigType = ModelConfig
 
     public var supportsLogits: Bool { true }
@@ -57,8 +57,9 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
     // I/O names from descriptor
     private let logitsName: String
 
-    // Input handling — handler owns allocation and fill logic
-    private var inputHandler: TokenInputHandler
+    // Builds a new input handler for each session, so sessions never share its reusable buffers.
+    // When the graph has an `attn_mask` input, the handler fills it as a composite extra.
+    private let makeInputHandler: () throws -> CompositeInputHandler<TokenInputHandler>
 
     /// Retained to build per-session state handlers in `makeSessionState()`.
     private let options: EngineOptions
@@ -68,28 +69,42 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
     // semantics. Idempotent callers pass their own session instead.
     private let session: GenerationSessionState
 
-    // Logits descriptor and buffer
+    // The single-row logits buffer lives on each session (`GenerationSessionState.logitsScratch`).
     private let logitsDescriptor: NDArrayDescriptor
-    private var logitsArray: NDArray
-    private var cachedLogitsBatchSize: Int
 
     // Ring buffer mode: handled by TokenInputHandler.useCompactPositionIds
 
     // Track processed tokens for incremental inference (delegated to the shim session).
     public var processedTokenCount: Int { session.processedTokenCount }
 
-    public private(set) var lastPrefixHitCount: Int = 0
+    /// The static batch size of the graph, read from the leading dimension of `input_ids`.
+    ///
+    /// A dynamic-batch graph reports 1 here, while `maxBatchCapacity` reports `Int.max` for it.
+    package var declaredBatchSize: Int {
+        guard let name = InputLayout.knownInputIdNames.first(where: functionDescriptor.inputNames.contains),
+            case .ndArray(let desc) = functionDescriptor.inputDescriptor(of: name),
+            let dim0 = desc.shape.first, dim0 > 0
+        else { return 1 }
+        return dim0
+    }
 
-    // Track in-flight generation via token (replaces simple bool lock)
-    private let tokenBox = GenerationTokenBox()
+    var maxBatchCapacity: Int {
+        guard let name = InputLayout.knownInputIdNames.first(where: functionDescriptor.inputNames.contains),
+            case .ndArray(let desc) = functionDescriptor.inputDescriptor(of: name),
+            let dim0 = desc.shape.first
+        else { return 1 }
+        if dim0 > 1 { return dim0 }  // Static batch=N graph.
+        if dim0 < 0 { return Int.max }  // Dynamic batch. The runtime enforces the exported maximum.
+        return 1  // Single-batch graph.
+    }
 
-    public var isBusy: Bool { tokenBox.isBusy }
+    /// The prefix hit count of the most recent generation. With several sessions in use, this is
+    /// the count of whichever session ran last.
+    public private(set) var lastPrefixHitCount: Int = 0
 
-    /// Clear the engine's active token if it matches the given token.
-    /// Called by the iterator when generation finishes or is cancelled.
-    func clearTokenIfActive(_ token: GenerationToken) {
-        tokenBox.clearIfActive(token)
-    }
+    /// Whether a generation started by `generate(with:samplingConfiguration:inferenceOptions:)` is
+    /// in flight. Generations on caller-supplied sessions are not counted.
+    public var isBusy: Bool { session.tokenBox.isBusy }
 
     // MARK: - Init
 
@@ -114,11 +129,14 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
         }
         self.functionDescriptor = descriptor
 
-        // Validate model architecture: 2 inputs, 1+ output, at least KV cache pair.
+        // Validate model architecture: 2 inputs (3 with `attn_mask`), 1+ output, at least KV cache pair.
         // Hybrid models may declare additional persistent fixed-shape states.
-        guard descriptor.inputNames.count == 2 else {
+        let hasAttnMask = descriptor.inputNames.contains(where: InputLayout.knownAttnMaskNames.contains)
+        guard descriptor.inputNames.count == 2 || (descriptor.inputNames.count == 3 && hasAttnMask)
+        else {
             throw InferenceRuntimeError.invalidInputType(
-                "Expected 2 inputs, got \(descriptor.inputNames.count): \(descriptor.inputNames)")
+                "Expected 2 inputs (or 3 with attn_mask), got \(descriptor.inputNames.count): "
+                    + "\(descriptor.inputNames)")
         }
         guard descriptor.outputNames.count >= 1 else {
             throw InferenceRuntimeError.invalidOutputType(
@@ -170,13 +188,31 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
         }
         self.logitsDescriptor = logitsDesc
 
-        self.inputHandler = TokenInputHandler(
-            inputIdsName: inputIdsName,
-            positionIdsName: positionIdsName,
-            inputIdsDescriptor: inputIdsDesc,
-            positionIdsDescriptor: posIdsDesc,
-            useCompactPositionIds: layout.positionPolicy == .compact
-        )
+        // With an `attn_mask` input, the handler fills a single-row causal mask each step.
+        // `decodeBatch` builds its own per-row masks and does not use this handler.
+        let compactPositions = layout.positionPolicy == .compact
+        let attnMaskName = layout.attnMaskName
+        self.makeInputHandler = {
+            var maskExtras: [CompositeInputHandler<TokenInputHandler>.ExtraInput] = []
+            if let maskName = attnMaskName {
+                guard case .ndArray(let maskDesc) = descriptor.inputDescriptor(of: maskName) else {
+                    throw InferenceRuntimeError.invalidInputType(
+                        "Cannot get descriptor for '\(maskName)'")
+                }
+                maskExtras.append(
+                    .init(name: maskName) { context in
+                        makeAdditiveCausalMask(context: context, descriptor: maskDesc)
+                    })
+            }
+            let tokenHandler = TokenInputHandler(
+                inputIdsName: inputIdsName,
+                positionIdsName: positionIdsName,
+                inputIdsDescriptor: inputIdsDesc,
+                positionIdsDescriptor: posIdsDesc,
+                useCompactPositionIds: compactPositions
+            )
+            return CompositeInputHandler(base: tokenHandler, extras: maskExtras)
+        }
 
         CLILogger.log(
             "KV cache: capacity=\(stateHandlers.kvCache.currentCapacity), states=\(stateHandlers.kvCache.stateNames)"
@@ -186,11 +222,6 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
                 "Additional persistent states: \(additional.stateNames.joined(separator: ", "))")
         }
 
-        // Allocate initial logits (1 token — will be reallocated per batch)
-        let initLogitsDesc = logitsDesc.resolvingDynamicDimensions([1, 1, config.vocabSize])
-        self.logitsArray = NDArray(descriptor: initLogitsDesc)
-        self.cachedLogitsBatchSize = 1
-
         // Load inference function
         self.prefillFunction = try loadPrefillGraph(
             from: model, matching: descriptor, mainName: config.function)
@@ -212,6 +243,9 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
         CLILogger.log(
             "CoreAI clean engine initialized — inputs: \(descriptor.inputNames), outputs: \(descriptor.outputNames), states: \(descriptor.stateNames)"
         )
+
+        // A session without a handler throws in `processTokenBatch`, so the default session needs one.
+        self.session.inputHandler = try makeInputHandler()
     }
 
     /// Convenience initializer with direct model URL.
@@ -265,14 +299,19 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
 
         let context = InputContext.dynamic(
             tokens: tokens, processedTokenCount: sessionState.processedTokenCount)
-        let inputs = try await inputHandler.prepare(context)
+        guard var handler = sessionState.inputHandler else {
+            throw InferenceRuntimeError.invalidState("session has no input handler")
+        }
+        let inputs = try await handler.prepare(context)
+        sessionState.inputHandler = handler  // Value type: store it back to keep its reusable buffers.
 
-        // Reuse pre-allocated logits when the batch size is unchanged.
-        if cachedLogitsBatchSize != batchSize {
+        // Reuse this session's logits buffer, and reallocate only when the step's token count changes.
+        if sessionState.logitsScratchBatchSize != batchSize || sessionState.logitsScratch == nil {
             let resolvedLogitsDesc = logitsDescriptor.resolvingDynamicDimensions([1, batchSize, config.vocabSize])
-            logitsArray = NDArray(descriptor: resolvedLogitsDesc)
-            cachedLogitsBatchSize = batchSize
+            sessionState.logitsScratch = NDArray(descriptor: resolvedLogitsDesc)
+            sessionState.logitsScratchBatchSize = batchSize
         }
+        var logitsArray = sessionState.logitsScratch!
 
         // Bind states, build output views, and execute
         try await runWithStates(
@@ -283,6 +322,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
             outputArray: &logitsArray,
             outputName: logitsName
         )
+        sessionState.logitsScratch = logitsArray  // Store the handle back so the next step reuses it.
 
         // Read logits from NDArray
         let totalLogits = batchSize * config.vocabSize
@@ -298,6 +338,213 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
         return logitBuffer
     }
 
+    // MARK: - Batched decode
+
+    /// Output of `lockstepBatchedGenerate`. `tokens` is `[N][generated]`, and a row that hit EOS is
+    /// shorter and ends with the EOS token. Times are wall-clock seconds.
+    package struct BatchedRunResult: Sendable {
+        package let tokens: [[Int32]]
+        package let prefillSeconds: Double
+        package let decodeSeconds: Double
+    }
+
+    /// Prefills and decodes `N` equal-length prompts in lockstep on a new batch=N session, sampling
+    /// each row with its own entry in `configs`.
+    ///
+    /// A row that samples an EOS token stops collecting tokens but stays in the batch until every
+    /// row is done. Pass an empty `eosTokenIds` to always generate `maxNewTokens` per row.
+    ///
+    /// The prompts are prefilled as one `decodeBatch` step on the decode graph, not through the
+    /// chunked prefill path that `prefillRow` uses. For rows at different cursors, use `prefillRow`
+    /// and `decodeBatch` directly.
+    package func lockstepBatchedGenerate(
+        promptRows: [[Int32]],
+        configs: [SamplingConfiguration],
+        maxNewTokens: Int,
+        eosTokenIds: Set<Int32>
+    ) async throws -> BatchedRunResult {
+        precondition(!promptRows.isEmpty && promptRows.count == configs.count)
+        let n = promptRows.count
+        let promptLen = promptRows[0].count
+        precondition(
+            promptRows.allSatisfy { $0.count == promptLen }, "lockstep requires equal-length rows")
+
+        let batchSession = try makeSessionState(batchSize: n)
+        _ = try batchSession.kvCache.ensureCapacity(forContextLength: promptLen + maxNewTokens)
+
+        var gen = Array(repeating: [Int32](), count: n)
+        var live = Array(repeating: true, count: n)
+
+        let prefillStart = SuspendingClock.now
+        var perRowLast = try await decodeBatch(
+            tokensPerRow: promptRows,
+            startPositions: Array(repeating: 0, count: n),
+            session: batchSession)
+        let prefillSeconds = (SuspendingClock.now - prefillStart).inSeconds
+
+        let decodeStart = SuspendingClock.now
+        for s in 0..<maxNewTokens {
+            let toks = BatchedSampler.sample(
+                rows: perRowLast,
+                configurations: configs,
+                histories: gen.map { ArraySlice($0) },
+                steps: (0..<n).map { gen[$0].count })
+            for b in 0..<n where live[b] {
+                gen[b].append(toks[b])
+                if eosTokenIds.contains(toks[b]) { live[b] = false }
+            }
+            if !live.contains(true) { break }
+            perRowLast = try await decodeBatch(
+                tokensPerRow: toks.map { [$0] },
+                startPositions: Array(repeating: promptLen + s, count: n),
+                session: batchSession)
+        }
+        let decodeSeconds = (SuspendingClock.now - decodeStart).inSeconds
+        return BatchedRunResult(
+            tokens: gen, prefillSeconds: prefillSeconds, decodeSeconds: decodeSeconds)
+    }
+
+    /// See ``BatchedDecodeStep`` for the contract and the admission pattern.
+    func decodeBatch(
+        tokensPerRow: [[Int32]],
+        startPositions: [Int],
+        session: GenerationSessionState
+    ) async throws -> [[LogitsScalarType]] {
+        precondition(!tokensPerRow.isEmpty, "decodeBatch requires at least one row")
+        let n = tokensPerRow.count
+        precondition(startPositions.count == n, "one start position per row")
+        let q = tokensPerRow[0].count
+        precondition(
+            tokensPerRow.allSatisfy { $0.count == q }, "lockstep decode requires equal-length rows")
+        let isRagged = !startPositions.allSatisfy { $0 == startPositions[0] }
+        precondition(
+            !isRagged || q == 1, "ragged (per-row) start positions require q == 1 (decode only)")
+
+        let names = functionDescriptor.inputNames
+        guard let inputIdsName = InputLayout.knownInputIdNames.first(where: names.contains),
+            let positionIdsName = InputLayout.knownPositionIdNames.first(where: names.contains)
+        else { throw InferenceRuntimeError.invalidState("missing input_ids/position_ids") }
+        let maskName = InputLayout.knownAttnMaskNames.first(where: names.contains)
+        guard case .ndArray(let idsDesc) = functionDescriptor.inputDescriptor(of: inputIdsName),
+            case .ndArray(let posDesc) = functionDescriptor.inputDescriptor(of: positionIdsName)
+        else { throw InferenceRuntimeError.invalidInputType("bad input descriptors") }
+        var maskDesc: NDArrayDescriptor?
+        if let maskName, case .ndArray(let d) = functionDescriptor.inputDescriptor(of: maskName) {
+            maskDesc = d
+        }
+        if isRagged {
+            guard maskName != nil, maskDesc != nil else {
+                throw InferenceRuntimeError.invalidState(
+                    "ragged decode needs an attn_mask graph input (export with --attention-mask batched)")
+            }
+            guard session.kvCache is RegionCopyable else {
+                throw InferenceRuntimeError.invalidState(
+                    "ragged decode requires a RegionCopyable KV cache")
+            }
+        }
+
+        let vocab = config.vocabSize
+        // Every row writes its new token at `writeSlot = max(startPositions)`. After the forward pass,
+        // shorter rows are copied down to their own cursors. Equal cursors need no copy.
+        let writeSlot = startPositions.max() ?? 0
+        let seqLen = writeSlot + q
+
+        // `BatchedInputBuilder` documents the flat buffer layouts. position_ids and the mask span the
+        // whole key axis `seqLen`, not just the `q` new tokens.
+        let idsFlat = BatchedInputBuilder.tokenIDs(rows: tokensPerRow)
+        var ids = NDArray(descriptor: idsDesc.resolvingDynamicDimensions([n, q]))
+        fillNDArray(&ids, as: Int32.self, count: n * q) { i in idsFlat[i] }
+
+        let posFlat =
+            isRagged
+            ? BatchedInputBuilder.raggedDecodePositionIDs(cursors: startPositions)
+            : BatchedInputBuilder.positionIDs(rowCount: n, startPos: writeSlot, queryLen: q)
+        var pos = NDArray(descriptor: posDesc.resolvingDynamicDimensions([n, seqLen]))
+        fillNDArray(&pos, as: Int32.self, count: n * seqLen) { i in posFlat[i] }
+
+        var inputs: [String: NDArray] = [inputIdsName: ids, positionIdsName: pos]
+        if let maskName, let maskDesc {
+            let maskFlat =
+                isRagged
+                ? BatchedInputBuilder.raggedDecodeMask(cursors: startPositions)
+                : BatchedInputBuilder.additiveCausalMask(
+                    rowCount: n, startPos: writeSlot, queryLen: q)
+            var mask = NDArray(descriptor: maskDesc.resolvingDynamicDimensions([n, 1, q, seqLen]))
+            fillNDArray(&mask, as: Float16.self, count: n * q * seqLen) { i in maskFlat[i] }
+            inputs[maskName] = mask
+        }
+
+        var out = NDArray(descriptor: logitsDescriptor.resolvingDynamicDimensions([n, q, vocab]))
+        // The batch cache starts small, so grow it before the graph writes up to `seqLen`.
+        _ = try session.kvCache.ensureCapacity(forContextLength: seqLen)
+        try await runWithStates(
+            function: function, inputs: inputs,
+            primary: session.kvCache, secondary: session.additionalStates,
+            outputArray: &out, outputName: logitsName)
+
+        // Move each shorter row's new KV entry from `writeSlot` down to its own cursor.
+        if isRagged, let region = session.kvCache as? RegionCopyable {
+            for b in 0..<n where startPositions[b] < writeSlot {
+                region.copyRegion(row: b, fromSeq: writeSlot, toSeq: startPositions[b])
+            }
+        }
+
+        let flat = readNDArray(out, as: LogitsScalarType.self, count: n * q * vocab)
+        return (0..<n).map { b in
+            let base = (b * q + (q - 1)) * vocab
+            return Array(flat[base..<base + vocab])
+        }
+    }
+
+    /// See ``BatchedDecodeStep`` for the contract.
+    ///
+    /// The prompt runs alone on a batch-1 scratch session through the normal prefill path, so the
+    /// row matches a standalone request. Its `[0, promptLen)` KV prefix is then copied into the row.
+    /// Throws for models with additional states (conv or sliding window), which are not per row.
+    func prefillRow(
+        row: Int,
+        promptTokens: [Int32],
+        session: GenerationSessionState
+    ) async throws -> [LogitsScalarType] {
+        precondition(!promptTokens.isEmpty, "prefillRow requires a non-empty prompt")
+        guard let target = session.kvCache as? GrowingNDArrayState else {
+            throw InferenceRuntimeError.invalidState(
+                "prefillRow requires an NDArray-backed (GrowingNDArrayState) KV cache")
+        }
+        if session.additionalStates != nil {
+            throw InferenceRuntimeError.invalidState(
+                "prefillRow does not yet support models with additional (conv/sliding) states")
+        }
+        let promptLen = promptTokens.count
+        // Use the normal prefill path, which runs chunks on the separate `prefill` graph when the asset
+        // has one. Running the prompt through `decodeBatch(q: promptLen)` would change the decode
+        // executable's input shapes on every admission, and each shape change leaks its heap.
+        // Assets without a prefill graph still run the prompt on the decode graph.
+        let scratch = try makeSessionState(batchSize: 1)
+        let promptSlice = promptTokens[...]
+        let logits: [LogitsScalarType]
+        switch selectPrefillStrategy(newTokenCount: promptLen) {
+        case .chunked(let chunkSize):
+            logits = try await processChunkedPrompt(
+                tokens: promptSlice, chunkSize: chunkSize, sessionState: scratch)
+        case .wholeBatch:
+            let allLogits = try await processTokenBatch(promptSlice, sessionState: scratch)
+            logits = lastTokenLogits(from: allLogits, vocabSize: config.vocabSize)
+        case .oneAtATime:
+            var last: [LogitsScalarType] = []
+            for j in promptSlice.indices {
+                last = try await processTokenBatch(promptSlice[j...j], sessionState: scratch)
+            }
+            logits = last
+        }
+        guard let scratchCache = scratch.kvCache as? GrowingNDArrayState else {
+            throw InferenceRuntimeError.invalidState("prefillRow scratch cache is not NDArray-backed")
+        }
+        _ = try target.ensureCapacity(forContextLength: promptLen)
+        target.copyRowPrefix(from: scratchCache, toRow: row, count: promptLen)
+        return logits
+    }
+
     /// Run one prefill chunk on the prefill graph: KV cache writes only, no logits.
     private func encodePrefillChunk(
         _ tokens: ArraySlice<Int32>, using prefillFn: InferenceFunction,
@@ -309,7 +556,11 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
 
         let context = InputContext.dynamic(
             tokens: tokens, processedTokenCount: sessionState.processedTokenCount)
-        let inputs = try await inputHandler.prepare(context)
+        guard var handler = sessionState.inputHandler else {
+            throw InferenceRuntimeError.invalidState("session has no input handler")
+        }
+        let inputs = try await handler.prepare(context)
+        sessionState.inputHandler = handler  // Value type: store it back to keep its reusable buffers.
 
         try await runWithStatesNoOutputs(
             function: prefillFn,
@@ -439,7 +690,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
     /// Wait for any in-flight generate() Task to finish.
     private func drain() {
         var attempts = 0
-        while tokenBox.isBusy {
+        while session.tokenBox.isBusy {
             attempts += 1
             if attempts > 5000 {
                 fatalError("Sequential engine drain() timeout — generation Task stuck?")
@@ -449,7 +700,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
     }
 
     public func cancel() async throws {
-        tokenBox.cancelActive()
+        session.tokenBox.cancelActive()
     }
 
     public func reset(to tokenIndex: Int) async throws {
@@ -468,7 +719,7 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
                 "Partial reset to \(tokenIndex) would read sliding-ring slots that have been "
                     + "overwritten. Use reset(to: 0) and replay the prefix.")
         }
-        tokenBox.cancelActive()
+        session.tokenBox.cancelActive()
         internalReset(to: tokenIndex, sessionState: session)
     }
 
@@ -502,15 +753,24 @@ public final class CoreAISequentialEngine: InferenceEngine, IdempotentEngine, @u
 extension CoreAISequentialEngine {
     /// Mint a fresh, independent per-request state (fresh KV cache + empty history).
     package func makeSessionState() throws -> GenerationSessionState {
+        try makeSessionState(batchSize: 1)
+    }
+
+    /// Creates a session whose KV cache has `batchSize` rows. A `batchSize` above 1 needs a graph
+    /// with `maxBatchCapacity >= batchSize`.
+    func makeSessionState(batchSize: Int) throws -> GenerationSessionState {
         let stateHandlers = try StateHandlerFactory.createSyncHandlers(
             descriptor: functionDescriptor,
             maxContextLength: config.maxContextLength,
             stateKinds: config.states,
-            options: options)
-        return GenerationSessionState(
+            options: options,
+            batchSize: batchSize)
+        let session = GenerationSessionState(
             kvCache: stateHandlers.kvCache,
             additionalStates: stateHandlers.additionalStates,
             hasNonTruncatableStates: stateHandlers.hasNonTruncatableStates)
+        session.inputHandler = try makeInputHandler()
+        return session
     }
 
     /// Drive generation over the caller-supplied state instead of engine-owned state.
@@ -520,8 +780,8 @@ extension CoreAISequentialEngine {
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
-        // Cancel any prior generation so its Iterator stops on next poll.
-        tokenBox.cancelActive()
+        // Cancel an earlier generation on this session only. Other sessions keep running.
+        sessionState.tokenBox.cancelActive()
 
         // Implicit prefix caching: resolve input against history.
         // For hybrid models with recurrent states, we must full-reset on any
@@ -529,7 +789,7 @@ extension CoreAISequentialEngine {
         // be truncated by moving a KV cursor. Future: checkpoint/restore.
         // Default to 0 so a fresh session (no history) always reports a clean
         // value instead of leaking the previous session's hit count.
-        lastPrefixHitCount = 0
+        sessionState.lastPrefixHitCount = 0
         if sessionState.history.count > 0 {
             let (commonPrefix, _) = sessionState.history.resolve(input: input)
             let plan = Self.prefixResetPlan(
@@ -543,11 +803,12 @@ extension CoreAISequentialEngine {
             if let resetTo = plan.resetTo {
                 internalReset(to: resetTo, sessionState: sessionState)
             }
-            lastPrefixHitCount = plan.prefixHitCount
+            sessionState.lastPrefixHitCount = plan.prefixHitCount
         }
+        lastPrefixHitCount = sessionState.lastPrefixHitCount
 
         let token = GenerationToken()
-        tokenBox.install(token)
+        sessionState.tokenBox.install(token)
         return GenerationSequence(
             engine: self,
             input: input,
@@ -686,7 +947,7 @@ extension CoreAISequentialEngine.GenerationSequence {
         }
 
         deinit {
-            engine.clearTokenIfActive(generationToken)
+            sessionState.tokenBox.clearIfActive(generationToken)
         }
 
         public func next() async throws -> InferenceOutput? {
@@ -840,7 +1101,7 @@ extension CoreAISequentialEngine.GenerationSequence {
                 return
             }
             finished = true
-            engine.clearTokenIfActive(generationToken)
+            sessionState.tokenBox.clearIfActive(generationToken)
         }
     }
 }
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationSessionState.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationSessionState.swift
index b3082e9a..f9bcf08f 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationSessionState.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationSessionState.swift
@@ -3,22 +3,19 @@
 // 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
+
 // MARK: - Generation Session State
 
 /// Single-owner container for the per-request mutable state of an idempotent engine.
 ///
-/// Holds everything a `CoreAISequentialEngine` generation mutates on a per-request
-/// basis — the KV cache, optional persistent (hybrid) states, the processed-token
-/// cursor, and the implicit prefix-cache history. Relocating this state off the
-/// engine lets one engine serve independent sessions sequentially: a fresh session
-/// per conversation, or a reused one for prefix reuse. The engine's single
-/// `GenerationTokenBox` serializes generation, so only one session is driven at a
-/// time; each carries its own cache and history, so `generate(with:sessionState:…)`
-/// is idempotent with respect to the engine.
+/// It holds the KV cache, any hybrid states, the processed-token cursor, the prefix-cache history,
+/// the in-flight generation token, and the last prefix-hit count. Use a new session for each
+/// conversation, or reuse one to keep its prefix.
 ///
-/// Not thread-safe — a session is owned by exactly one active generation at a time.
-/// The engine's `GenerationTokenBox` enforces single-active-generation, so only one
-/// task ever mutates a given session; concurrent access is a programming error.
+/// Not thread-safe. Only one generation drives a session at a time: starting a generation cancels
+/// the previous one on the same session through its own `GenerationTokenBox`. Other sessions are
+/// not affected.
 ///
 /// It is deliberately not `Sendable`: it holds non-`Sendable`, NDArray-backed
 /// state (`SyncStateHandler` / `FixedNDArrayState`) and is confined to the engine's
@@ -47,6 +44,22 @@ package final class GenerationSessionState {
     /// Token history backing implicit prefix caching.
     var history = TokenHistory()
 
+    /// This session's in-flight generation. A new generation cancels only this session's previous one.
+    let tokenBox = GenerationTokenBox()
+
+    /// Prefix-cache hit count of this session's last generation. The engine-level
+    /// `lastPrefixHitCount` is copied from whichever session ran last.
+    var lastPrefixHitCount: Int = 0
+
+    /// Reusable logits output for the single-row path, reallocated when the step's token count
+    /// changes. `decodeBatch` does not use it.
+    var logitsScratch: NDArray?
+    /// Token count `logitsScratch` was sized for.
+    var logitsScratchBatchSize: Int = 0
+
+    /// This session's own input handler, so sessions never share input buffers.
+    var inputHandler: (any SyncInputHandler)?
+
     init(
         kvCache: any SyncStateHandler,
         additionalStates: FixedNDArrayState?,
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/IdempotentEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/IdempotentEngine.swift
index 59164b21..79058d6b 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/IdempotentEngine.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/IdempotentEngine.swift
@@ -24,10 +24,9 @@
 ///
 /// Caveats:
 /// - A session is single-owner: only one active generation may drive a given state
-///   at a time (enforced by `GenerationTokenBox`; see `GenerationSessionState`).
-/// - Not all engine state is per-session yet. `lastPrefixHitCount` and the shared
-///   `GenerationTokenBox` stay engine-scoped, so sessions run sequentially, not
-///   concurrently; a concurrent scheduler would need to relocate those too.
+///   at a time. Starting a generation cancels the previous one on the same session.
+/// - Sessions do not affect each other. The engine-level `lastPrefixHitCount` reports only the
+///   session that ran last.
 package protocol IdempotentEngine: InferenceEngine {
     /// Mint a fresh, independent per-request state (fresh KV cache + empty history).
     func makeSessionState() throws -> GenerationSessionState
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift
index 0c2e40f4..a7780a59 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/SequentialIterator.swift
@@ -38,7 +38,11 @@ enum SequentialIterator {
         if let forced {
             return forced[step]
         }
-        var mutableLogits = logits
-        return sampling.fallbackSampler(from: &mutableLogits, tokenHistory: tokenHistory, step: step)
+        return BatchedSampler.sample(
+            rows: [logits],
+            configurations: [sampling],
+            histories: [tokenHistory],
+            steps: [step]
+        )[0]
     }
 }
diff --git a/swift/Sources/CoreAILanguageModels/Samplers/BatchedSampler.swift b/swift/Sources/CoreAILanguageModels/Samplers/BatchedSampler.swift
new file mode 100644
index 00000000..4e7f4776
--- /dev/null
+++ b/swift/Sources/CoreAILanguageModels/Samplers/BatchedSampler.swift
@@ -0,0 +1,46 @@
+// 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 Foundation
+
+/// Samples one token per row, each row with its own ``SamplingConfiguration`` and step.
+///
+/// Rows are independent, so one batch can mix greedy and seeded temperature rows. A seeded row's
+/// draw depends only on its own `(seed, step)`. Each configuration is normalized before sampling.
+struct BatchedSampler {
+    /// Samples the next token for every row.
+    ///
+    /// - Parameters:
+    ///   - rows: Logits for each row, each `[vocab]`, such as the rows of a `[N, vocab]` output.
+    ///   - configurations: Sampling configuration for each row.
+    ///   - histories: Recent tokens for each row's repetition penalty. Pass an empty slice when unused.
+    ///   - steps: Generation step for each row. Required for rows whose configuration sets `seed`,
+    ///     ignored otherwise.
+    /// - Precondition: All four arrays have the same count.
+    /// - Returns: One token id per row, in row order.
+    static func sample(
+        rows: [[LogitsScalarType]],
+        configurations: [SamplingConfiguration],
+        histories: [ArraySlice<Int32>],
+        steps: [Int?]
+    ) -> [Int32] {
+        precondition(
+            rows.count == configurations.count
+                && rows.count == histories.count
+                && rows.count == steps.count,
+            "BatchedSampler: rows (\(rows.count)), configurations (\(configurations.count)), "
+                + "histories (\(histories.count)) and steps (\(steps.count)) must have equal counts")
+
+        var tokens = [Int32]()
+        tokens.reserveCapacity(rows.count)
+        for row in rows.indices {
+            var logits = rows[row]
+            tokens.append(
+                configurations[row].normalized().fallbackSampler(
+                    from: &logits, tokenHistory: histories[row], step: steps[row]))
+        }
+        return tokens
+    }
+}
diff --git a/swift/Sources/CoreAIShared/Bundle/ModelBundle.swift b/swift/Sources/CoreAIShared/Bundle/ModelBundle.swift
index 2f2e99d6..0196dbdc 100644
--- a/swift/Sources/CoreAIShared/Bundle/ModelBundle.swift
+++ b/swift/Sources/CoreAIShared/Bundle/ModelBundle.swift
@@ -41,6 +41,9 @@ public struct ModelBundle: Sendable {
     /// their own blocks without re-reading the file.
     public let raw: Data
 
+    /// `metadata_version` values this runtime loads. Any other version is rejected.
+    private static let supportedMetadataVersions = ["0.3", "0.2"]
+
     // MARK: - Component Keys
 
     public enum ComponentKey {
@@ -124,7 +127,8 @@ public struct ModelBundle: Sendable {
             case .malformedMetadata(let url, let err):
                 return "malformed metadata.json at \(url.path): \(err)"
             case .unsupportedVersion(let v):
-                return "unsupported metadata_version '\(v)' (known: 0.2)"
+                let known = ModelBundle.supportedMetadataVersions.joined(separator: ", ")
+                return "unsupported metadata_version '\(v)' (known: \(known))"
             case .kindMismatch(let expected, let got):
                 return "expected bundle kind \(expected), got \(got)"
             case .missingField(let name):
@@ -185,7 +189,7 @@ public struct ModelBundle: Sendable {
         }
 
         let version = envelope.metadataVersion ?? "0.1"
-        guard version == "0.2" else {
+        guard Self.supportedMetadataVersions.contains(version) else {
             throw BundleError.unsupportedVersion(version)
         }
 
@@ -196,7 +200,7 @@ public struct ModelBundle: Sendable {
             throw BundleError.malformedMetadata(
                 bundlePath.appending(path: "metadata.json"), underlying: error)
         }
-        self.metadataVersion = "0.2"
+        self.metadataVersion = version
         self.kind = common.kind
         self.name = common.name
         self.userData = common.userData
diff --git a/swift/Sources/Tools/benchmark/BenchmarkMain.swift b/swift/Sources/Tools/benchmark/BenchmarkMain.swift
index 186b529c..05c08e73 100644
--- a/swift/Sources/Tools/benchmark/BenchmarkMain.swift
+++ b/swift/Sources/Tools/benchmark/BenchmarkMain.swift
@@ -62,10 +62,23 @@ struct LLMBenchmark: AsyncParsableCommand {
     )
     var clearCoreAICache: Bool = false
 
+    @Option(
+        name: .customLong("inference-engine-variant"),
+        help: "Engine variant. The default picks one from the model. Batch sizes above 1 need coreai-sequential.")
+    var inferenceEngineVariant: String = "default"
+
+    @Option(
+        name: .customLong("batch-size"),
+        help:
+            "Number of requests to run as one batch. Defaults to the graph's static batch dimension. A dynamic-batch graph reports 1, so set this explicitly. Needs --inference-engine-variant coreai-sequential."
+    )
+    var batchSize: Int?
+
     func validate() throws {
         if promptTokens < 1 { throw ValidationError("--prompt-tokens must be >= 1") }
         if generationTokens < 1 { throw ValidationError("--generation-tokens must be >= 1") }
         if numTrials < 1 { throw ValidationError("--num-trials must be >= 1") }
+        if let b = batchSize, b < 1 { throw ValidationError("--batch-size must be >= 1") }
         if !FileManager.default.fileExists(atPath: model) {
             throw ValidationError("Model path not found: \(model)")
         }
@@ -100,6 +113,7 @@ struct LLMBenchmark: AsyncParsableCommand {
         let resolvedChunkSize = chunkSize ?? bundle.language.prefillChunkSize
         let resolvedChunkThreshold = chunkThreshold ?? bundle.language.prefillChunkThreshold
         let engineOptions = EngineOptions(
+            variant: inferenceEngineVariant,
             prefillChunkSize: resolvedChunkSize,
             prefillChunkThreshold: resolvedChunkThreshold,
             tensorData: bundle.tensorData
@@ -117,11 +131,37 @@ struct LLMBenchmark: AsyncParsableCommand {
         let prompt = randomPrompt(vocabSize: vocabSize, count: promptTokens, seed: seed)
         let sampling = SamplingConfiguration(temperature: 0)
 
+        // A dynamic-batch graph's batch dimension is symbolic and detects as 1, so it needs an
+        // explicit --batch-size.
+        let detectedBatch = (engine as? CoreAISequentialEngine)?.declaredBatchSize ?? 1
+        let effectiveBatch = batchSize ?? detectedBatch
+        if effectiveBatch > bundle.maxBatchSize {
+            throw ValidationError(
+                "--batch-size \(effectiveBatch) exceeds the asset's max batch size \(bundle.maxBatchSize). "
+                    + "Re-export with --dynamic-batch-size >= \(effectiveBatch), or lower --batch-size.")
+        }
+        if effectiveBatch > 1 {
+            guard engine is CoreAISequentialEngine else {
+                throw ValidationError(
+                    "batched benchmarking with batch size \(effectiveBatch) requires "
+                        + "--inference-engine-variant coreai-sequential")
+            }
+            let source = batchSize != nil ? "requested" : "auto-detected from model"
+            print("🧵 Batch size: \(effectiveBatch) (\(source); benchmarking \(effectiveBatch) rows together)")
+        }
+
+        func oneTrial() async throws -> TrialResult {
+            if effectiveBatch > 1, let seq = engine as? CoreAISequentialEngine {
+                return try await runBatchedTrial(engine: seq, prompt: prompt, batchSize: effectiveBatch)
+            }
+            return try await runTrial(engine: engine, prompt: prompt, sampling: sampling)
+        }
+
         // Warmup
         print("\n⚙️  Warming up engine...", terminator: "")
         fflush(stdout)
         let warmupStart = SuspendingClock.now
-        _ = try await runTrial(engine: engine, prompt: prompt, sampling: sampling)
+        _ = try await oneTrial()
         let warmupSeconds = (SuspendingClock.now - warmupStart).inSeconds
         print(" done in \(fmt(warmupSeconds))s")
 
@@ -130,7 +170,7 @@ struct LLMBenchmark: AsyncParsableCommand {
         var trials: [TrialResult] = []
 
         for i in 0..<numTrials {
-            let r = try await runTrial(engine: engine, prompt: prompt, sampling: sampling)
+            let r = try await oneTrial()
             trials.append(r)
             if i > 0 { print() }
             print("🧪 Trial \(i + 1)")
@@ -145,6 +185,9 @@ struct LLMBenchmark: AsyncParsableCommand {
         print(String(repeating: "=", count: 50))
         print("Prepare:    \(fmt(prepareSeconds))s\(cacheSuffix)")
         print("Warmup:     \(fmt(warmupSeconds))s")
+        if effectiveBatch > 1 {
+            print("Batch:      \(effectiveBatch) rows (the prompt and generation rates below are totals across rows)")
+        }
         print("Prompt:     \(fmt(avgPrompt)) tokens/sec")
         print("Generation: \(fmt(avgGen)) tokens/sec")
         print(String(repeating: "=", count: 50))
@@ -155,6 +198,7 @@ struct LLMBenchmark: AsyncParsableCommand {
                 promptTokens: promptTokens,
                 generationTokens: generationTokens,
                 numTrials: numTrials,
+                batchSize: effectiveBatch,
                 prepareSeconds: prepareSeconds,
                 cacheHit: cacheHit,
                 warmupSeconds: warmupSeconds,
@@ -204,6 +248,24 @@ struct LLMBenchmark: AsyncParsableCommand {
         return TrialResult(promptTps: promptTps, genTps: genTps)
     }
 
+    /// Runs `batchSize` copies of the prompt as one batch. Rates are totals across rows. No EOS is
+    /// set, so every row generates `generationTokens` tokens.
+    private func runBatchedTrial(
+        engine: CoreAISequentialEngine,
+        prompt: [Int32],
+        batchSize: Int
+    ) async throws -> TrialResult {
+        try? await Task.sleep(for: .milliseconds(50))
+        let rows = Array(repeating: prompt, count: batchSize)
+        let configs = Array(repeating: SamplingConfiguration(temperature: 0), count: batchSize)
+        let r = try await engine.lockstepBatchedGenerate(
+            promptRows: rows, configs: configs, maxNewTokens: generationTokens, eosTokenIds: [])
+        let generatedPerRow = r.tokens.first?.count ?? 0
+        let promptTps = r.prefillSeconds > 0 ? Double(batchSize * promptTokens) / r.prefillSeconds : 0
+        let genTps = r.decodeSeconds > 0 ? Double(batchSize * generatedPerRow) / r.decodeSeconds : 0
+        return TrialResult(promptTps: promptTps, genTps: genTps)
+    }
+
     // MARK: - Helpers
 
     private func randomPrompt(vocabSize: Int, count: Int, seed: UInt64) -> [Int32] {
@@ -257,6 +319,7 @@ struct BenchmarkReport: Codable {
     let promptTokens: Int
     let generationTokens: Int
     let numTrials: Int
+    let batchSize: Int
     let prepareSeconds: Double
     let cacheHit: Bool
     let warmupSeconds: Double
diff --git a/swift/Tests/CoreAISharedTests/ModelBundleTests.swift b/swift/Tests/CoreAISharedTests/ModelBundleTests.swift
index 4dcc9baf..b2a59bc5 100644
--- a/swift/Tests/CoreAISharedTests/ModelBundleTests.swift
+++ b/swift/Tests/CoreAISharedTests/ModelBundleTests.swift
@@ -68,6 +68,47 @@ struct ModelBundleTests {
         }
     }
 
+    @Test("0.3 bundle decodes and metadataVersion reports the file's version")
+    func decodes03AndPreservesVersion() throws {
+        let url = try Self.tempBundle(
+            """
+            {
+              "metadata_version": "0.3",
+              "kind": "llm",
+              "name": "qwen3-0.6b",
+              "assets": { "main": "model.aimodel" },
+              "language": {
+                "tokenizer": "Qwen/Qwen3-0.6B",
+                "vocab_size": 151936,
+                "max_context_length": 8192
+              }
+            }
+            """, named: "v03")
+        let bundle = try ModelBundle(at: url)
+        #expect(bundle.metadataVersion == "0.3")
+    }
+
+    @Test("Unknown future version throws unsupportedVersion")
+    func unknownFutureVersionThrows() throws {
+        let url = try Self.tempBundle(
+            """
+            {
+              "metadata_version": "0.4",
+              "kind": "llm",
+              "name": "from-the-future",
+              "assets": { "main": "model.aimodel" }
+            }
+            """, named: "v04")
+        let error = #expect(throws: ModelBundle.BundleError.self) {
+            _ = try ModelBundle(at: url)
+        }
+        guard case .unsupportedVersion(let v) = error else {
+            Issue.record("expected unsupportedVersion, got \(String(describing: error))")
+            return
+        }
+        #expect(v == "0.4")
+    }
+
     @Test("Missing metadata.json throws")
     func missingMetadataThrows() throws {
         let dir = FileManager.default.temporaryDirectory.appending(
diff --git a/swift/Tests/LanguageModelsTests/BatchedInputBuilderTests.swift b/swift/Tests/LanguageModelsTests/BatchedInputBuilderTests.swift
new file mode 100644
index 00000000..d2914103
--- /dev/null
+++ b/swift/Tests/LanguageModelsTests/BatchedInputBuilderTests.swift
@@ -0,0 +1,94 @@
+// 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
+import Testing
+
+@testable import CoreAILanguageModels
+
+/// Model-free checks of the batched (N-row) input fills against hand-built tensors. The ragged cases
+/// give each row its own decode cursor.
+@Suite("Batched input builder")
+struct BatchedInputBuilderTests {
+    private let s = causalMaskSentinel
+
+    @Test("token ids are row-major flattened")
+    func tokenIDs() {
+        #expect(BatchedInputBuilder.tokenIDs(rows: [[1, 2, 3], [4, 5, 6]]) == [1, 2, 3, 4, 5, 6])
+    }
+
+    @Test("uniform position ids count 0..<keyLen per row")
+    func uniformPositions() {
+        // N=2 decode at startPos 1: keyLen is 2, so each row is [0, 1].
+        #expect(BatchedInputBuilder.positionIDs(rowCount: 2, startPos: 1, queryLen: 1) == [0, 1, 0, 1])
+        // N=1 prefill of 3 tokens from startPos 0.
+        #expect(BatchedInputBuilder.positionIDs(rowCount: 1, startPos: 0, queryLen: 3) == [0, 1, 2])
+    }
+
+    @Test("uniform causal mask matches the hand-built reference tensor for N=1 and N=2")
+    func uniformMask() {
+        // Decode at startPos 1 (keyLen 2): the query attends both keys.
+        #expect(BatchedInputBuilder.additiveCausalMask(rowCount: 1, startPos: 1, queryLen: 1) == [0, 0])
+        // Prefill of 2 tokens from startPos 0: [[0, sentinel], [0, 0]].
+        #expect(BatchedInputBuilder.additiveCausalMask(rowCount: 1, startPos: 0, queryLen: 2) == [0, s, 0, 0])
+        // With N=2, each row repeats the N=1 block.
+        #expect(
+            BatchedInputBuilder.additiveCausalMask(rowCount: 2, startPos: 0, queryLen: 2)
+                == [0, s, 0, 0, 0, s, 0, 0])
+    }
+
+    // MARK: - Ragged decode (per-row cursors)
+
+    @Test("ragged decode positions: real slots keep their index, the write slot gets the cursor")
+    func raggedPositions() {
+        // Cursors [2, 5] give keyLen 6 with the write slot at 5.
+        // Row 0: real slots 0-1, gap 2-4 pinned to 0, write slot holds position 2.
+        // Row 1: real slots 0-4, write slot holds position 5.
+        #expect(
+            BatchedInputBuilder.raggedDecodePositionIDs(cursors: [2, 5])
+                == [0, 1, 0, 0, 0, 2] + [0, 1, 2, 3, 4, 5])
+    }
+
+    @Test("ragged decode mask: attend real keys and the shared write slot, mask the gap")
+    func raggedMask() {
+        // Cursors [2, 5]: row 0 masks the gap at slots 2-4, and row 1 gets the full causal mask.
+        #expect(
+            BatchedInputBuilder.raggedDecodeMask(cursors: [2, 5])
+                == [0, 0, s, s, s, 0] + [0, 0, 0, 0, 0, 0])
+    }
+
+    @Test("each ragged row matches a single-row decode at its cursor (positions and mask)")
+    func raggedMatchesPerRowSingleBatch() {
+        let cursors = [1, 4, 7]
+        let maxCursor = cursors.max()!
+        let keyLen = maxCursor + 1
+        let pos = BatchedInputBuilder.raggedDecodePositionIDs(cursors: cursors)
+        let mask = BatchedInputBuilder.raggedDecodeMask(cursors: cursors)
+        for (b, c) in cursors.enumerated() {
+            // The new token sits at slot maxCursor instead of c, so compare the positions of the
+            // attended keys, which must be 0...c as in a single-row decode.
+            var attendedPositions: [Int32] = []
+            for key in 0..<keyLen where mask[b * keyLen + key] == 0 {
+                attendedPositions.append(pos[b * keyLen + key])
+            }
+            #expect(
+                attendedPositions == (0...Int32(c)).map { $0 },
+                "row \(b) (cursor \(c)) must attend exactly positions 0...\(c)")
+        }
+    }
+
+    @Test("equal cursors reproduce the uniform decode tensors")
+    func raggedEqualCursorsMatchUniform() {
+        let c = 3
+        let rows = 4
+        #expect(
+            BatchedInputBuilder.raggedDecodePositionIDs(cursors: Array(repeating: c, count: rows))
+                == BatchedInputBuilder.positionIDs(rowCount: rows, startPos: c, queryLen: 1))
+        #expect(
+            BatchedInputBuilder.raggedDecodeMask(cursors: Array(repeating: c, count: rows))
+                == BatchedInputBuilder.additiveCausalMask(rowCount: rows, startPos: c, queryLen: 1))
+    }
+}
diff --git a/swift/Tests/LanguageModelsTests/BatchedSamplerTests.swift b/swift/Tests/LanguageModelsTests/BatchedSamplerTests.swift
new file mode 100644
index 00000000..80255ec0
--- /dev/null
+++ b/swift/Tests/LanguageModelsTests/BatchedSamplerTests.swift
@@ -0,0 +1,58 @@
+// 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 Foundation
+import Testing
+
+@testable import CoreAILanguageModels
+
+/// Each row samples with its own ``SamplingConfiguration`` and step, so greedy and seeded rows can
+/// share a batch.
+@Suite("Batched Sampler Tests")
+struct BatchedSamplerTests {
+    /// A logit row of length `vocab` with a single clear argmax at `peak`.
+    private func peakedRow(vocab: Int, peak: Int) -> [LogitsScalarType] {
+        var row = [LogitsScalarType](repeating: -10, count: vocab)
+        row[peak] = 10
+        return row
+    }
+
+    private let noHistory = ArraySlice<Int32>()
+
+    @Test("N=2, both greedy: each row samples its own argmax")
+    func batchedGreedyPerRow() {
+        let rows = [peakedRow(vocab: 32, peak: 3), peakedRow(vocab: 32, peak: 7)]
+        let tokens = BatchedSampler.sample(
+            rows: rows,
+            configurations: [.greedy, .greedy],
+            histories: [noHistory, noHistory],
+            steps: [nil, nil])
+        #expect(tokens == [3, 7])
+    }
+
+    @Test("N=2, greedy and seeded temperature rows: greedy is deterministic, seeded is reproducible")
+    func batchedMixedSampling() {
+        let rows = [peakedRow(vocab: 64, peak: 2), peakedRow(vocab: 64, peak: 50)]
+        let configs: [SamplingConfiguration] = [
+            .greedy,
+            SamplingConfiguration(temperature: 0.7, seed: 42),
+        ]
+        let histories = [noHistory, noHistory]
+        let steps: [Int?] = [nil, 5]
+
+        let first = BatchedSampler.sample(
+            rows: rows, configurations: configs, histories: histories, steps: steps)
+        let second = BatchedSampler.sample(
+            rows: rows, configurations: configs, histories: histories, steps: steps)
+
+        // Greedy row 0 returns its argmax.
+        #expect(first[0] == 2)
+        #expect(second[0] == 2)
+        // Seeded row 1 repeats across identical calls and stays in the vocab range.
+        #expect(first[1] == second[1])
+        #expect(first[1] >= 0 && first[1] < 64)
+    }
+}
diff --git a/swift/Tests/LanguageModelsTests/IdempotentEngineTests.swift b/swift/Tests/LanguageModelsTests/IdempotentEngineTests.swift
index 4524884e..19997e56 100644
--- a/swift/Tests/LanguageModelsTests/IdempotentEngineTests.swift
+++ b/swift/Tests/LanguageModelsTests/IdempotentEngineTests.swift
@@ -168,6 +168,41 @@ struct IdempotentEngineRoutingTests {
     }
 }
 
+// MARK: - Per-session generation state (model-free)
+
+/// The token box and `lastPrefixHitCount` live on `GenerationSessionState`, so sessions sharing one
+/// engine cannot cancel or overwrite each other.
+@Suite("Per-session generation state")
+struct PerSessionStateTests {
+    private func makeSession() -> GenerationSessionState {
+        GenerationSessionState(
+            kvCache: MockStateHandler(names: ["k"], shape: [1, 1, 1, 1, 1]),
+            additionalStates: nil,
+            hasNonTruncatableStates: false)
+    }
+
+    @Test("each session owns an independent in-flight token box")
+    func independentTokenBoxes() {
+        let a = makeSession()
+        let b = makeSession()
+        let tokenA = GenerationToken()
+        a.tokenBox.install(tokenA)
+        #expect(a.tokenBox.isBusy)
+        #expect(!b.tokenBox.isBusy)  // installing on A must not make B busy
+        b.tokenBox.cancelActive()
+        #expect(a.tokenBox.isBusy)  // cancelling B must not clear A's token
+    }
+
+    @Test("each session tracks its own prefix hit count")
+    func independentPrefixHitCount() {
+        let a = makeSession()
+        let b = makeSession()
+        a.lastPrefixHitCount = 7
+        #expect(a.lastPrefixHitCount == 7)
+        #expect(b.lastPrefixHitCount == 0)
+    }
+}
+
 // MARK: - Parity Scaffold (Phase 2 — requires an on-device compiled model)
 
 /// Model-backed parity checks for the idempotent-engine refactor.
diff --git a/swift/Tests/LanguageModelsTests/InputHandlerHarnessTests.swift b/swift/Tests/LanguageModelsTests/InputHandlerHarnessTests.swift
index 4fba6afd..4fc62a63 100644
--- a/swift/Tests/LanguageModelsTests/InputHandlerHarnessTests.swift
+++ b/swift/Tests/LanguageModelsTests/InputHandlerHarnessTests.swift
@@ -298,6 +298,29 @@ struct InputHandlerFuzzTests {
         }
     }
 
+    @Test(
+        "Additive causal attn_mask [1,1,query,key] correctness",
+        arguments: [
+            (4, 0),  // prefill from the start
+            (1, 0),  // single-token prompt
+            (1, 5),  // decode step after 5 processed tokens
+            (8, 100),  // prefill chunk mid-sequence
+        ] as [(Int, Int)])
+    func additiveCausalMaskFill(queryLen: Int, processedTokenCount: Int) {
+        let keyLen = processedTokenCount + queryLen
+        var mask = NDArray(shape: [1, 1, queryLen, keyLen], scalarType: .float16)
+
+        fillAdditiveCausalMask(&mask, processedTokenCount: processedTokenCount)
+
+        let result = readNDArray(mask, as: Float16.self, count: queryLen * keyLen)
+        for query in 0..<queryLen {
+            for key in 0..<keyLen {
+                let attends = key <= processedTokenCount + query
+                #expect(result[query * keyLen + key] == (attends ? 0 : causalMaskSentinel), "q=\(query) k=\(key)")
+            }
+        }
+    }
+
     @Test("Rapid re-fill does not corrupt")
     func rapidRefill() {
         var array = NDArray(shape: [1, 256], scalarType: .float16)
diff --git a/swift/Tests/LanguageModelsTests/LanguageModelBundleTests.swift b/swift/Tests/LanguageModelsTests/LanguageModelBundleTests.swift
index 4511e450..02e0cba1 100644
--- a/swift/Tests/LanguageModelsTests/LanguageModelBundleTests.swift
+++ b/swift/Tests/LanguageModelsTests/LanguageModelBundleTests.swift
@@ -48,6 +48,28 @@ struct LanguageModelBundleTests {
         #expect(bundle.language.embeddedTokenizer == true)  // default
         #expect(bundle.language.functionMap?.names(for: "extend") == ["extend_512", "extend_1024"])
         #expect(bundle.language.functionMap?.name(for: "main") == "main")
+        #expect(bundle.maxBatchSize == 1)  // default when max_batch_size is absent
+    }
+
+    @Test("max_batch_size decodes from the language block")
+    func decodesMaxBatchSize() throws {
+        let url = try Self.tempBundle(
+            """
+            {
+              "metadata_version": "0.3",
+              "kind": "llm",
+              "name": "qwen3-1.7b-batched",
+              "assets": { "main": "model.aimodel" },
+              "language": {
+                "tokenizer": "Qwen/Qwen3-1.7B",
+                "vocab_size": 151936,
+                "max_context_length": 8192,
+                "max_batch_size": 8
+              }
+            }
+            """)
+        let bundle = try LanguageModelBundle(at: url)
+        #expect(bundle.maxBatchSize == 8)
     }
 
     @Test("0.1 legacy bundle throws unsupportedVersion")
diff --git a/swift/Tests/LanguageModelsTests/StateHandlerTests.swift b/swift/Tests/LanguageModelsTests/StateHandlerTests.swift
index d03fae56..75ce0259 100644
--- a/swift/Tests/LanguageModelsTests/StateHandlerTests.swift
+++ b/swift/Tests/LanguageModelsTests/StateHandlerTests.swift
@@ -237,6 +237,26 @@ struct StateHandlerConformanceTests {
     }
 }
 
+// MARK: - KV Cache Shape Resolution (batch-aware)
+
+@Suite("KV cache dynamic-shape resolution")
+struct KVCacheShapeResolutionTests {
+    /// Layout `[layers, batch, kv_heads, seq, head_dim]`: the seq dim grows to `sequenceLength`
+    /// and a dynamic batch dim resolves to `batchSize`.
+    @Test(
+        "resolves the seq dim and a dynamic batch dim",
+        arguments: [
+            ([2, -1, 4, -1, 8], 3, [2, 3, 4, 256, 8]),
+            ([2, 1, 4, -1, 8], 1, [2, 1, 4, 256, 8]),
+            ([2, -1, 4, -1, 8], 1, [2, 1, 4, 256, 8]),
+        ] as [([Int], Int, [Int])])
+    func resolves(shape: [Int], batchSize: Int, resolved: [Int]) {
+        #expect(
+            GrowingNDArrayState.resolveKVShape(
+                shape, sequenceDimIndex: 3, sequenceLength: 256, batchSize: batchSize) == resolved)
+    }
+}
+
 // MARK: - withBoundStates Tests
 
 /// Minimal state handler for testing the binding API.
@@ -328,3 +348,120 @@ struct StateKindValidationTests {
         #expect(StateHandlerFactory.dynamicStatesWithFixedSizeKind(valid, isDynamic: isDynamic).isEmpty)
     }
 }
+
+// MARK: - Per-row KV copies
+
+/// Batch layouts the engine sees: batch leading (`[B, H, seq, D]`) and batch at dim 1
+/// (`[G, B, H, seq, D]`), where code that assumes a leading batch dim copies the wrong slab.
+enum RowCopyLayout: CaseIterable, Sendable {
+    case leadingBatch, nonLeadingBatch
+
+    var batchDimIndex: Int { self == .leadingBatch ? 0 : 1 }
+    var sequenceDimIndex: Int { self == .leadingBatch ? 2 : 3 }
+
+    /// Shape with `batch` rows and `seq` slots.
+    func shape(batch: Int, seq: Int) -> [Int] {
+        self == .leadingBatch ? [batch, 2, seq, 3] : [2, batch, 2, seq, 3]
+    }
+}
+
+@Suite("Per-row KV copies")
+struct RowCopyTests {
+    /// An array whose bytes count up, so every element has a distinct bit pattern.
+    private func countingArray(_ shape: [Int], _ type: NDArray.ScalarType, offset: Int = 0) -> NDArray {
+        var array = NDArray(shape: shape, scalarType: type)
+        let byteCount = shape.reduce(1, *) * type.byteSize
+        array.mutableRawView().withUnsafeMutableBytes { ptr, _, _ in
+            let bytes = ptr.assumingMemoryBound(to: UInt8.self)
+            for i in 0..<byteCount { bytes[i] = UInt8((i + offset) % 251) }
+        }
+        return array
+    }
+
+    private func bytes(_ array: NDArray) -> [UInt8] {
+        let byteCount = array.shape.reduce(1, *) * array.scalarType.byteSize
+        var out = [UInt8](repeating: 0, count: byteCount)
+        array.rawView().withUnsafeBytes { ptr, _, _ in
+            let p = ptr.assumingMemoryBound(to: UInt8.self)
+            for i in 0..<byteCount { out[i] = p[i] }
+        }
+        return out
+    }
+
+    /// Reference copy: every destination element for which `source(index)` returns a source index
+    /// takes that source element's bytes.
+    private func expected(
+        destination: NDArray, from source: NDArray, source sourceIndex: ([Int]) -> [Int]?
+    ) -> [UInt8] {
+        var out = bytes(destination)
+        let src = bytes(source)
+        let width = destination.scalarType.byteSize
+        func flat(_ index: [Int], _ shape: [Int]) -> Int {
+            zip(index, shape).reduce(0) { $0 * $1.1 + $1.0 }
+        }
+        let shape = destination.shape
+        for element in 0..<shape.reduce(1, *) {
+            var index = [Int](repeating: 0, count: shape.count)
+            var rest = element
+            for dim in shape.indices.reversed() {
+                index[dim] = rest % shape[dim]
+                rest /= shape[dim]
+            }
+            guard let from = sourceIndex(index) else { continue }
+            let s = flat(from, source.shape) * width
+            out.replaceSubrange(element * width..<(element + 1) * width, with: src[s..<s + width])
+        }
+        return out
+    }
+
+    @Test(
+        "copyRegion moves one row's slot and leaves everything else untouched",
+        arguments: [NDArray.ScalarType.float16, .float32, .int8], RowCopyLayout.allCases)
+    func copyRegion(_ type: NDArray.ScalarType, _ layout: RowCopyLayout) {
+        var array = countingArray(layout.shape(batch: 3, seq: 6), type)
+        let (b, seq) = (layout.batchDimIndex, layout.sequenceDimIndex)
+        // Row 1's new token moves from the shared write slot 5 down to its cursor 2.
+        let want = expected(destination: array, from: array) { index in
+            guard index[b] == 1, index[seq] == 2 else { return nil }
+            var from = index
+            from[seq] = 5
+            return from
+        }
+
+        GrowingNDArrayState.copyRegion(
+            in: &array, sequenceDimIndex: seq, batchDimIndex: b, row: 1, fromSeq: 5, toSeq: 2)
+
+        #expect(bytes(array) == want)
+    }
+
+    @Test("copyRegion with fromSeq == toSeq is a no-op")
+    func copyRegionSameSlot() {
+        var array = countingArray([2, 1, 4, 2], .float16)
+        let before = bytes(array)
+        GrowingNDArrayState.copyRegion(
+            in: &array, sequenceDimIndex: 2, batchDimIndex: 0, row: 0, fromSeq: 3, toSeq: 3)
+        #expect(bytes(array) == before)
+    }
+
+    @Test(
+        "copyRowPrefix copies a batch-1 prefix into one row across different seq capacities",
+        arguments: [NDArray.ScalarType.float16, .float32, .int8], RowCopyLayout.allCases)
+    func copyRowPrefix(_ type: NDArray.ScalarType, _ layout: RowCopyLayout) {
+        let source = countingArray(layout.shape(batch: 1, seq: 4), type, offset: 100)
+        var destination = countingArray(layout.shape(batch: 3, seq: 6), type)
+        let (b, seq) = (layout.batchDimIndex, layout.sequenceDimIndex)
+        // Slots 0..<3 of row 2 come from the source; everything else keeps its bytes.
+        let want = expected(destination: destination, from: source) { index in
+            guard index[b] == 2, index[seq] < 3 else { return nil }
+            var from = index
+            from[b] = 0
+            return from
+        }
+
+        GrowingNDArrayState.copyRowPrefix(
+            from: source, sourceRow: 0, to: &destination, destRow: 2,
+            sequenceDimIndex: seq, batchDimIndex: b, count: 3)
+
+        #expect(bytes(destination) == want)
+    }
+}