VLM: fix EmbeddedInput shape handling, error recovery, and parallel loading (#79)

* VLM: fix tokenCount for 2D tensors, replace precondition with throw

- EmbeddedInput.tokenCount: handle 2D [seq_len, hidden_dim] vs 3D [batch, seq_len, hidden_dim]
- scatterMerge: replace precondition(float16) with guard/throw for bfloat16 compatibility

* Address review: remove premature maxTiles flag, centralize embedding shape handling

- Remove --max-tiles CLI option (tiling model not implemented yet)
- Move seq_len extraction into EmbeddedInput.seqLen(of:) so scatterMerge
  and other call sites use one canonical shape resolution path
- tokenCount is now a stored property computed once at init

* Simplify EmbeddedInput: assume 3D [batch, seq_len, hidden_dim] layout

All current VLM models produce 3D embeddings. Drop the 2D fallback
and the seqLen helper -- tokenCount is just shape[1].

* EmbeddedInput: validate exactly 3D shape at init, throw on mismatch

* Fix test: EmbeddedInput.init is now throwing

* Fix test: mark embeddedInputBasics as throws

Author: stikves

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift

@@ -385,7 +385,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
 
         CLILogger.log("VLM encodeImage complete: \(tokenCount) embedding tokens")
 
-        return EmbeddedInput(
+        return try EmbeddedInput(
             embeddings: projectedEmbeddings,
             embeddingPositions: placeholderRange
         )
@@ -560,8 +560,8 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
                     + "expected \(imageTokenCount) from config. Check prompt template.")
         }
 
-        let seqLen = textEmbeddings.shape.count >= 2 ? textEmbeddings.shape[1] : 0
-        let imgSeqLen = imageEmbeddings.shape.count >= 2 ? imageEmbeddings.shape[1] : 0
+        let seqLen = textEmbeddings.shape[1]
+        let imgSeqLen = imageEmbeddings.shape[1]
         guard imgSeqLen >= imageTokenCount else {
             throw InferenceRuntimeError.invalidArgument(
                 "scatterMerge: image embeddings have \(imgSeqLen) tokens, need \(imageTokenCount)")
@@ -574,10 +574,10 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
         }
 
         // Copy image embeddings into placeholder positions.
-        precondition(
-            imageEmbeddings.scalarType == .float16,
-            "scatterMerge only supports float16 embeddings; got \(imageEmbeddings.scalarType)"
-        )
+        guard imageEmbeddings.scalarType == .float16 else {
+            throw InferenceRuntimeError.invalidInputType(
+                "scatterMerge only supports float16 embeddings; got \(imageEmbeddings.scalarType)")
+        }
         imageEmbeddings.view(as: Float16.self).withUnsafePointer { imgPtr, _, _ in
             var mutableView = merged.mutableView(as: Float16.self)
             mutableView.withUnsafeMutablePointer { mergedPtr, _, _ in

swift/Sources/CoreAILanguageModels/InferenceEngines/EmbeddedInput.swift

@@ -12,22 +12,25 @@ import Foundation
 /// language model. The engine performs scatter-merge: replacing placeholder
 /// token positions with these embeddings before the first forward pass.
 public struct EmbeddedInput: Sendable {
-    /// The embedding tensor, typically shape [1, seq_len, hidden_dim].
+    /// The embedding tensor, shape [batch, seq_len, hidden_dim].
     /// Scalar type matches the LLM's expected input (float16, bFloat16, etc.).
     public let embeddings: NDArray
 
     /// Positions in the token sequence where embeddings replace placeholders.
     public let embeddingPositions: Range<Int>
 
-    public init(embeddings: NDArray, embeddingPositions: Range<Int>) {
+    public init(embeddings: NDArray, embeddingPositions: Range<Int>) throws {
+        guard embeddings.shape.count == 3 else {
+            throw InferenceRuntimeError.invalidArgument(
+                "EmbeddedInput requires 3D embeddings [batch, seq_len, hidden_dim], "
+                    + "got shape with \(embeddings.shape.count) dimensions")
+        }
         self.embeddings = embeddings
         self.embeddingPositions = embeddingPositions
     }
 
     /// Number of embedding tokens (seq_len dimension).
-    public var tokenCount: Int {
-        embeddings.shape.count >= 2 ? embeddings.shape[1] : 0
-    }
+    public var tokenCount: Int { embeddings.shape[1] }
 
     // TODO: Multi-turn support — allow multiple image regions per input,
     // persistent across generate() calls (keep in KV cache on reset).

swift/Sources/Tools/llm-runner/LLMRunnerMain.swift

@@ -371,6 +371,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             )
             let vlmConfig = VLMModelConfig(base: baseConfig, visionConfig: visionConfig)
 
+            // Sequential to avoid runtime errors with concurrent model preparation.
             let visionModel = try await PreparedModel.prepare(at: visionURL)
             let embedModel = try await PreparedModel.prepare(at: embedURL)
             let llmModel = try await PreparedModel.prepare(at: mainURL)

swift/Tests/LanguageModelsTests/VLMProtocolTests.swift

@@ -16,12 +16,12 @@ import CoreAI
 struct MultimodalTypeTests {
     #if canImport(CoreAI)
     @Test("EmbeddedInput wraps NDArray with positions")
-    func embeddedInputBasics() {
+    func embeddedInputBasics() throws {
         let embeddings = NDArray(
             shape: [1, 256, 2048],
             scalarType: .float16
         )
-        let input = EmbeddedInput(
+        let input = try EmbeddedInput(
             embeddings: embeddings,
             embeddingPositions: 5..<261
         )