Merge branch 'main' into sukru/fix-offline-tests

Author: stikves

python/src/coreai_models/diffusion/pipeline.py

@@ -162,8 +162,6 @@ def _load_hf_pipeline(model_id: str, pipeline_type: str, model_dtype: torch.dtyp
         from diffusers import Flux2KleinPipeline
 
         hf_pipe = Flux2KleinPipeline.from_pretrained(model_id, torch_dtype=model_dtype)
-        # Text encoder needs float32 for token embedding precision
-        hf_pipe.text_encoder = hf_pipe.text_encoder.float()
         return hf_pipe
 
     if pipeline_type == "sd3":

swift/Sources/CoreAILanguageModels/GuidedGeneration/ConstrainedGenerationSession.swift

@@ -17,6 +17,8 @@ import Tokenizers
 /// Each session is tied to a specific JSON schema and vocabulary. It tracks
 /// the generation state and produces token masks that enforce schema compliance.
 public struct ConstrainedGenerationSession: ~Copyable {
+    static let maxRollbackTokens = 64
+
     private let tokenizerInfo: TokenizerInfo
     private let compiler: GrammarCompiler
     private let compiledGrammar: CompiledGrammar
@@ -93,7 +95,7 @@ public struct ConstrainedGenerationSession: ~Copyable {
         self.tokenizerInfo = tokenizerInfo
         self.compiler = GrammarCompiler(tokenizerInfo: tokenizerInfo)
         self.compiledGrammar = try compiler.compileJSONSchema(jsonSchema)
-        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar)
+        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar, maxRollbackTokens: Self.maxRollbackTokens)
         self.vocabularySize = tokenizerInfo.vocabularySize
         self.bitmaskSize = (vocabularySize + 31) / 32
         self.bitmaskBuffer = Array(repeating: 0, count: bitmaskSize)
@@ -193,6 +195,50 @@ public struct ConstrainedGenerationSession: ~Copyable {
         matcher.reset()
         allTokensBlocked = false
     }
+
+    /// Rollback the grammar state by N tokens. Returns false if rollback failed
+    /// (e.g., exceeds maxRollbackTokens budget).
+    @discardableResult
+    public mutating func rollback(_ numTokens: Int = 1) -> Bool {
+        guard numTokens >= 0 else { return false }
+        return matcher.rollback(numTokens)
+    }
+
+    /// Find the longest deterministic string from the current grammar state.
+    /// Does not change the matcher state. Returns nil if no jump-forward is possible.
+    public func findJumpForwardString() -> String? {
+        matcher.findJumpForwardString()
+    }
+
+    /// Result of filling a bitmask for the next token.
+    public enum BitmaskResult: Equatable {
+        /// Grammar is terminated or all tokens are blocked — generation should stop.
+        case terminated
+        /// All tokens are allowed — no mask needed, generate unconstrained.
+        case unconstrained
+        /// Bitmask was written — apply it to constrain sampling.
+        case constrained
+    }
+
+    /// Fill the bitmask directly into a caller-provided buffer (e.g., a GPU-visible MTLBuffer).
+    ///
+    /// The caller must ensure the pointer has room for at least `(vocabularySize + 31) / 32`
+    /// Int32 words.
+    public mutating func fillBitmask(into pointer: UnsafeMutablePointer<Int32>) -> BitmaskResult {
+        if isTerminated { return .terminated }
+
+        let needsApplication = matcher.fillNextTokenBitmask(pointer)
+        if !needsApplication {
+            // xgrammar signals all tokens are allowed — no mask needed
+            return .unconstrained
+        }
+        // Check for all-zeros (no tokens allowed — grammar done)
+        for i in 0..<bitmaskSize {
+            if pointer[i] != 0 { return .constrained }
+        }
+        allTokensBlocked = true
+        return .terminated
+    }
 }
 
 // MARK: - Float16 Masking

swift/Sources/CoreAILanguageModels/GuidedGeneration/XGrammarWrapper.swift

@@ -136,9 +136,29 @@ public final class GrammarMatcher {
         return xgrammar_matcher_is_terminated(handle)
     }
 
+    public var isCompleted: Bool {
+        return xgrammar_matcher_is_completed(handle)
+    }
+
     public func reset() {
         xgrammar_matcher_reset(handle)
     }
+
+    @discardableResult
+    public func rollback(_ numTokens: Int = 1) -> Bool {
+        xgrammar_matcher_rollback(handle, Int32(numTokens))
+    }
+
+    /// Returns the longest deterministic string from the current grammar state,
+    /// or nil if no jump-forward is possible. Does not change matcher state.
+    public func findJumpForwardString() -> String? {
+        guard let cStr = xgrammar_matcher_find_jump_forward_string(handle) else {
+            return nil
+        }
+        let result = String(cString: cStr)
+        free(UnsafeMutablePointer(mutating: cStr))
+        return result.isEmpty ? nil : result
+    }
 }
 
 // MARK: - Errors

swift/Sources/lib/CXGrammar/include/xgrammar_c_bridge.h

@@ -101,6 +101,16 @@ bool xgrammar_matcher_is_terminated(const XGrammarMatcher* matcher);
 // Reset matcher
 void xgrammar_matcher_reset(XGrammarMatcher* matcher);
 
+// Rollback N tokens (restores grammar state to N tokens ago)
+bool xgrammar_matcher_rollback(XGrammarMatcher* matcher, int num_tokens);
+
+// Find the longest deterministic string from the current grammar state.
+// Returns a malloc'd C string (caller must free), or NULL if no jump-forward is possible.
+const char* xgrammar_matcher_find_jump_forward_string(XGrammarMatcher* matcher);
+
+// Check if the grammar's root rule has been fully matched (without requiring stop token)
+bool xgrammar_matcher_is_completed(const XGrammarMatcher* matcher);
+
 // Free grammar matcher
 void xgrammar_matcher_free(XGrammarMatcher* matcher);
 

swift/Sources/lib/CXGrammar/xgrammar_c_bridge.cpp

@@ -222,6 +222,39 @@ void xgrammar_matcher_reset(XGrammarMatcher* matcher) {
     } catch (...) {}
 }
 
+bool xgrammar_matcher_rollback(XGrammarMatcher* matcher, int num_tokens) {
+    try {
+        if (matcher) {
+            matcher->cpp_obj.Rollback(num_tokens);
+            return true;
+        }
+    } catch (...) {}
+    return false;
+}
+
+const char* xgrammar_matcher_find_jump_forward_string(XGrammarMatcher* matcher) {
+    try {
+        if (!matcher) return nullptr;
+        std::string result = matcher->cpp_obj.FindJumpForwardString();
+        if (result.empty()) return nullptr;
+        char* buf = static_cast<char*>(malloc(result.size() + 1));
+        if (!buf) return nullptr;
+        memcpy(buf, result.c_str(), result.size() + 1);
+        return buf;
+    } catch (...) {
+        return nullptr;
+    }
+}
+
+bool xgrammar_matcher_is_completed(const XGrammarMatcher* matcher) {
+    try {
+        if (!matcher) return false;
+        return matcher->cpp_obj.IsCompleted();
+    } catch (...) {
+        return false;
+    }
+}
+
 void xgrammar_matcher_free(XGrammarMatcher* matcher) {
     delete matcher;
 }

swift/Tests/GuidedGenerationTests/ConstrainedGenerationSessionTests.swift

@@ -84,6 +84,43 @@ struct ConstrainedGenerationSessionTests {
         #expect(allowedTokenIDs(session: &session).isEmpty)
     }
 
+    @Test func fillBitmaskIntoPointerMatchesNextTokenBitmask() throws {
+        var session = try createTestSession()
+
+        // Get bitmask via the array-returning method
+        let arrayBitmask = session.nextTokenBitmask()
+        #expect(arrayBitmask != nil)
+
+        // Reset and get bitmask via the pointer method
+        session.reset()
+        let bitmaskSize = (session.vocabularySize + 31) / 32
+        var pointerBitmask = [Int32](repeating: 0, count: bitmaskSize)
+        let result = pointerBitmask.withUnsafeMutableBufferPointer { buf in
+            session.fillBitmask(into: buf.baseAddress!)
+        }
+        #expect(result == .constrained)
+
+        // They should be identical
+        #expect(arrayBitmask!.count == pointerBitmask.count)
+        for i in 0..<arrayBitmask!.count {
+            #expect(
+                arrayBitmask![i] == pointerBitmask[i],
+                "Mismatch at word \(i): array=\(arrayBitmask![i]) pointer=\(pointerBitmask[i])")
+        }
+    }
+
+    @Test func fillBitmaskReturnsTerminatedWhenDone() throws {
+        var session = try createTestSession()
+        driveToCompletion(session: &session)
+
+        let bitmaskSize = (session.vocabularySize + 31) / 32
+        var buffer = [Int32](repeating: 0, count: bitmaskSize)
+        let result = buffer.withUnsafeMutableBufferPointer { buf in
+            session.fillBitmask(into: buf.baseAddress!)
+        }
+        #expect(result == .terminated)
+    }
+
     // MARK: - Token Acceptance Tests
 
     @Test func acceptToken() throws {
@@ -109,6 +146,64 @@ struct ConstrainedGenerationSessionTests {
         )
     }
 
+    // MARK: - Rollback Tests
+
+    @Test func rollbackRestoresPriorState() throws {
+        var session = try createTestSession()
+
+        let initialAllowed = allowedTokenIDs(session: &session)
+
+        // Accept "{" then rollback — should return to initial state
+        _ = session.acceptToken(TestConstants.openBraceToken)
+        let afterAccept = allowedTokenIDs(session: &session)
+        #expect(afterAccept != initialAllowed)
+
+        let success = session.rollback(1)
+        #expect(success, "rollback(1) should succeed")
+
+        let afterRollback = allowedTokenIDs(session: &session)
+        #expect(afterRollback == initialAllowed, "After rollback, allowed tokens should match initial state")
+    }
+
+    @Test func rollbackMultipleTokens() throws {
+        var session = try createTestSession()
+
+        let initialAllowed = allowedTokenIDs(session: &session)
+
+        // Accept "{" then "\"" then rollback both
+        _ = session.acceptToken(TestConstants.openBraceToken)  // "{"
+        _ = session.acceptToken(TestConstants.quoteToken)  // "\""
+
+        let success = session.rollback(2)
+        #expect(success)
+
+        let afterRollback = allowedTokenIDs(session: &session)
+        #expect(afterRollback == initialAllowed)
+    }
+
+    // MARK: - Jump Forward Tests
+
+    @Test func findJumpForwardStringReturnsNilAtChoice() throws {
+        var session = try createTestSession()
+
+        // At the start, multiple tokens are valid ("{") — no deterministic jump
+        let jump = session.findJumpForwardString()
+        // The grammar may or may not have a deterministic prefix at the very start.
+        // After "{", the next required token is "\"" (to start a key), but there could
+        // be whitespace options. Just verify the API returns without crashing.
+        _ = jump  // no assertion on value — grammar-dependent
+    }
+
+    @Test func findJumpForwardStringDoesNotMutateState() throws {
+        var session = try createTestSession()
+
+        let beforeAllowed = allowedTokenIDs(session: &session)
+        _ = session.findJumpForwardString()
+        let afterAllowed = allowedTokenIDs(session: &session)
+
+        #expect(beforeAllowed == afterAllowed, "findJumpForwardString should not change grammar state")
+    }
+
     // MARK: - Reset Tests
 
     @Test func reset() throws {