From 7e2db01edec9d2ca0405cae7ca96f1321b4e0162 Mon Sep 17 00:00:00 2001
From: sukru tikves <sukrutikves@gmail.com>
Date: Tue, 22 Sep 2026 07:16:37 -0700
Subject: [PATCH] =?UTF-8?q?Accept=202=E2=80=934=20states=20in=20VLM=20sequ?=
 =?UTF-8?q?ential=20engine=20for=20hybrid=20decoders=20(#276)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The LLM decoder guard rejected any function with a state count other than
exactly two, blocking hybrid (conv/recurrent) 4-state decoders even though
the shared state-handler machinery already supports them. Relax it to accept
2–4 states, matching CoreAISequentialEngine.

(cherry picked from commit 17febaa2a5a8c25535dacd8f93a9610431cdd1e3)
---
 .../CoreAISequentialVLMEngine.swift           | 17 +++++---
 .../CoreAISequentialVLMEngineTests.swift      | 40 +++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 swift/Tests/LanguageModelsTests/CoreAISequentialVLMEngineTests.swift

diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift
index 46353fda..de2d095c 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift
@@ -136,6 +136,17 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
 
     // MARK: - Init
 
+    /// Validates the LLM decoder's state-name contract. Accepts 2–4 states (KV cache + optional
+    /// persistent conv/recurrent states), matching `CoreAISequentialEngine`; hybrid decoders carry
+    /// extra conv/recurrent states beyond the two KV-cache states.
+    static func validateLLMStateNames(_ stateNames: [String]) throws {
+        guard stateNames.count >= 2 && stateNames.count <= 4 else {
+            throw InferenceRuntimeError.invalidOutputType(
+                "VLM LLM function expected 2–4 states (KV cache + optional persistent states), "
+                    + "got \(stateNames.count): \(stateNames)")
+        }
+    }
+
     /// Initialize the VLM engine with separate model assets for vision, embed, and LLM.
     ///
     /// - Parameters:
@@ -232,11 +243,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
                 "VLM LLM function expected 2 inputs (in_embeddings, position_ids), "
                     + "got \(llmDesc.inputNames.count): \(llmDesc.inputNames)")
         }
-        guard llmDesc.stateNames.count == 2 else {
-            throw InferenceRuntimeError.invalidOutputType(
-                "VLM LLM function expected 2 states (KV cache), "
-                    + "got \(llmDesc.stateNames.count): \(llmDesc.stateNames)")
-        }
+        try Self.validateLLMStateNames(llmDesc.stateNames)
         guard llmDesc.outputNames.count >= 1 else {
             throw InferenceRuntimeError.invalidOutputType(
                 "VLM LLM function expected at least 1 output (logits), "
diff --git a/swift/Tests/LanguageModelsTests/CoreAISequentialVLMEngineTests.swift b/swift/Tests/LanguageModelsTests/CoreAISequentialVLMEngineTests.swift
new file mode 100644
index 00000000..4f701c21
--- /dev/null
+++ b/swift/Tests/LanguageModelsTests/CoreAISequentialVLMEngineTests.swift
@@ -0,0 +1,40 @@
+// 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
+import Testing
+
+@testable import CoreAILanguageModels
+
+@Suite("CoreAISequentialVLMEngine state-name contract")
+struct CoreAISequentialVLMEngineTests {
+    @Test("accepts a full-attention two-state KV decoder")
+    func acceptsTwoStates() throws {
+        try CoreAISequentialVLMEngine.validateLLMStateNames(["k_cache", "v_cache"])
+    }
+
+    @Test("accepts a hybrid four-state decoder (conv + recurrent states)")
+    func acceptsFourStates() throws {
+        // Hybrid decoders carry conv/recurrent states beyond the two KV-cache states; the
+        // engine previously rejected anything but exactly two states.
+        try CoreAISequentialVLMEngine.validateLLMStateNames(
+            ["k_cache", "v_cache", "conv_states", "recurrent_states"])
+    }
+
+    @Test("rejects fewer than two states")
+    func rejectsOneState() {
+        #expect(throws: InferenceRuntimeError.self) {
+            try CoreAISequentialVLMEngine.validateLLMStateNames(["k_cache"])
+        }
+    }
+
+    @Test("rejects more than four states")
+    func rejectsFiveStates() {
+        #expect(throws: InferenceRuntimeError.self) {
+            try CoreAISequentialVLMEngine.validateLLMStateNames(
+                ["s0", "s1", "s2", "s3", "s4"])
+        }
+    }
+}