From 0815f680d6d07021a714944cdade62605236a8ef Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Sat, 12 Sep 2026 17:03:07 -0700 Subject: [PATCH 1/2] Fix guided generation refused after engine load on static-shape engines isGuidedGenerationSupported checked loadedEngineIsConstrainedCapable first and returned immediately once non-nil, so the loadedEngineSupportsLogits fallback never ran after an engine loaded. StaticShapeEngine (chunked-static iOS exports) supports logits but isn't ConstrainedGenerationCapable, so guided generation was reported unsupported after the first request even though the executor's guard would have allowed it. Moved the check into InferenceEngine.supportsGuidedGeneration, used by both the executor guard and capability reporting. Fixes apple/coreai-models#247 --- .../InferenceEngines/InferenceEngine.swift | 9 +++++++++ .../LanguageModel/CoreAILanguageModel.swift | 13 ++++--------- .../LanguageModel/ModelResources.swift | 16 ++++------------ .../ModelResourcesTests.swift | 19 +++++++++++++++++++ 4 files changed, 36 insertions(+), 21 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index 5b0a3f1e..bd58dfd5 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -193,6 +193,15 @@ extension InferenceEngine { public var supportsLogits: Bool { false } } +extension InferenceEngine { + /// Guided/structured generation needs either per-step logits (CPU-side + /// constrained decoding) or GPU-side constrained sampling + /// (`ConstrainedGenerationCapable`). + public var supportsGuidedGeneration: Bool { + supportsLogits || self is any ConstrainedGenerationCapable + } +} + extension InferenceEngine { /// Default: no prefix hits (engine doesn't track history). public var lastPrefixHitCount: Int { 0 } diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift index 95d3fd28..3a8a8174 100644 --- a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift +++ b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift @@ -199,15 +199,10 @@ public struct CoreAILanguageModel: LanguageModel { resources.unloadResources() } - /// Whether guided generation is available for this model. + /// Whether guided generation is available for this model. Assumes yes + /// before an engine is loaded. private var isGuidedGenerationSupported: Bool { - if let isConstrainedCapable = resources.loadedEngineIsConstrainedCapable { - return isConstrainedCapable - } - if let supportsLogits = resources.loadedEngineSupportsLogits { - return supportsLogits - } - return true + resources.loadedEngineSupportsGuidedGeneration ?? true } // MARK: - Executor @@ -283,7 +278,7 @@ public struct CoreAILanguageModel: LanguageModel { // Check if guided generation is requested if let schema = request.schema { - guard engine.supportsLogits || engine is any ConstrainedGenerationCapable else { + guard engine.supportsGuidedGeneration else { throw LanguageModelError.unsupportedCapability( .init( capability: .guidedGeneration, diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift index bc050652..7d5c4aa3 100644 --- a/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift +++ b/swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift @@ -36,18 +36,10 @@ final class ModelResources: ResourceManaging { var isLoaded: Bool { state.withLock { $0.loaded != nil } } - /// `supportsLogits` of the resident engine, or `nil` when nothing is loaded. - /// Used only for best-effort capability reporting before a load. - var loadedEngineSupportsLogits: Bool? { - state.withLock { $0.loaded?.supportsLogits } - } - - /// Whether the loaded engine supports GPU-side constrained generation, or `nil` when unloaded. - var loadedEngineIsConstrainedCapable: Bool? { - state.withLock { engine in - guard let loaded = engine.loaded else { return nil } - return loaded is any ConstrainedGenerationCapable - } + /// Guided-generation capability of the resident engine, or `nil` when nothing + /// is loaded. Mirrors the check the executor enforces per-request. + var loadedEngineSupportsGuidedGeneration: Bool? { + state.withLock { $0.loaded?.supportsGuidedGeneration } } /// Returns the engine, loading it on first use. Concurrent callers share one diff --git a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift index 334390c5..c3bd6338 100644 --- a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift +++ b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift @@ -107,6 +107,25 @@ struct ModelResourcesTests { #expect(resources.isLoaded) } + @Test( + "loadedEngineSupportsGuidedGeneration stays true after loading a logits-only engine" + ) + func guidedGenerationSupportedForLogitsOnlyEngineAfterLoad() async throws { + // MockEngine(vocabSize: 100) mirrors StaticShapeEngine: supportsLogits == true, + // but it does not conform to ConstrainedGenerationCapable. + let resources = ModelResources { MockEngine(vocabSize: 100) } + + // Nothing loaded yet — best-effort reporting defaults to "supported". + #expect(resources.loadedEngineSupportsGuidedGeneration == nil) + + _ = try await resources.engine() + + // Regression check: before the fix this incorrectly returned `false` once + // an engine was loaded, because `loadedEngineIsConstrainedCapable` (false) + // shadowed `loadedEngineSupportsLogits` (true). + #expect(resources.loadedEngineSupportsGuidedGeneration == true) + } + @Test("unloadResources during an active borrow defers teardown until it finishes") func unloadDeferredDuringActiveBorrow() async throws { let resources = ModelResources { MockEngine() } From 8667ee607d4154098981f6b2fda4b5f77fa82f91 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Mon, 14 Sep 2026 11:06:46 -0700 Subject: [PATCH 2/2] Add constrained-only branch test for supportsGuidedGeneration Covers the second operand of the capability OR: an engine that is ConstrainedGenerationCapable with supportsLogits == false must still report guided generation supported. Reuses MockConstrainedEngine. --- .../ModelResourcesTests.swift | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift index c3bd6338..02e00801 100644 --- a/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift +++ b/swift/Tests/LanguageModelsTests/ModelResourcesTests.swift @@ -126,6 +126,23 @@ struct ModelResourcesTests { #expect(resources.loadedEngineSupportsGuidedGeneration == true) } + @Test( + "loadedEngineSupportsGuidedGeneration is true for a constrained-capable engine without logits" + ) + func guidedGenerationSupportedForConstrainedOnlyEngineAfterLoad() async throws { + // MockConstrainedEngine conforms to ConstrainedGenerationCapable but has + // supportsLogits == false — the other branch of the capability OR. Guided + // generation must still be reported supported via GPU-side constrained + // sampling. + let resources = ModelResources { MockConstrainedEngine(scriptedTokens: []) } + + #expect(resources.loadedEngineSupportsGuidedGeneration == nil) + + _ = try await resources.engine() + + #expect(resources.loadedEngineSupportsGuidedGeneration == true) + } + @Test("unloadResources during an active borrow defers teardown until it finishes") func unloadDeferredDuringActiveBorrow() async throws { let resources = ModelResources { MockEngine() }