diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml new file mode 100644 index 000000000..9e9cb5d8c --- /dev/null +++ b/.github/workflows/integration_tests.yml @@ -0,0 +1,72 @@ +name: Integration Tests + +# Heavy integration tests (Hugging Face model downloads, Metal GPU, long-running). +# Kept out of the PR path so they never block merges +on: + workflow_dispatch: + +permissions: + contents: read + +jobs: + integration_tests: + if: github.repository == 'ml-explore/mlx-swift-lm' + runs-on: [self-hosted, macos] + timeout-minutes: 120 + steps: + - uses: actions/checkout@v6 + with: + submodules: recursive + + - name: Select Xcode + shell: bash + run: | + # The FoundationModels integration tests only build against the macOS 27 + # SDK (canImport(FoundationModels, _version: 2)); on the macOS 26 SDK the + # MLXFoundationModels adapter compiles out and those tests are skipped. + # Prefer Xcode 27 when the runner has it so the full suite runs; otherwise + # fall back to the default toolchain and run the SDK-agnostic tests. + dev="" + for app in /Applications/Xcode_27*.app /Applications/Xcode-27*.app /Applications/Xcode.app; do + [ -d "$app" ] || continue + v=$("$app/Contents/Developer/usr/bin/xcodebuild" -version 2>/dev/null | head -1) + case "$v" in "Xcode 27"*) dev="$app/Contents/Developer" ;; esac + [ -n "$dev" ] && break + done + if [ -n "$dev" ]; then + echo "Using Xcode 27 at $dev (full suite, incl. FoundationModels tests)" + echo "DEVELOPER_DIR=$dev" >> "$GITHUB_ENV" + else + echo "Xcode 27 not found; using default $(xcode-select -p)." + echo "FoundationModels tests will be compiled out (macOS 27 SDK required)." + fi + + - name: Verify MetalToolchain installed + shell: bash + run: xcodebuild -showComponent MetalToolchain + + - name: Run IntegrationTesting tests (Xcode, macOS) + shell: sh + run: | + xcodebuild -version + swift --version + # Model-dependent tests must run in a single xctest worker: concurrent + # workers race on the shared on-disk Hugging Face cache. Disable parallel + # testing (see IntegrationTestingTests/.../Support/FMTestHelpers.swift). + xcodebuild test \ + -project IntegrationTesting/IntegrationTesting.xcodeproj \ + -scheme IntegrationTesting \ + -destination 'platform=macOS' \ + -parallel-testing-enabled NO \ + -skipPackagePluginValidation \ + -resultBundlePath IntegrationTesting.xcresult + + - name: Upload test results + if: failure() + uses: actions/upload-artifact@v4 + with: + name: integration-test-results + path: | + IntegrationTesting.xcresult + ~/Library/Logs/DiagnosticReports/* + retention-days: 7 diff --git a/IntegrationTesting/IntegrationTesting.xcodeproj/xcshareddata/xcschemes/IntegrationTesting.xcscheme b/IntegrationTesting/IntegrationTesting.xcodeproj/xcshareddata/xcschemes/IntegrationTesting.xcscheme index a1bea13f4..840b80185 100644 --- a/IntegrationTesting/IntegrationTesting.xcodeproj/xcshareddata/xcschemes/IntegrationTesting.xcscheme +++ b/IntegrationTesting/IntegrationTesting.xcodeproj/xcshareddata/xcschemes/IntegrationTesting.xcscheme @@ -39,6 +39,13 @@ + + + + URL? { // on MLX runtime state and produce non-deterministic numeric divergence in // the Rung 2/3 forward parity assertion. +/// Model ID of the 31B assistant drafter checkpoint shared by the tests below. +private let drafter31BModelId = "mlx-community/gemma-4-31B-it-assistant-bf16" + +/// Pinned checkpoint revision matching the weights that were live when the +/// Rung 4 `drafter_block` fixtures were generated. Kept in sync with +/// `MTPRung4TokenParityTests`. +private let drafter31BRevision = "28e92270316e89288579ec59c17939541d9ca433" + +/// Shared downloader for the drafter checkpoint. Fetches to the local HF +/// cache on first use; subsequent tests and runs reuse the cache. +private let downloader: any Downloader = #hubDownloader() + +private func drafter31BDirectory() async throws -> URL { + try await downloader.download( + id: drafter31BModelId, + revision: drafter31BRevision, + matching: ["*.safetensors", "*.json"], + useLatest: false, + progressHandler: { _ in } + ) +} + @Suite(.serialized) struct Gemma4AssistantIntegrationTests { @Test - func testGemma4AssistantConfigurationDecodesRealCheckpoint() throws { - let drafterDir = hfSnapshotDir(modelId: "mlx-community/gemma-4-31B-it-assistant-bf16") - guard let drafterDir else { - Issue.record("31B drafter checkpoint not present in HF cache; skipping") - return - } + func testGemma4AssistantConfigurationDecodesRealCheckpoint() async throws { + let drafterDir = try await drafter31BDirectory() let configURL = drafterDir.appendingPathComponent("config.json") let data = try Data(contentsOf: configURL) let cfg = try JSONDecoder().decode(Gemma4AssistantConfiguration.self, from: data) @@ -64,13 +84,8 @@ struct Gemma4AssistantIntegrationTests { } @Test - func testRung1WeightsLoadFrom31BCheckpoint() throws { - guard - let drafterDir = hfSnapshotDir(modelId: "mlx-community/gemma-4-31B-it-assistant-bf16") - else { - Issue.record("31B drafter checkpoint not in HF cache; skipping Rung 1") - return - } + func testRung1WeightsLoadFrom31BCheckpoint() async throws { + let drafterDir = try await drafter31BDirectory() let configURL = drafterDir.appendingPathComponent("config.json") let cfg = try JSONDecoder().decode( Gemma4AssistantConfiguration.self, from: Data(contentsOf: configURL)) @@ -86,12 +101,7 @@ struct Gemma4AssistantIntegrationTests { guard let fixturesDir = await drafterForwardFixturesOrSkip(name: "case_01") else { return } - guard - let drafterDir = hfSnapshotDir(modelId: "mlx-community/gemma-4-31B-it-assistant-bf16") - else { - Issue.record("31B drafter checkpoint not in HF cache; skipping Rung 2/3") - return - } + let drafterDir = try await drafter31BDirectory() let configURL = drafterDir.appendingPathComponent("config.json") let cfg = try JSONDecoder().decode( diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenFixtureManifestTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenFixtureManifestTests.swift index 48e244b66..ad6e4ca49 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenFixtureManifestTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenFixtureManifestTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenReplayTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenReplayTests.swift index 533e3115a..bf9f51736 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenReplayTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/GoldenReplayTests.swift @@ -66,7 +66,7 @@ // Gated on both traits because the tokenizer path routes through the // same `loadTestModelContainer` as the bridge tests. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/MalformedSchemaErrorParityTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/MalformedSchemaErrorParityTests.swift index f6997ca34..268e478f9 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/MalformedSchemaErrorParityTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GoldenReplay/MalformedSchemaErrorParityTests.swift @@ -34,7 +34,7 @@ // Gated on both traits because the tokenizer path routes through // `loadTestModelContainer` the same as the other integration tests. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift index bcc57d164..8d6663d93 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -33,10 +33,8 @@ struct GenerableRoundTripTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } return text diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift index f83a574ad..16502f054 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift @@ -21,7 +21,7 @@ // text because the schema -- a single `const` string field -- forces // the entire body as FF after the opening `{`. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift index cae962a94..d2e79aab8 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift @@ -95,7 +95,7 @@ // `loadTestModelContainer`. The GrammarConstraint type lives in the // MLXGuidedGeneration library and is always available alongside the adapter. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/ForkIndependenceTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/ForkIndependenceTests.swift index 38b240021..d71310e28 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/ForkIndependenceTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/ForkIndependenceTests.swift @@ -21,7 +21,7 @@ // Gated on both traits because the tokenizer path routes through // `loadTestModelContainer`. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/LoopInvariantsOnXGrammarTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/LoopInvariantsOnXGrammarTests.swift index 74bbfc770..c7eaaa3f3 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/LoopInvariantsOnXGrammarTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/LoopInvariantsOnXGrammarTests.swift @@ -37,7 +37,7 @@ // `loadTestModelContainer`. `GrammarConstraint` lives in the // MLXGuidedGeneration library and is always available alongside the adapter. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/RollbackDeterminismTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/RollbackDeterminismTests.swift index 6c9aad247..aef057a80 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/RollbackDeterminismTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/RollbackDeterminismTests.swift @@ -19,7 +19,7 @@ // Gated on both traits because the tokenizer path routes through // the same `loadTestModelContainer` as the other tests. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/TokenizerVocabExtractorTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/TokenizerVocabExtractorTests.swift index d421e728d..dfaccd254 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/TokenizerVocabExtractorTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/TokenizerVocabExtractorTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/XGrammarBridgeTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/XGrammarBridgeTests.swift index 5fa4418ee..899648dd0 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/XGrammarBridgeTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/XGrammarBridgeTests.swift @@ -45,7 +45,7 @@ // fixture is not yet available. Llama-3 coverage is pending its // `tokenizer_llama3.json` fixture. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift index 66a1720c3..b930382ab 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationIntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationIntegrationTests.swift index b57ed0c1f..fea0407b7 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationIntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationIntegrationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -53,33 +53,56 @@ struct GuidedGenerationIntegrationTests { let stream = try await executeResponse(executor, request: request, model: model) - var events: [LanguageModelExecutorGenerationChannel.Event] = [] + var events: [MLXLanguageModel.Executor.GenerationEvent] = [] for try await event in stream { events.append(event) } #expect(events.count >= 2, "Should produce metadata and text events") - guard - let firstResponse = events.first - as? LanguageModelExecutorGenerationChannel.Response, - case .updateMetadata = firstResponse.action - else { + guard case .updateMetadata = events.first else { Issue.record("First event should be metadataUpdate") return } let hasText = events.contains { event in - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { - return true - } + if case .appendText(_, _, .response) = event { return true } return false } #expect(hasText, "Should produce text deltas") } + @Test + func schemaGuidedCancellationPropagates() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt( + segments: [.text(Transcript.TextSegment(content: "Return one integer."))], + responseFormat: nil)) + ]) + let request = makeExecutorRequest(transcript: transcript, schema: Int.generationSchema) + let sink = GuidedGenerationDiagnosticSink(cancelAfterEmitCount: 1) + let stream = try await executeResponse( + executor, + request: request, + model: model, + guidedGenerationSink: sink) + + var sawCancellation = false + do { + for try await _ in stream {} + } catch is CancellationError { + sawCancellation = true + } + await stream.cancelAndWait() + + #expect(sink.emitCount >= 1, "schema generation must reach its guided emit") + #expect(sawCancellation, "schema guided cancellation must not become normal EOS") + } + @Test func noSchemaUsesUnconstrainedPath() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } @@ -100,9 +123,7 @@ struct GuidedGenerationIntegrationTests { var hasText = false for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = event { hasText = true break } @@ -139,10 +160,8 @@ struct GuidedGenerationIntegrationTests { let stream1 = try await executeResponse(executor, request: request1, model: model) var text1 = "" for try await event in stream1 { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text1 += delta.content + if case .appendText(let chunk, _, .response) = event { + text1 += chunk } } #expect(!text1.isEmpty, "Turn 1 (unconstrained) should produce text") @@ -159,10 +178,8 @@ struct GuidedGenerationIntegrationTests { let stream2 = try await executeResponse(executor, request: request2, model: model) var text2 = "" for try await event in stream2 { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text2 += delta.content + if case .appendText(let chunk, _, .response) = event { + text2 += chunk } } let trimmed2 = text2.trimmingCharacters(in: .whitespacesAndNewlines) @@ -187,10 +204,8 @@ struct GuidedGenerationIntegrationTests { let stream3 = try await executeResponse(executor, request: request3, model: model) var text3 = "" for try await event in stream3 { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text3 += delta.content + if case .appendText(let chunk, _, .response) = event { + text3 += chunk } } #expect(!text3.isEmpty, "Turn 3 (unconstrained) should produce text") @@ -220,10 +235,8 @@ struct GuidedGenerationIntegrationTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } return text @@ -244,10 +257,8 @@ struct GuidedGenerationIntegrationTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } return text @@ -286,16 +297,14 @@ struct GuidedGenerationIntegrationTests { let stream = try await executeResponse(executor, request: request, model: model) - var events: [LanguageModelExecutorGenerationChannel.Event] = [] + var events: [MLXLanguageModel.Executor.GenerationEvent] = [] for try await event in stream { events.append(event) } let incompleteIdx = events.firstIndex { event in - guard let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .updateMetadata(let metadata) = response.action - else { return false } - return (metadata.values["incompleteOutput"] as? Bool) == true + guard case .updateMetadata(let metadata, _) = event else { return false } + return (metadata["incompleteOutput"] as? Bool) == true } #expect( incompleteIdx != nil, @@ -304,9 +313,7 @@ struct GuidedGenerationIntegrationTests { if let incompleteIdx, let lastTextIdx = events.lastIndex(where: { - if let response = $0 as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = $0 { return true } else { return false diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift index 7ae9b1f18..78af32594 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift @@ -1,6 +1,6 @@ // Copyright (c) 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationUsageTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationUsageTests.swift index 62f2b1915..4fddb5322 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationUsageTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationUsageTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -185,10 +185,8 @@ struct GuidedGenerationUsageTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } return text diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift index 6a1445f48..73e548670 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift @@ -1,6 +1,6 @@ // Copyright (c) 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift index ee414064c..05cd656d8 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -101,10 +101,8 @@ struct MaxTokenTruncationTests { var fullText = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - fullText += delta.content + if case .appendText(let chunk, _, .response) = event { + fullText += chunk } } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift index 2b2f66ffb..d5c17953a 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -402,10 +402,8 @@ struct MultiModelGuidedGenerationTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } return text diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift index 1d6d8e6e9..50226193e 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -45,9 +45,7 @@ struct PrewarmGrammarTests { var hasText = false for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = event { hasText = true break } @@ -78,9 +76,7 @@ struct PrewarmGrammarTests { var hasText = false for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = event { hasText = true break } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/StopTokenRegressionIntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/StopTokenRegressionIntegrationTests.swift index 04fcb2956..d6bc028ed 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/StopTokenRegressionIntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/StopTokenRegressionIntegrationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/ColdHubClientFetchTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/ColdHubClientFetchTests.swift index ecfcd4307..2c88a9577 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/ColdHubClientFetchTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/ColdHubClientFetchTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/HuggingFaceLanguageModelMacroSmoke.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/HuggingFaceLanguageModelMacroSmoke.swift index 27ecef210..08b03f11b 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/HuggingFaceLanguageModelMacroSmoke.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/PlatformAvailability/HuggingFaceLanguageModelMacroSmoke.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ConfigurationResolverRoutingTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ConfigurationResolverRoutingTests.swift index dcdb75e6e..96117757b 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ConfigurationResolverRoutingTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ConfigurationResolverRoutingTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels @@ -47,14 +47,10 @@ struct ConfigurationResolverRoutingTests { var reasoning = "" var response = "" for try await event in stream { - if let r = event as? LanguageModelExecutorGenerationChannel.Reasoning, - case .appendText(let fragment) = r.action - { - reasoning += fragment.content - } else if let r = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let fragment) = r.action - { - response += fragment.content + if case .appendText(let chunk, _, .reasoning) = event { + reasoning += chunk + } else if case .appendText(let chunk, _, .response) = event { + response += chunk } } return (reasoning, response) diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningCapabilityGateTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningCapabilityGateTests.swift index b52cba4c5..e6c811a3c 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningCapabilityGateTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningCapabilityGateTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels @@ -67,14 +67,10 @@ struct ReasoningCapabilityGateTests { var response = "" var reasoning = "" for try await event in stream { - if let r = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let fragment) = r.action - { - response += fragment.content - } else if let r = event as? LanguageModelExecutorGenerationChannel.Reasoning, - case .appendText(let fragment) = r.action - { - reasoning += fragment.content + if case .appendText(let chunk, _, .response) = event { + response += chunk + } else if case .appendText(let chunk, _, .reasoning) = event { + reasoning += chunk } } // No leak. diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningFamilyVerificationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningFamilyVerificationTests.swift index 6af5d8fea..a96725438 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningFamilyVerificationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningFamilyVerificationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningIntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningIntegrationTests.swift index e35d75f7c..7fb9a325c 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningIntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Reasoning/ReasoningIntegrationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels @@ -53,14 +53,10 @@ struct ReasoningIntegrationTests { var reasoning = "" var response = "" for try await event in stream { - if let r = event as? LanguageModelExecutorGenerationChannel.Reasoning, - case .appendText(let fragment) = r.action - { - reasoning += fragment.content - } else if let r = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let fragment) = r.action - { - response += fragment.content + if case .appendText(let chunk, _, .reasoning) = event { + reasoning += chunk + } else if case .appendText(let chunk, _, .response) = event { + response += chunk } } return (reasoning, response) diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/FMTestHelpers.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/FMTestHelpers.swift index d840a51d3..d358fb324 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/FMTestHelpers.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/FMTestHelpers.swift @@ -40,7 +40,7 @@ var fixturesBundle: Bundle { Bundle(for: FixturesBundleToken.self) } // SmallTokenizer — all of which live outside the gate — for tests that // exercise xgrammar / MLXLMCommon directly. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) /// Constructs an `MLXLanguageModel` using the production test downloader / /// tokenizer loader and a `HubCache.default`-backed `weightsLocation:` closure. @@ -183,15 +183,16 @@ func makeExecutorRequest( /// public `finish()`. In production the framework closes the channel after /// respond returns; tests bypass the framework, so iterating the channel /// directly hangs forever. We relay events into an `AsyncThrowingStream` -/// that we own. A producer task runs `respond()`, then cancels a collector -/// task (which relays channel events into our stream). Our stream's +/// that we own. A producer task attaches an observer that yields readable +/// `GenerationEvent`s into our stream, then runs `respond()` and cancels a +/// collector task that drains and discards the channel's opaque events (just +/// enough to keep `respond()`'s sends from stalling). Our stream's /// continuation is finished once both tasks settle, so `for try await` -/// terminates naturally. Early break from iteration cancels both tasks via -/// `deinit`, so tests that stop reading mid-generation don't waste GPU -/// compute on tokens nobody wants. +/// terminates naturally. Tests that must prove cancellation can call +/// `cancelAndWait()`; ordinary early breaks still cancel both tasks via `deinit`. @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) final class TestResponseStream: AsyncSequence, @unchecked Sendable { - typealias Element = LanguageModelExecutorGenerationChannel.Event + typealias Element = MLXLanguageModel.Executor.GenerationEvent typealias AsyncIterator = AsyncThrowingStream.AsyncIterator private let stream: AsyncThrowingStream @@ -201,33 +202,52 @@ final class TestResponseStream: AsyncSequence, @unchecked Sendable { init( executor: MLXLanguageModel.Executor, request: LanguageModelExecutorGenerationRequest, - model: MLXLanguageModel + model: MLXLanguageModel, + cancelProducerWhen: (@Sendable (Element) -> Bool)? = nil, + guidedGenerationSink: GuidedGenerationDiagnosticSink? = nil ) { let channel = LanguageModelExecutorGenerationChannel() let (stream, continuation) = AsyncThrowingStream.makeStream() self.stream = stream - // Collector: relay events from the framework channel into our stream. + // The SDK's channel events are opaque, so we read generated content via + // the executor's observation hook instead. We still must drain the + // channel: respond() sends into it, and without a consumer those sends + // would stall. The drained events are discarded; content comes from the + // observer below. let collector = Task { do { - for try await event in channel { - continuation.yield(event) - } + for try await _ in channel {} } catch { // Including CancellationError; we don't depend on cancellation here. } } self.collectorTask = collector - // Producer: run respond(), then finish our stream so the test's - // iteration terminates. + // Producer: attach the observer (task-local, so it also reaches the + // guided-path forwarder child task), run respond(), then finish our + // stream so the test's iteration terminates. self.producerTask = Task { defer { collector.cancel() } - do { - try await executor.respond(to: request, model: model, streamingInto: channel) - continuation.finish() - } catch { - continuation.finish(throwing: error) + let resolvedGuidedGenerationSink = + guidedGenerationSink ?? GuidedGenerationDiagnosticSink.current + await GuidedGenerationDiagnosticSink.$current.withValue( + resolvedGuidedGenerationSink + ) { + await MLXLanguageModel.Executor.$generationObserver.withValue({ event in + continuation.yield(event) + if cancelProducerWhen?(event) == true { + withUnsafeCurrentTask { $0?.cancel() } + } + }) { + do { + try await executor.respond( + to: request, model: model, streamingInto: channel) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } } } } @@ -240,6 +260,16 @@ final class TestResponseStream: AsyncSequence, @unchecked Sendable { func makeAsyncIterator() -> AsyncIterator { stream.makeAsyncIterator() } + + /// Cancels an in-flight response and waits until both the producer and its + /// channel collector have unwound. The stream continuation is finished by + /// the producer with the cancellation error from `respond()`. + func cancelAndWait() async { + producerTask.cancel() + await producerTask.value + collectorTask.cancel() + await collectorTask.value + } } /// Starts executor.respond(...) on a background task and returns a wrapper that @@ -253,6 +283,38 @@ func executeResponse( TestResponseStream(executor: executor, request: request, model: model) } +/// Test-only variant that binds guided-generation diagnostics to the producer +/// task, including its synchronous guided emit closure. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +func executeResponse( + _ executor: MLXLanguageModel.Executor, + request: LanguageModelExecutorGenerationRequest, + model: MLXLanguageModel, + guidedGenerationSink: GuidedGenerationDiagnosticSink +) async throws -> TestResponseStream { + TestResponseStream( + executor: executor, + request: request, + model: model, + guidedGenerationSink: guidedGenerationSink) +} + +/// Test-only variant that synchronously cancels the producer from its event +/// observer, allowing cancellation tests to avoid racing buffered output. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +func executeResponse( + _ executor: MLXLanguageModel.Executor, + request: LanguageModelExecutorGenerationRequest, + model: MLXLanguageModel, + cancelProducerWhen: @escaping @Sendable (MLXLanguageModel.Executor.GenerationEvent) -> Bool +) async throws -> TestResponseStream { + TestResponseStream( + executor: executor, + request: request, + model: model, + cancelProducerWhen: cancelProducerWhen) +} + // MARK: - GPU Memory Management /// Releases all GPU memory: synchronizes pending GPU work, evicts cached models, @@ -340,6 +402,17 @@ enum TestFixtures { static let gemmaModelID = "mlx-community/gemma-3-270m-it-4bit" + /// Gemma 4 (E2B) instruction-tuned. Its chat template natively renders + /// tool calls and `tool` responses, so it exercises the multi-turn + /// tool-calling replay path (assistant tool-call + tool result). + static let gemma4ModelID = "mlx-community/gemma-4-e2b-it-4bit" + + /// Qwen3 (4B) instruction-tuned. A second model family whose chat template + /// renders tool calls and `tool` responses, used alongside `gemma4ModelID` + /// to verify the multi-turn tool-calling path is model-agnostic rather than + /// tuned to any one template dialect. + static let qwen3ModelID = "mlx-community/Qwen3-4B-4bit" + /// Default model ID for tests that don't care which specific MLX model runs, /// but do need a model known to exercise the full guided-generation and /// tool-calling paths. diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/TestabilityProbe.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/TestabilityProbe.swift index d661673ad..9b7b3b3eb 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/TestabilityProbe.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/Support/TestabilityProbe.swift @@ -1,7 +1,7 @@ // Compile-only probe. Proves at COMPILE time, against the macOS-27 SDK, that: // 1. `@testable import MLXFoundationModels` resolves from this xcodeproj // test target against the local SwiftPM package. -// 2. An `internal` symbol (`MLXLanguageModel.Executor.samplingMode(from:)`) +// 2. An `internal` symbol (`MLXLanguageModel.Executor.samplingConfiguration(from:)`) // is reachable, i.e. the package was built with testability enabled and // the FoundationModelsIntegration trait came in enabled (the module is // not the empty trait-disabled variant). @@ -9,13 +9,13 @@ // If this file COMPILES, the gate is green. It is never executed — the // function is unreferenced and `@available`-gated. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import FoundationModels @testable import MLXFoundationModels @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) func _testabilityProbe() { // Internal static on the bridge-local Executor — only visible via @testable. - _ = MLXLanguageModel.Executor.samplingMode(from: nil) + _ = MLXLanguageModel.Executor.samplingConfiguration(from: nil) } #endif diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/IntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/IntegrationTests.swift index ade6b5725..8359289b4 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/IntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/IntegrationTests.swift @@ -89,31 +89,23 @@ struct IntegrationTests { contextOptions: ContextOptions(), metadata: [:] ) - let channel = LanguageModelExecutorGenerationChannel() - let respondTask = Task { - try await executor.respond(to: request, model: model, streamingInto: channel) - } - - var events: [LanguageModelExecutorGenerationChannel.Event] = [] - for try await event in channel { + let stream = try await executeResponse(executor, request: request, model: model) + var events: [MLXLanguageModel.Executor.GenerationEvent] = [] + for try await event in stream { events.append(event) if events.count >= 3 { // Get a few events break } } - respondTask.cancel() - try? await respondTask.value // First event should be metadata - guard let response = events.first as? LanguageModelExecutorGenerationChannel.Response, - case .updateMetadata(let metadata) = response.action - else { + guard case .updateMetadata(let metadata, _) = events.first else { Issue.record("First event should be metadataUpdate") return } #expect( - metadata.values["modelID"] != nil, + metadata["modelID"] != nil, "Metadata should contain model identifier" ) } @@ -306,16 +298,11 @@ struct IntegrationTests { contextOptions: ContextOptions(), metadata: [:] ) - let channel = LanguageModelExecutorGenerationChannel() - let respondTask = Task { - try await executor.respond(to: request, model: model, streamingInto: channel) - } + let stream = try await executeResponse(executor, request: request, model: model) var tokenCount = 0 - for try await event in channel { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + for try await event in stream { + if case .appendText(_, _, .response) = event { tokenCount += 1 } // Cancel early after a few tokens @@ -323,8 +310,6 @@ struct IntegrationTests { break } } - // Cancel the respond task since we broke out early - respondTask.cancel() #expect(tokenCount >= 5, "Should have received at least 5 tokens before cancellation") } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/PlainChatGenerationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/PlainChatGenerationTests.swift index 8c7ebc09d..e8f4f9405 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/PlainChatGenerationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/PlainChatGenerationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -37,9 +37,7 @@ struct PlainChatGenerationTests { let stream = try await executeResponse(executor, request: request, model: model) var sawTextDelta = false for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = event { sawTextDelta = true } } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/SamplingModeBehaviorTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/SamplingModeBehaviorTests.swift index e20ecefb0..2884c5c46 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/SamplingModeBehaviorTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/SamplingModeBehaviorTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import FoundationModels @@ -43,10 +43,8 @@ struct SamplingModeBehaviorTests { private func responseText(_ stream: TestResponseStream) async throws -> String { var response = "" for try await event in stream { - if let r = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let fragment) = r.action - { - response += fragment.content + if case .appendText(let chunk, _, .response) = event { + response += chunk } } return response diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/StreamingDeltaTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/StreamingDeltaTests.swift index b2a4f0b44..1d2fee866 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/StreamingDeltaTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/StreamingDeltaTests.swift @@ -1,6 +1,6 @@ // Copyright © 2025 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -30,9 +30,7 @@ struct StreamingDeltaTests { let stream = try await executeResponse(executor, request: request, model: model) var deltaCount = 0 for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText = response.action - { + if case .appendText(_, _, .response) = event { deltaCount += 1 } } @@ -56,10 +54,8 @@ struct StreamingDeltaTests { let stream = try await executeResponse(executor, request: request, model: model) var text = "" for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - text += delta.content + if case .appendText(let chunk, _, .response) = event { + text += chunk } } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/UpdateUsageEmissionTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/UpdateUsageEmissionTests.swift index 67e476dbc..a4a02e3ba 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/UpdateUsageEmissionTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/TextGeneration/UpdateUsageEmissionTests.swift @@ -15,7 +15,7 @@ // Suite is `.serialized` and gated on both traits because the schema/ // tool-calling tests load `ModelContainer` and require xgrammar. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -135,10 +135,8 @@ private func collectFinalUsage( ) async throws -> (input: Int, output: Int) { var lastUsage: (input: Int, output: Int)? for try await event in stream { - if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .updateUsage(let usage) = response.action - { - lastUsage = (usage.input.totalTokenCount, usage.output.totalTokenCount) + if case .updateUsage(let input, let output, _) = event { + lastUsage = (input.totalTokenCount, output.totalTokenCount) } } return try #require( diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/FoundationModelsToolCallingTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/FoundationModelsToolCallingTests.swift index a0c6f9971..6c95d2dce 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/FoundationModelsToolCallingTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/FoundationModelsToolCallingTests.swift @@ -1,12 +1,13 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation import MLX import FoundationModels @testable import MLXFoundationModels +@testable import MLXGuidedGeneration @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) @Generable @@ -15,17 +16,57 @@ private struct WeatherArgs { var location: String } -/// End-to-end test for tool calling via guided generation. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct Greeting { + var message: String +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private enum DeveloperMarkerValue: String { + case recorded +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct DeveloperMarkerArguments { + var value: DeveloperMarkerValue +} + +/// End-to-end tests for Foundation Models tool-calling modes. /// /// This suite validates that when a request has `enabledTools`, the /// executor (1) formats tools into the prompt via the tokenizer's native -/// tool-aware chat template, (2) constrains the model's output to a -/// union-of-tools JSON envelope via xgrammar, and (3) parses the result -/// into either a `toolCallDelta` (real tool) or `textDelta` (synthetic -/// final-answer tool). +/// tool-aware chat template, (2) routes native `.allowed` output to either +/// a response or a real tool call, and (3) constrains `.required` output to +/// a developer-tool JSON envelope via xgrammar. @Suite(.serialized, .timeLimit(.minutes(5))) struct FoundationModelsToolCallingTests { + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func developerTool(named name: String) -> Transcript.ToolDefinition { + Transcript.ToolDefinition( + name: name, + description: "Get the current weather in a given location.", + parameters: WeatherArgs.generationSchema) + } + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func weatherTool() -> Transcript.ToolDefinition { + developerTool(named: "get_weather") + } + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func singlePromptTranscript(_ text: String) -> Transcript { + Transcript(entries: [ + .prompt( + Transcript.Prompt( + segments: [.text(Transcript.TextSegment(content: text))], + responseFormat: nil)) + ]) + } + @Test("Setup: release GPU state from prior suites") func clearGPUBeforeToolCalling() async { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } @@ -37,108 +78,88 @@ struct FoundationModelsToolCallingTests { print("[ToolCallingSetup] freed \(freed)MB active, \(cache)MB cache") } - @Test - func toolsEnabledEmitsToolCallOrText() async throws { + @Test func allowedCanAnswerWithoutCallingATool() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } let model = makeTestModel(TestFixtures.defaultModelID) let executor = try makeMLXExecutor(for: model) - - let weatherTool = Transcript.ToolDefinition( - name: "get_weather", - description: "Get the current weather in a given location.", - parameters: WeatherArgs.generationSchema - ) - - let transcript = Transcript(entries: [ - .prompt( - Transcript.Prompt( - segments: [ - .text(Transcript.TextSegment(content: "What's the weather in Tokyo?")) - ], responseFormat: nil)) - ]) - let request = makeExecutorRequest( - transcript: transcript, - enabledTools: [weatherTool] - ) + transcript: singlePromptTranscript("Say hello in one short sentence."), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 64, + toolCallingMode: .allowed)) let stream = try await executeResponse(executor, request: request, model: model) + var calls: [String] = [] + var response = "" + for try await event in stream { + if case .toolCall(_, let name, _) = event { calls.append(name) } + if case .appendText(let text, _, .response) = event { response += text } + } + #expect(calls.isEmpty) + #expect(!response.isEmpty) + } - var sawWeatherToolCall = false - var sawText = false - var textContent = "" + @Test func allowedCanChooseARealTool() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("What is the current weather in Tokyo?"), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .allowed)) + let stream = try await executeResponse(executor, request: request, model: model) + var calls: [String] = [] + var response = "" for try await event in stream { - if let toolCalls = event as? LanguageModelExecutorGenerationChannel.ToolCalls, - case .toolCall(let toolCall) = toolCalls.action, - case .appendArguments(let argsDelta) = toolCall.action - { - if toolCall.name == "get_weather" { - sawWeatherToolCall = true - let data = Data(argsDelta.content.utf8) - let parsed = try? JSONSerialization.jsonObject(with: data) - #expect( - parsed != nil, - "Tool call arguments should be valid JSON: \(argsDelta.content)") - } - } else if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - sawText = true - textContent += delta.content - } + if case .toolCall(_, let name, _) = event { calls.append(name) } + if case .appendText(let text, _, .response) = event { response += text } } + #expect(calls == ["get_weather"]) + #expect(response.isEmpty) + } - // Exactly one of the two paths should have produced output. - #expect( - sawWeatherToolCall || sawText, - "Executor with enabled tools must emit either a toolCallDelta or a textDelta" - ) + @Test func allowedResponseStillHonorsResponseSchema() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("Say hello in one short sentence."), + enabledTools: [weatherTool()], + schema: Greeting.generationSchema, + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .allowed)) - if sawWeatherToolCall { - #expect( - textContent.isEmpty, - "When a real tool call fires, no text deltas should be emitted" - ) - } else { - #expect( - !textContent.isEmpty, - "When the synthetic final-answer tool fires, text should be non-empty" - ) + let stream = try await executeResponse(executor, request: request, model: model) + var calls: [String] = [] + var responseJSON = "" + for try await event in stream { + if case .toolCall(_, let name, _) = event { calls.append(name) } + if case .appendText(let text, _, .response) = event { responseJSON += text } } + #expect(calls.isEmpty) + let content = try GeneratedContent(json: responseJSON) + _ = try Greeting(content) } - /// With tool-aware prompt formatting plus the tool-call grammar - /// that allows ``-wrapped output, the model can both *see* the - /// available tools in the prompt and emit them in its trained format. + /// With tool-aware prompt formatting plus native allowed generation, the + /// model can both *see* the available tools in the prompt and emit them in + /// its trained format. /// For a weather query, Qwen should pick `get_weather` rather than - /// hallucinating via the synthetic final-answer path. + /// answering without the requested live data. @Test func toolAwarePromptRoutesWeatherQueryToGetWeather() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } let model = makeTestModel(TestFixtures.defaultModelID) let executor = try makeMLXExecutor(for: model) - let weatherTool = Transcript.ToolDefinition( - name: "get_weather", - description: - "Get the current weather in a given location. Use this whenever the user asks about weather, temperature, or conditions anywhere.", - parameters: WeatherArgs.generationSchema - ) - - let transcript = Transcript(entries: [ - .prompt( - Transcript.Prompt( - segments: [ - .text( - Transcript.TextSegment( - content: "What's the current weather in Tokyo, Japan?")) - ], responseFormat: nil)) - ]) - let request = makeExecutorRequest( - transcript: transcript, - enabledTools: [weatherTool] + transcript: singlePromptTranscript("What's the current weather in Tokyo, Japan?"), + enabledTools: [weatherTool()] ) let stream = try await executeResponse(executor, request: request, model: model) @@ -148,16 +169,11 @@ struct FoundationModelsToolCallingTests { var textContent = "" for try await event in stream { - if let toolCalls = event as? LanguageModelExecutorGenerationChannel.ToolCalls, - case .toolCall(let toolCall) = toolCalls.action, - case .appendArguments(let argsDelta) = toolCall.action - { - toolCallName = toolCall.name - toolCallArguments = argsDelta.content - } else if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - textContent += delta.content + if case .toolCall(_, let name, let arguments) = event { + toolCallName = name + toolCallArguments = arguments + } else if case .appendText(let chunk, _, .response) = event { + textContent += chunk } } @@ -175,6 +191,133 @@ struct FoundationModelsToolCallingTests { ) } } + + @Test func requiredCannotAnswerWithoutCallingARealTool() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("Say hello without using a tool."), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .required)) + + let stream = try await executeResponse(executor, request: request, model: model) + var names: [String] = [] + var response = "" + for try await event in stream { + if case .toolCall(_, let name, _) = event { names.append(name) } + if case .appendText(let text, _, .response) = event { response += text } + } + + #expect(names == ["get_weather"]) + #expect(response.isEmpty) + } + + @Test func requiredGuidedCancellationPropagates() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("What is the weather in Tokyo?"), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .required)) + let sink = GuidedGenerationDiagnosticSink(cancelAfterEmitCount: 1) + let stream = try await executeResponse( + executor, + request: request, + model: model, + guidedGenerationSink: sink) + + var sawCancellation = false + do { + for try await _ in stream {} + } catch is CancellationError { + sawCancellation = true + } + await stream.cancelAndWait() + + #expect(sink.emitCount >= 1, "required generation must reach its guided emit") + #expect(sawCancellation, "required guided cancellation must not become normal EOS") + } + + @Test func disallowedIgnoresManuallyEnabledTools() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("Say hello in one sentence."), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 64, + toolCallingMode: .disallowed)) + + let stream = try await executeResponse(executor, request: request, model: model) + var sawToolCall = false + var response = "" + for try await event in stream { + if case .toolCall = event { sawToolCall = true } + if case .appendText(let text, _, .response) = event { response += text } + } + + #expect(!sawToolCall) + #expect(!response.isEmpty) + } + + @Test func formerSyntheticNameIsADeveloperTool() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let developerTool = Transcript.ToolDefinition( + name: "mlx_final_answer", + description: "Record a developer-owned marker.", + parameters: DeveloperMarkerArguments.generationSchema) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("Record the marker."), + enabledTools: [developerTool], + generationOptions: GenerationOptions( + maximumResponseTokens: 64, + toolCallingMode: .required)) + + let stream = try await executeResponse(executor, request: request, model: model) + var names: [String] = [] + var response = "" + for try await event in stream { + if case .toolCall(_, let name, _) = event { names.append(name) } + if case .appendText(let text, _, .response) = event { response += text } + } + + #expect(names == ["mlx_final_answer"]) + #expect(response.isEmpty) + } + + @Test func requiredTruncationNeverEmitsResponseText() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeTestModel(TestFixtures.defaultModelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: singlePromptTranscript("What's the weather in Tokyo?"), + enabledTools: [weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 1, + toolCallingMode: .required)) + + let stream = try await executeResponse(executor, request: request, model: model) + var response = "" + var sawIncompleteOutput = false + for try await event in stream { + if case .appendText(let text, _, .response) = event { response += text } + if case .updateMetadata(let metadata, _) = event { + sawIncompleteOutput = (metadata["incompleteOutput"] as? Bool) == true + } + } + + #expect(response.isEmpty) + #expect(sawIncompleteOutput) + } } #endif diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/MultiTurnToolCallingTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/MultiTurnToolCallingTests.swift new file mode 100644 index 000000000..cbc4111a4 --- /dev/null +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/MultiTurnToolCallingTests.swift @@ -0,0 +1,290 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration + +import Foundation +import FoundationModels +import MLX +import MLXLLM +import MLXLMCommon +import Testing + +@testable import MLXFoundationModels + +/// The on/off choice the model fills in for the flashlight tool, mirroring the +/// `FlashlightTool` in the FMFeatures sample. A `@Generable` enum becomes a +/// fixed set of choices in the argument schema. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private enum FlashlightState: String { + case on, off +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct FlashlightArguments { + @Guide(description: "Whether to turn the flashlight on or off.") + var state: FlashlightState +} + +/// Validates multi-turn tool calling with the flashlight tool from the +/// FMFeatures sample: a device action the user drives across turns ("turn it +/// on", then "turn it off"). Two things must hold once a tool has run: +/// +/// - Continuation round: after the session executes a tool and re-invokes us +/// with the call + output appended, we replay both into the model's context +/// so its answer reflects the result rather than re-issuing the call. +/// - Multi-round: with tools still enabled on a continuation round, native +/// `.allowed` routing lets the model chain another tool or answer from the +/// replayed result. Whether it chains is model-dependent, so it is not +/// asserted; the deterministic guarantee that each round's context is +/// assembled correctly lives in the converter test +/// ``testTwoToolRoundsProduceCorrelatedMessages``. +/// +/// The tool loop is faked here: instead of letting `LanguageModelSession` run +/// the flashlight tool, the transcript is seeded with the call and output the +/// session would have appended. This is exactly the continuation-round request +/// the SDK builds (`enabledTools` are passed every round). +/// +/// The behavioral tests run against every model in ``models`` so the replay +/// path is exercised across template dialects, not tuned to one model family. +@Suite(.serialized, .timeLimit(.minutes(5))) +struct MultiTurnToolCallingTests { + + /// Models whose chat templates render tool calls and `tool` responses. + /// Both families are run so the multi-turn path is verified as + /// model-agnostic rather than fitted to a single template dialect. + static let models = [TestFixtures.gemma4ModelID, TestFixtures.qwen3ModelID] + + @Test("Setup: release GPU state from prior suites") + func clearGPUBeforeMultiTurn() async { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + await releaseAllGPUMemory() + } + + private static let toolName = "set_flashlight" + private static let flashlightOnResult = "Flashlight turned on." + private static let structuredConfirmationCode = "fm-structured-4821" + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func flashlightTool() -> Transcript.ToolDefinition { + Transcript.ToolDefinition( + name: Self.toolName, + description: "Turn the device flashlight (torch) on or off.", + parameters: FlashlightArguments.generationSchema) + } + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func instructions() -> Transcript.Instructions { + Transcript.Instructions( + segments: [ + .text( + Transcript.TextSegment( + content: + "You control the device flashlight. When the user asks, call the flashlight tool to turn the light on or off. After the tool returns a result, tell the user what happened." + )) + ], + toolDefinitions: [flashlightTool()]) + } + + /// A completed "turn it on" round: user request + the flashlight call and + /// its output. `ToolOutput.id` matches the call id, as the SDK's + /// ToolCallCoordinator sets it, so the template correlates them. + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func turnOnRound() throws -> [Transcript.Entry] { + let callID = "call_flashlight_on" + let toolCall = Transcript.ToolCall( + id: callID, + toolName: Self.toolName, + arguments: try GeneratedContent(json: #"{"state":"on"}"#)) + let toolOutput = Transcript.ToolOutput( + id: callID, + toolName: Self.toolName, + segments: [.text(Transcript.TextSegment(content: Self.flashlightOnResult))]) + return [ + .prompt( + Transcript.Prompt( + segments: [.text(Transcript.TextSegment(content: "Turn on the flashlight."))], + responseFormat: nil)), + .toolCalls(Transcript.ToolCalls(id: "toolcalls_on", [toolCall])), + .toolOutput(toolOutput), + ] + } + + /// The continuation-round transcript: instructions + the "turn it on" round. + /// The newest entry is a tool output, so this is the round the session + /// re-invokes us with immediately after the tool ran. + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func continuationTranscript() throws -> Transcript { + Transcript(entries: [.instructions(instructions())] + (try turnOnRound())) + } + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func structuredContinuationTranscript() throws -> Transcript { + let callID = "call_flashlight_structured" + let result = try GeneratedContent( + json: #"{"state":"on","confirmationCode":"fm-structured-4821"}"#) + let toolCall = Transcript.ToolCall( + id: callID, + toolName: Self.toolName, + arguments: try GeneratedContent(json: #"{"state":"on"}"#)) + let toolOutput = Transcript.ToolOutput( + id: callID, + toolName: Self.toolName, + segments: [ + .structure( + Transcript.StructuredSegment( + schemaName: "FlashlightResult", + content: result)) + ]) + + return Transcript(entries: [ + .instructions(instructions()), + .prompt( + Transcript.Prompt( + segments: [ + .text(Transcript.TextSegment(content: "Turn on the flashlight.")) + ], + responseFormat: nil)), + .toolCalls(Transcript.ToolCalls(id: "toolcalls_structured", [toolCall])), + .toolOutput(toolOutput), + ]) + } + + /// Constructs each model the way a real app would. Reasoning-first models + /// (Qwen3) must declare `.reasoning` so the tool path runs the + /// think-then-call phase (letting the model reason before the grammar + /// constrains it), and are built from their `LLMRegistry` configuration so + /// they carry the right `extraEOSTokens`. Other families keep the default + /// tool-calling capabilities. Returns the reasoning level to thread into the + /// request (nil for non-reasoning models). + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func makeMultiTurnModel( + _ modelID: String + ) -> (model: MLXLanguageModel, level: ContextOptions.ReasoningLevel?) { + switch modelID { + case TestFixtures.qwen3ModelID: + let model = MLXLanguageModel( + configuration: LLMRegistry.qwen3_4b_4bit, + capabilities: [.reasoning, .guidedGeneration, .toolCalling], + weightsLocation: testWeightsLocation(modelID:), + load: testLoad()) + return (model, .moderate) + default: + return (makeTestModel(modelID), nil) + } + } + + @Test(arguments: MultiTurnToolCallingTests.models) + func continuationRoundIncorporatesToolOutput(modelID: String) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let (model, reasoningLevel) = makeMultiTurnModel(modelID) + let executor = try makeMLXExecutor(for: model) + + let request = makeExecutorRequest( + transcript: try continuationTranscript(), + enabledTools: [flashlightTool()], + generationOptions: GenerationOptions(maximumResponseTokens: 128), + contextOptions: ContextOptions(reasoningLevel: reasoningLevel)) + + let stream = try await executeResponse(executor, request: request, model: model) + + var textContent = "" + var madeToolCall = false + + for try await event in stream { + if case .toolCall = event { + madeToolCall = true + } else if case .appendText(let chunk, _, .response) = event { + textContent += chunk + } + } + + let lowered = textContent.lowercased() + let reflectsOutput = lowered.contains("light") || lowered.contains("on") + + // Multi-round: `.allowed` mode controls termination on a continuation + // round, so the model may answer from the replayed result or call a tool + // again. The regression this guards is the model ignoring the tool + // output: a plain answer that neither reflects the result nor issues a + // tool call. + #expect( + reflectsOutput || madeToolCall, + "Continuation must reflect the tool output or issue a tool call, not ignore the result. Got: \"\(textContent)\" (\(modelID))" + ) + } + + @Test(arguments: MultiTurnToolCallingTests.models) + func requiredContinuationCallsAnotherRealTool(modelID: String) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let (model, reasoningLevel) = makeMultiTurnModel(modelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: try continuationTranscript(), + enabledTools: [flashlightTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .required), + contextOptions: ContextOptions(reasoningLevel: reasoningLevel)) + + let stream = try await executeResponse(executor, request: request, model: model) + var names: [String] = [] + var response = "" + for try await event in stream { + if case .toolCall(_, let name, _) = event { names.append(name) } + if case .appendText(let text, _, .response) = event { response += text } + } + + #expect(names == [Self.toolName]) + #expect(response.isEmpty) + } + + @Test(arguments: MultiTurnToolCallingTests.models) + func continuationPromptContainsToolOutput(modelID: String) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let model = makeTestModel(modelID) + let container = try await model.loadContainer() + let transcript = try continuationTranscript() + + let decoded = try await container.perform { context in + let messages = TranscriptConverter.mlxMessages(for: transcript) + let raw = DefaultMessageGenerator().generate(messages: messages) + let tokens = try context.tokenizer.applyChatTemplate(messages: raw) + return context.tokenizer.decode(tokenIds: tokens) + } + + print("[MultiTurn][RenderedPrompt][\(modelID)]\n\(decoded)\n[/RenderedPrompt]") + + #expect( + decoded.lowercased().contains("flashlight turned on"), + "Rendered continuation prompt must contain the tool output text (\(modelID))" + ) + } + + @Test(arguments: MultiTurnToolCallingTests.models) + func continuationPromptContainsStructuredToolOutput(modelID: String) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let model = makeTestModel(modelID) + let container = try await model.loadContainer() + let transcript = try structuredContinuationTranscript() + + let decoded = try await container.perform { context in + let messages = TranscriptConverter.mlxMessages(for: transcript) + let raw = DefaultMessageGenerator().generate(messages: messages) + let tokens = try context.tokenizer.applyChatTemplate(messages: raw) + return context.tokenizer.decode(tokenIds: tokens) + } + + #expect( + decoded.contains(Self.structuredConfirmationCode), + "Rendered continuation prompt must contain the structured tool result (\(modelID))") + #expect( + decoded.contains("confirmationCode"), + "Rendered continuation prompt must preserve the structured result field (\(modelID))") + } +} + +#endif // FoundationModelsIntegration diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/StructuredToolOutputSessionTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/StructuredToolOutputSessionTests.swift new file mode 100644 index 000000000..0bc3a0a5e --- /dev/null +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/StructuredToolOutputSessionTests.swift @@ -0,0 +1,143 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration + +import Foundation +import FoundationModels +import MLX +import Testing + +@testable import MLXFoundationModels + +private let structuredToolOutputSentinel = "fm-structured-output-9D7E-4821" + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +extension SessionPropertyValues { + @SessionPropertyEntry + var structuredToolOutputCallCount: Int = 0 +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct StructuredLookupArguments { + @Guide(description: "The record identifier to retrieve.") + var recordID: String +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct StructuredLookupResult { + var requiredToken: String + var summary: String +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +private struct StructuredLookupTool: Tool { + let name = "lookup_structured_record" + let description = + "Retrieves a record whose structured result contains the exact token for the response." + + func call(arguments: StructuredLookupArguments) async throws -> StructuredLookupResult { + StructuredLookupResult( + requiredToken: structuredToolOutputSentinel, + summary: "Retrieved record \(arguments.recordID).") + } +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +private struct StructuredToolOutputProfile: LanguageModelSession.DynamicProfile { + let model: MLXLanguageModel + + @SessionProperty(\.structuredToolOutputCallCount) + var toolCallCount + + var body: some LanguageModelSession.DynamicProfile { + if toolCallCount == 0 { + Profile { + Instructions { + "Call the lookup tool once. After it returns, answer with the value of its requiredToken field exactly." + } + StructuredLookupTool() + } + .model(model) + .toolCallingMode(.required) + .onToolCall { + toolCallCount += 1 + } + } else { + Profile { + Instructions { + "Use the latest tool output. Return its requiredToken field exactly and no other text." + } + } + .model(model) + .toolCallingMode(.disallowed) + } + } +} + +/// Opt-in live Foundation Models tool round trip. The two model rows load +/// multi-GB checkpoints, so default sweeps skip this suite. When launching +/// through `xcodebuild`, set `TEST_RUNNER_MLX_RUN_FM_TOOL_INTEGRATION=1` on +/// macOS or iOS; Xcode strips the `TEST_RUNNER_` prefix before launching the +/// test process. +@Suite( + .serialized, + .timeLimit(.minutes(10)), + .enabled( + if: ProcessInfo.processInfo.environment[ + "MLX_RUN_FM_TOOL_INTEGRATION" + ] == "1") +) +struct StructuredToolOutputSessionTests { + static let models = [TestFixtures.gemma4ModelID, TestFixtures.qwen3ModelID] + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func makeSessionModel(_ modelID: String) -> MLXLanguageModel { + if modelID == TestFixtures.qwen3ModelID { + return makeReasoningTestModel(modelID) + } + return makeTestModel(modelID) + } + + @Test("Setup: release GPU state before live structured-output sessions") + func clearGPUBeforeStructuredSessions() async { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + await releaseAllGPUMemory() + } + + @Test(arguments: StructuredToolOutputSessionTests.models) + func languageModelSessionUsesGenerableToolOutput(modelID: String) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let model = makeSessionModel(modelID) + let session = LanguageModelSession( + profile: StructuredToolOutputProfile(model: model)) + + let response = try await session.respond( + to: "Retrieve record alpha and follow the profile instructions.") + + #expect(session.properties.structuredToolOutputCallCount == 1) + + let outputs = session.transcript.compactMap { entry -> Transcript.ToolOutput? in + guard case .toolOutput(let output) = entry else { return nil } + return output + } + #expect(outputs.count == 1) + + let structuredSegments = outputs.flatMap { $0.segments }.compactMap { + segment -> Transcript.StructuredSegment? in + guard case .structure(let structured) = segment else { return nil } + return structured + } + let structured = try #require(structuredSegments.first) + #expect(structured.content.jsonString.contains(structuredToolOutputSentinel)) + + #expect( + response.content.contains(structuredToolOutputSentinel), + "The live \(modelID) response must use the token available only in the structured tool output. Got: \(response.content)" + ) + } +} + +#endif // FoundationModelsIntegration diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallGrammarCharacterizationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallGrammarCharacterizationTests.swift new file mode 100644 index 000000000..cee22f73d --- /dev/null +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallGrammarCharacterizationTests.swift @@ -0,0 +1,161 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration + +import Foundation +import FoundationModels +import MLX +import MLXLMCommon +import Testing + +@testable import MLXFoundationModels +@testable import MLXGuidedGeneration + +/// The fixed `on` state the model fills in for this characterization tool. +/// A one-case `@Generable` enum bounds the generated argument. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private enum CharacterizationFlashlightState: String { + case on +} + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private struct CharacterizationFlashlightArguments { + @Guide(description: "The fixed on state for this characterization.") + var state: CharacterizationFlashlightState +} + +/// Characterizes the explicit `.required` guided tool-call grammar/tokenizer +/// boundary on real models. +/// +/// This drives the real executor's required guided path on each target model on +/// a fresh turn and captures the grammar/tokenizer boundary through +/// `GuidedGenerationDiagnosticSink` (token IDs, whether the grammar terminated, +/// the exact buffer handed to the parser, and whether that buffer parses), +/// recording a token-level trace while pinning valid parsing as the stable +/// invariant. +/// +/// Both Gemma-4 E2B and Qwen-3 are expected to terminate the guided grammar and +/// hand the parser a valid developer tool call. Add rows when another model +/// needs this same required-mode boundary coverage. +@Suite(.serialized, .timeLimit(.minutes(5))) +struct ToolCallGrammarCharacterizationTests { + + struct Case: CustomStringConvertible, Sendable { + let modelID: String + /// The pinned invariant: whether the buffer handed to the parser is + /// expected to parse as a valid tool call today. + let expectsValidToolCall: Bool + var description: String { modelID } + } + + static let cases: [Case] = [ + // Required guided grammar must terminate and parse a developer tool for + // each supported family. + Case(modelID: TestFixtures.gemma4ModelID, expectsValidToolCall: true), + Case(modelID: TestFixtures.qwen3ModelID, expectsValidToolCall: true), + ] + + private static let toolName = "set_flashlight" + + @Test("Setup: release GPU state from prior suites") + func clearGPUBeforeCharacterization() async { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + await releaseAllGPUMemory() + } + + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func flashlightTool() -> Transcript.ToolDefinition { + Transcript.ToolDefinition( + name: Self.toolName, + description: "Turn the device flashlight (torch) on.", + parameters: CharacterizationFlashlightArguments.generationSchema) + } + + /// A fresh turn: instructions plus a single user request that warrants a + /// tool call. No prior tool call/output, so the model's first action this + /// turn is expected to be a same-turn tool call, which is the path we trace. + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) + private func freshTurnTranscript() -> Transcript { + let instructions = Transcript.Instructions( + segments: [ + .text( + Transcript.TextSegment( + content: + "You control the device flashlight. Call the flashlight tool to turn the light on." + )) + ], + toolDefinitions: [flashlightTool()]) + let prompt = Transcript.Prompt( + segments: [.text(Transcript.TextSegment(content: "Turn on the flashlight."))], + responseFormat: nil) + return Transcript(entries: [.instructions(instructions), .prompt(prompt)]) + } + + @Test(arguments: ToolCallGrammarCharacterizationTests.cases) + func characterizeToolCallBoundary(_ testCase: Case) async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let model = makeTestModel(testCase.modelID) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: freshTurnTranscript(), + enabledTools: [flashlightTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .required)) + + let sink = GuidedGenerationDiagnosticSink() + + // Bind the sink so the real executor's generation loop and tool path + // record into it. The unstructured producer Task created inside + // executeResponse inherits this task-local binding. + try await GuidedGenerationDiagnosticSink.$current.withValue(sink) { + let stream = try await executeResponse(executor, request: request, model: model) + for try await _ in stream {} // drain to completion + } + + // Decode the sampled token IDs back to text for the trace. + let sampledText = try await model.loadContainer().perform { context in + context.tokenizer.decode(tokenIds: sink.sampledTokenIDs) + } + + // Record the token-level trace without affecting the assertions below. + print( + """ + [ToolCallCharacterization][\(testCase.modelID)] + sampledTokenIDs (\(sink.sampledTokenIDs.count)): \(sink.sampledTokenIDs) + sampledDecoded: \(sampledText) + fastForwardTokenIDs (\(sink.fastForwardTokenIDs.count)): \(sink.fastForwardTokenIDs) + grammarTerminated: \(sink.grammarTerminated) + generatedTokenCount: \(sink.generatedTokenCount) + incompleteOutput: \(sink.incompleteOutput) + parsedAsToolCall: \(String(describing: sink.parsedAsToolCall)) + parsedName: \(String(describing: sink.parsedName)) + finalBuffer: \(String(describing: sink.finalBuffer)) + [/ToolCallCharacterization] + """) + + // Sanity: generation actually ran and the tool path handed off a buffer. + #expect( + sink.generatedTokenCount >= 1, + "No tokens were generated (\(testCase.modelID))") + #expect( + sink.finalBuffer != nil, + "Tool path never handed a buffer to the parser (\(testCase.modelID))") + + // Pinned invariant: required guided output must parse as a tool call. + #expect( + sink.parsedAsToolCall == testCase.expectsValidToolCall, + """ + Parse verdict for \(testCase.modelID) was \ + \(String(describing: sink.parsedAsToolCall)), expected \ + \(testCase.expectsValidToolCall). Buffer: \ + \(String(describing: sink.finalBuffer)). A false verdict is a \ + grammar-path regression. + """) + } +} + +#endif // FoundationModelsIntegration diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallRoundTripTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallRoundTripTests.swift index 96405c7fb..c323d22ee 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallRoundTripTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallRoundTripTests.swift @@ -31,7 +31,7 @@ // `loadTestModelContainer` and the schema path requires `@Generable`, // which is behind `FoundationModelsIntegration`. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningCharacterizationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningCharacterizationTests.swift index b4399e0f6..9864744bd 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningCharacterizationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningCharacterizationTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Testing import Foundation @@ -9,23 +9,29 @@ import FoundationModels @testable import MLXFoundationModels import MLXLMCommon +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +@Generable +private enum WeatherLocation: String { + case tokyo = "Tokyo" +} + @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) @Generable private struct WeatherArgs { - @Guide(description: "City and state, e.g. 'San Francisco, CA'.") - var location: String + @Guide(description: "Location to query.") + var location: WeatherLocation } /// Characterizes two empirical facts about today's tool-calling path /// (device/manual-gated, requires a device running iOS 27.0+). Touches no production code. /// /// What it answers: -/// 1. SEPARABILITY (`qwen3WithToolsSuppressesThink`): does the tool-calling -/// grammar suppress `` today? The structural tag is compiled in the -/// *non-triggered* form (`xg_compile_structural_tag` with `nullopt`, no -/// `token_triggered_tags`; see `XGrammarBridge.swift:409`), so the model is -/// masked to ``/`{` from generated-token zero and cannot emit -/// ``. If a marker leaks, the grammar does not in fact suppress it. +/// 1. REQUIRED ENVELOPE (`qwen3RequiredModeEmitsOnlyDeveloperToolCall`): does +/// `.required` mode emit only its enabled developer tool? The structural tag +/// is compiled in the *non-triggered* form (`xg_compile_structural_tag` with +/// `nullopt`, no `token_triggered_tags`; see `XGrammarBridge.swift:409`), so +/// guided generation starts at generated-token zero and must produce a tool +/// envelope rather than response text or reasoning markers. /// 2. TOOL-AWARE THINKING SEED (`toolAwareTemplateHonorsEnableThinking`): does the /// 3-arg `applyChatTemplate(messages:tools:additionalContext:)` produce a /// *distinct* thinking-primed prompt on the tool path, and what `primedInside` @@ -51,14 +57,14 @@ struct ToolCallingReasoningCharacterizationTests { await releaseAllGPUMemory() } - // MARK: - 1. Separability: does the tool grammar suppress `` today? + // MARK: - 1. Required mode: does guidance emit only the developer tool? - /// Drives Qwen3 (thinking-on by template default) + a weather tool through the - /// CURRENT single-phase tool path. Qwen3 wants to emit `` on the - /// unconstrained path, but the token-zero structural-tag grammar should mask it - /// out. The falsifiable assertion: no response/tool-call delta contains `` - /// or ``. - @Test func qwen3WithToolsSuppressesThink() async throws { + /// Drives Qwen3 plus a weather tool through explicit `.required` mode. This + /// fixture does not declare reasoning, so production disables thinking in the + /// tool-aware template before token-zero structural guidance begins. The + /// falsifiable assertion is that guidance emits the enabled developer tool, + /// its arguments contain no reasoning markers, and no response text appears. + @Test func qwen3RequiredModeEmitsOnlyDeveloperToolCall() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } let model = makeTestModel(Self.qwen3) let executor = try makeMLXExecutor(for: model) @@ -75,23 +81,23 @@ struct ToolCallingReasoningCharacterizationTests { .text(Transcript.TextSegment(content: "What's the weather in Tokyo?")) ], responseFormat: nil)) ]) - let request = makeExecutorRequest(transcript: transcript, enabledTools: [weatherTool]) + let request = makeExecutorRequest( + transcript: transcript, + enabledTools: [weatherTool], + generationOptions: GenerationOptions( + maximumResponseTokens: 128, + toolCallingMode: .required)) let stream = try await executeResponse(executor, request: request, model: model) var responseText = "" var toolCallName: String? = nil var toolArgs = "" for try await event in stream { - if let toolCalls = event as? LanguageModelExecutorGenerationChannel.ToolCalls, - case .toolCall(let toolCall) = toolCalls.action, - case .appendArguments(let argsDelta) = toolCall.action - { - toolCallName = toolCall.name - toolArgs += argsDelta.content - } else if let response = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let delta) = response.action - { - responseText += delta.content + if case .toolCall(_, let name, let arguments) = event { + toolCallName = name + toolArgs += arguments + } else if case .appendText(let chunk, _, .response) = event { + responseText += chunk } } @@ -100,21 +106,14 @@ struct ToolCallingReasoningCharacterizationTests { + "responseText=<<<\(responseText.prefix(200))>>> args=<<<\(toolArgs.prefix(200))>>>" ) - // THE confirmation: the grammar must have suppressed the think markers. - let leakedInResponse = - responseText.contains("") || responseText.contains("") + // Required guidance must emit the enabled developer tool rather than + // reasoning markers or response text. let leakedInArgs = toolArgs.contains("") || toolArgs.contains("") - #expect( - !leakedInResponse, - "Grammar-suppression hypothesis failed: / reached .response on the tool path." - ) #expect( !leakedInArgs, "Tool-call arguments must not contain reasoning markers.") - // Sanity: something was produced (tool call or synthetic-final-answer text). - #expect( - toolCallName != nil || !responseText.isEmpty, - "The tool path must emit a tool call or text.") + #expect(toolCallName == "get_weather") + #expect(responseText.isEmpty) } // MARK: - 2. Tool-aware thinking seed: does the 3-arg template honor enable_thinking? diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningTests.swift index 0a395541e..0629cc91b 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/ToolCalling/ToolCallingReasoningTests.swift @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) import Foundation import MLX @@ -8,6 +8,7 @@ import FoundationModels import Testing @testable import MLXFoundationModels +@testable import MLXGuidedGeneration import MLXLMCommon @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) @@ -17,13 +18,12 @@ private struct WeatherArgs { var location: String } -/// Think-then-call: a reasoning model given tools reasons unconstrained -/// first, then emits a grammar-constrained tool call. +/// Native allowed tool use: a reasoning model given tools can reason first, +/// then either answer or emit a model-native tool call. /// /// Device-only (requires a device running iOS 27.0+): loads real models. v1 family scope is /// Qwen3/QwQ (template renders tools AND honors `enable_thinking`); R1-Distill is -/// de-scoped (tool-blind template) and must fall through to the existing -/// single-phase tool path unchanged. +/// de-scoped (tool-blind template) and stays on its single-phase behavior. @Suite(.serialized, .timeLimit(.minutes(15))) struct ToolCallingReasoningTests { @@ -79,23 +79,16 @@ struct ToolCallingReasoningTests { private func collect(_ stream: TestResponseStream) async throws -> Collected { var c = Collected() for try await event in stream { - if let r = event as? LanguageModelExecutorGenerationChannel.Reasoning, - case .appendText(let fragment) = r.action - { - c.reasoning += fragment.content - } else if let t = event as? LanguageModelExecutorGenerationChannel.ToolCalls, - case .toolCall(let toolCall) = t.action, - case .appendArguments(let argsDelta) = toolCall.action - { + if case .appendText(let chunk, _, .reasoning) = event { + c.reasoning += chunk + } else if case .toolCall(_, let name, let arguments) = event { if c.toolCallName == nil { - c.toolCallName = toolCall.name + c.toolCallName = name c.reasoningBeforeToolCall = !c.reasoning.isEmpty } - c.toolArgs += argsDelta.content - } else if let r = event as? LanguageModelExecutorGenerationChannel.Response, - case .appendText(let fragment) = r.action - { - c.response += fragment.content + c.toolArgs += arguments + } else if case .appendText(let chunk, _, .response) = event { + c.response += chunk } } return c @@ -197,19 +190,69 @@ struct ToolCallingReasoningTests { /// crashing the serialized suite. @Test func cancellationDuringReasoningUnwinds() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } - let model = makeTestModel(Models.qwen3) + let model = makeReasoningTestModel(Models.qwen3) let executor = try makeMLXExecutor(for: model) let request = makeExecutorRequest( transcript: weatherTranscript(), enabledTools: [Self.weatherTool()], generationOptions: GenerationOptions(maximumResponseTokens: 1024)) - let stream = try await executeResponse(executor, request: request, model: model) - var events = 0 - for try await _ in stream { - events += 1 - if events >= 2 { break } // early break → respond is cancelled mid-flight + let stream = try await executeResponse( + executor, + request: request, + model: model, + cancelProducerWhen: { event in + if case .appendText(let text, _, .reasoning) = event { + return !text.isEmpty + } + return false + }) + var iterator = stream.makeAsyncIterator() + var sawReasoning = false + var sawCancellation = false + do { + while let event = try await iterator.next() { + if case .appendText(let text, _, .reasoning) = event, !text.isEmpty { + sawReasoning = true + } + } + } catch is CancellationError { + sawCancellation = true + } + await stream.cancelAndWait() + #expect(sawReasoning, "must cancel from inside native allowed reasoning") + #expect(sawCancellation, "native generation must terminate by cancellation, not EOS") + } + + /// Cancellation at the exact reasoning-close boundary must be observed + /// before required guided tool generation begins. + @Test func requiredCancellationAtReasoningClosePropagates() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + let model = makeReasoningTestModel(Models.qwen3) + let executor = try makeMLXExecutor(for: model) + let request = makeExecutorRequest( + transcript: weatherTranscript(), + enabledTools: [Self.weatherTool()], + generationOptions: GenerationOptions( + maximumResponseTokens: 1024, + toolCallingMode: .required)) + let sink = GuidedGenerationDiagnosticSink(cancelOnToolReasoningClose: true) + let stream = try await executeResponse( + executor, + request: request, + model: model, + guidedGenerationSink: sink) + + var sawCancellation = false + do { + for try await _ in stream {} + } catch is CancellationError { + sawCancellation = true } - #expect(events >= 1) + await stream.cancelAndWait() + + #expect(sink.toolReasoningCloseCount == 1, "reasoning must reach its close boundary") + #expect(sink.emitCount == 0, "cancellation must stop before guided tool generation emits") + #expect(sawCancellation, "reasoning-close cancellation must not become normal EOS") } } diff --git a/IntegrationTesting/IntegrationTestingTests/MTPDrafterModelFactoryIntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/MTPDrafterModelFactoryIntegrationTests.swift index d9ee9a8f5..07ad4e366 100644 --- a/IntegrationTesting/IntegrationTestingTests/MTPDrafterModelFactoryIntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MTPDrafterModelFactoryIntegrationTests.swift @@ -1,30 +1,30 @@ // Copyright © 2026 Apple Inc. import Foundation +import HuggingFace import IntegrationTestHelpers +import MLXHuggingFace import MLXLMCommon import MLXVLM import Testing -// MARK: - Factory load (gated on checkpoint presence) +// MARK: - Factory load (auto-downloads the checkpoint if not cached) + +/// Pinned checkpoint revision matching the weights that were live when the +/// Rung 4 `drafter_block` fixtures were generated. Kept in sync with +/// `MTPRung4TokenParityTests`. +private let drafter31BRevision = "28e92270316e89288579ec59c17939541d9ca433" @Test func testMTPDrafterFactoryLoadFromDirectoryWhenCheckpointPresent() async throws { - // Look for the 31B-assistant-bf16 checkpoint in the HF cache; skip - // gracefully when absent. - guard - let snapshot = hfSnapshotDir( - modelId: "mlx-community/gemma-4-31B-it-assistant-bf16") - else { - Issue.record("31B-assistant-bf16 checkpoint not in HF cache; skipping factory load test") - return - } - await Gemma4AssistantRegistration.register() let factory = MTPDrafterModelFactory.shared let container = try await factory.loadContainer( - from: snapshot, using: NoOpTokenizerLoader() + from: #hubDownloader(), + using: NoOpTokenizerLoader(), + configuration: .init( + id: "mlx-community/gemma-4-31B-it-assistant-bf16", revision: drafter31BRevision) ) let isDrafter = await container.perform { ctx in ctx.model is Gemma4AssistantDraftModel diff --git a/IntegrationTesting/IntegrationTestingTests/MTPIteratorEndToEndDiagnosticTests.swift b/IntegrationTesting/IntegrationTestingTests/MTPIteratorEndToEndDiagnosticTests.swift index bca0caca7..c443fe94f 100644 --- a/IntegrationTesting/IntegrationTestingTests/MTPIteratorEndToEndDiagnosticTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MTPIteratorEndToEndDiagnosticTests.swift @@ -47,27 +47,47 @@ private struct EmissionLog: LogitProcessor { } } +private let target31BModelId = "mlx-community/gemma-4-31b-it-8bit" +private let drafter31BModelId = "mlx-community/gemma-4-31B-it-assistant-bf16" + +/// Pinned checkpoint revisions matching the weights that were live when the +/// Rung 4 `drafter_block` fixtures were generated. Kept in sync with +/// `MTPRung4TokenParityTests` so this diagnostic suite exercises the same +/// bit-exact-verified checkpoint pair. +private let target31BRevision = "fe92291011fc698452920c0b558b52f790dff711" +private let drafter31BRevision = "28e92270316e89288579ec59c17939541d9ca433" + +/// Shared downloader for the 31B target+drafter pair. Fetches to the local +/// HF cache on first use (~35GB); subsequent tests and runs reuse the cache. +private let downloader: any Downloader = #hubDownloader() + private func loadTargetAndDrafter( targetModelId: String, drafterModelId: String, + targetRevision: String? = nil, + drafterRevision: String? = nil, drafterConfigType: any Codable.Type = Gemma4AssistantConfiguration.self -) async throws -> LoadedPair? { - guard let targetDir = hfSnapshotDir(modelId: targetModelId) else { return nil } - guard let drafterDir = hfSnapshotDir(modelId: drafterModelId) else { return nil } - - // Target — VLM factory's directory-load path gives us a full ModelContext - // (model + tokenizer + processor + configuration) without needing a - // Downloader. The 8-bit variant isn't in VLMRegistry but the factory - // resolves model type from `config.json` at the directory root. +) async throws -> LoadedPair { + // Target — VLM factory's downloader-based load path gives us a full + // ModelContext (model + tokenizer + processor + configuration), + // auto-fetching the snapshot from the HF Hub if it isn't cached yet. let context = try await VLMModelFactory.shared.load( - from: targetDir, - using: #huggingFaceTokenizerLoader() + from: downloader, + using: #huggingFaceTokenizerLoader(), + configuration: .init(id: targetModelId, revision: targetRevision ?? "main") ) // Drafter — same pattern as `MTPRung4TokenParityTests.loadRung4Drafter`: - // decode the Gemma4Assistant config, instantiate, then `loadWeights`. - // No explicit binding needed: the iterator passes the target through - // `draftBlock(target:...)` per round. + // fetch the snapshot, decode the Gemma4Assistant config, instantiate, + // then `loadWeights`. No explicit binding needed: the iterator passes + // the target through `draftBlock(target:...)` per round. + let drafterDir = try await downloader.download( + id: drafterModelId, + revision: drafterRevision, + matching: ["*.safetensors", "*.json"], + useLatest: false, + progressHandler: { _ in } + ) let cfg = try JSONDecoder().decode( Gemma4AssistantConfiguration.self, from: Data(contentsOf: drafterDir.appendingPathComponent("config.json"))) @@ -95,17 +115,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// after a real number is on the table. @Test func testMTP31BPairProducesAcceptedDrafts() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") @@ -187,17 +202,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// (n=66 vs n=29) and is the better regression gate. @Test func testMTP31BBlockSize4AcceptanceLifted() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") @@ -261,17 +271,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// pre-fix baseline while leaving ~8pp headroom for variance. @Test func testMTP31BBlockSize6AcceptanceLifted() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") @@ -332,17 +337,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// output stream silently started 1 or 2 positions ahead of baseline. @Test func testMTP31BMatchesBaselineByteIdentical() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let prompt = "Why is the sky blue? Explain in one paragraph." let userInput = UserInput(chat: [.user(prompt)]) @@ -461,17 +461,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// configuration. @Test func testMTP31BLongerStreamProducesValidContinuation() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let prompt = "Why is the sky blue? Explain in one paragraph." let userInput = UserInput(chat: [.user(prompt)]) @@ -595,17 +590,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// this stream length and the Phase B "failure" is reframed. @Test func testBaselineSelfDeterminism128() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let prompt = "Why is the sky blue? Explain in one paragraph." let userInput = UserInput(chat: [.user(prompt)]) @@ -661,17 +651,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// quantization triggers around round 3. @Test func testMTP31BQuantizationOnsetEngagesPassthrough() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") @@ -762,17 +747,12 @@ struct MTPIteratorEndToEndDiagnosticTests { /// position from the recorded-sequence start. @Test func testMTPLogitProcessorReceivesOnlyEmittedTokens() async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") @@ -912,17 +892,12 @@ struct MTPIteratorEndToEndDiagnosticTests { maxTokens: Int, blockSize: Int = 4 ) async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: "mlx-community/gemma-4-31b-it-8bit", - drafterModelId: "mlx-community/gemma-4-31B-it-assistant-bf16" - ) - else { - Issue.record( - "required checkpoint not in HF cache (31B 8-bit target or 31B drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: target31BModelId, + drafterModelId: drafter31BModelId, + targetRevision: target31BRevision, + drafterRevision: drafter31BRevision + ) let userInput = UserInput(chat: [.user(prompt)]) let lmInput = try await loaded.context.processor.prepare(input: userInput) @@ -995,17 +970,10 @@ struct MTPIteratorEndToEndDiagnosticTests { targetModelId: String, drafterModelId: String ) async throws { - guard - let loaded = try await loadTargetAndDrafter( - targetModelId: targetModelId, - drafterModelId: drafterModelId - ) - else { - Issue.record( - "required checkpoint not in HF cache (\(label) target or \(label) drafter); skipping" - ) - return - } + let loaded = try await loadTargetAndDrafter( + targetModelId: targetModelId, + drafterModelId: drafterModelId + ) let userInput = UserInput(chat: [ .user("Why is the sky blue? Explain in one paragraph.") diff --git a/IntegrationTesting/IntegrationTestingTests/MTPRung4TokenParityTests.swift b/IntegrationTesting/IntegrationTestingTests/MTPRung4TokenParityTests.swift index 2ed92e80e..ebf00c7c3 100644 --- a/IntegrationTesting/IntegrationTestingTests/MTPRung4TokenParityTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MTPRung4TokenParityTests.swift @@ -1,8 +1,10 @@ // Copyright © 2026 Apple Inc. import Foundation +import HuggingFace import IntegrationTestHelpers import MLX +import MLXHuggingFace import MLXLMCommon import MLXVLM import Testing @@ -49,9 +51,9 @@ private struct Rung4BoundDrafter { let target: Gemma4 } -nonisolated(unsafe) private var _rung4Cache: Result? +nonisolated(unsafe) private var _rung4Cache: Result? -private func sharedBoundDrafter() throws -> Rung4BoundDrafter? { +private func sharedBoundDrafter() async throws -> Rung4BoundDrafter { if let cached = _rung4Cache { switch cached { case .success(let value): return value @@ -59,7 +61,7 @@ private func sharedBoundDrafter() throws -> Rung4BoundDrafter? { } } do { - let result = try loadRung4Drafter() + let result = try await loadRung4Drafter() _rung4Cache = .success(result) return result } catch { @@ -68,15 +70,28 @@ private func sharedBoundDrafter() throws -> Rung4BoundDrafter? { } } -private func loadRung4Drafter() throws -> Rung4BoundDrafter? { - guard let drafterDir = hfSnapshotDir(modelId: "mlx-community/gemma-4-31B-it-assistant-bf16") - else { - return nil - } - guard let targetDir = hfSnapshotDir(modelId: "mlx-community/gemma-4-31b-it-8bit") - else { - return nil - } +private let rung4DrafterModelId = "mlx-community/gemma-4-31B-it-assistant-bf16" +private let rung4TargetModelId = "mlx-community/gemma-4-31b-it-8bit" + +/// Pinned checkpoint revisions matching the weights that were live when the +/// `drafter_block` fixtures were generated. Pinning keeps the bit-exact +/// parity assertions reproducible as the published checkpoints move. +private let rung4DrafterRevision = "28e92270316e89288579ec59c17939541d9ca433" +private let rung4TargetRevision = "fe92291011fc698452920c0b558b52f790dff711" + +/// Shared downloader for the Rung 4 target+drafter pair. Fetches to the +/// local HF cache on first use; subsequent tests and runs reuse the cache. +private let downloader: any Downloader = #hubDownloader() + +private func loadRung4Drafter() async throws -> Rung4BoundDrafter { + let drafterDir = try await downloader.download( + id: rung4DrafterModelId, revision: rung4DrafterRevision, + matching: ["*.safetensors", "*.json"], + useLatest: false, progressHandler: { _ in }) + let targetDir = try await downloader.download( + id: rung4TargetModelId, revision: rung4TargetRevision, + matching: ["*.safetensors", "*.json"], + useLatest: false, progressHandler: { _ in }) // Drafter — bf16, no quantization. let drafterCfg = try JSONDecoder().decode( @@ -137,12 +152,7 @@ private func assertDraftBlockMatchesFixture(name: String) async throws { guard let fixturesDir = await mtpFixturesDirOrSkip(name: name) else { return } - guard let bound = try sharedBoundDrafter() else { - Issue.record( - "required checkpoint not in HF cache (drafter or 31B 8-bit target); skipping Rung 4 \(name)" - ) - return - } + let bound = try await sharedBoundDrafter() let model = bound.drafter let target = bound.target diff --git a/IntegrationTesting/IntegrationTestingTests/VisionIntegrationTests.swift b/IntegrationTesting/IntegrationTestingTests/VisionIntegrationTests.swift index 51abd092f..043e78dd3 100644 --- a/IntegrationTesting/IntegrationTestingTests/VisionIntegrationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/VisionIntegrationTests.swift @@ -8,7 +8,7 @@ import Testing @testable import MLXFoundationModels -#if FoundationModelsIntegration +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) /// Opt-in end-to-end VLM test: drives a real `Qwen3-VL-4B-Instruct-4bit` through /// the FoundationModels adapter with a labeled image attachment and `.vision` diff --git a/Libraries/MLXFoundationModels/GuidedGeneration/SchemaConverter.swift b/Libraries/MLXFoundationModels/GuidedGeneration/SchemaConverter.swift index 32fd3da0c..010d2894e 100644 --- a/Libraries/MLXFoundationModels/GuidedGeneration/SchemaConverter.swift +++ b/Libraries/MLXFoundationModels/GuidedGeneration/SchemaConverter.swift @@ -48,10 +48,17 @@ enum SchemaConverter { /// } /// }, /// ... - /// ] + /// ], + /// "$defs": { "__": ... } /// } /// ``` /// + /// If a tool's parameters schema carries `$defs` (named sub-schemas such + /// as nested `@Generable` types), they are hoisted to the envelope root + /// under per-tool namespaced keys, with that tool's `$ref`s rewritten to + /// match — JSON Pointers resolve from the document root, so defs left + /// nested inside `arguments` would leave every ref dangling. + /// /// This is the *inner* schema -- it describes one tool call JSON object. /// For end-to-end grammar generation that also encodes the model's native /// tool-call wrapper (e.g. Qwen's `...`), see @@ -73,12 +80,12 @@ enum SchemaConverter { /// Builds an xgrammar structural-tag JSON that constrains the model /// to emit a tool call either wrapped in Qwen-style - /// `...` delimiters or as bare JSON. The - /// inner JSON is the envelope produced by - /// `toolCallingEnvelopeObject` (and serialized by - /// `encodeToolCallingEnvelopeJSON`). + /// `...` delimiters or as bare JSON. /// - /// Structural-tag shape: + /// Each tool is its own `tag` whose `begin` is the literal prefix + /// `{"name": "", "arguments": ` and whose `content` is that + /// tool's parameter schema, closed by an `end` of `}`. Structural-tag + /// shape: /// ```json /// { /// "type": "structural_tag", @@ -88,60 +95,85 @@ enum SchemaConverter { /// { /// "type": "tag", /// "begin": "\n", - /// "content": { "type": "json_schema", "json_schema": }, + /// "content": , /// "end": ["\n"] /// }, - /// { "type": "json_schema", "json_schema": } + /// /// ] /// } /// } /// ``` + /// where `` is an `or` over one `tag` per tool: + /// ```json + /// { + /// "type": "tag", + /// "begin": "{\"name\": \"set_flashlight\", \"arguments\": ", + /// "content": { "type": "json_schema", "json_schema": }, + /// "end": ["}"] + /// } + /// ``` /// - /// Accepting both alternatives lets the model stay in its trained - /// distribution — Qwen-family models overwhelmingly prefer the - /// wrapped form; the bare arm is a defensive fallback for models - /// that were trained on raw JSON and happen to share the envelope - /// shape. + /// **Why per-tool tags instead of one `oneOf` json_schema.** The + /// earlier shape embedded a single `{oneOf: [{name, arguments}, …]}` + /// json_schema in each arm. The structural-tag path compiles that + /// embedded schema with xgrammar's default (non-strict) property + /// ordering, so greedy decoding could open `"arguments"` before + /// `"name"` and dive into an unbounded free-text field before ever + /// committing to a tool -- producing a nameless, unparseable buffer + /// that ran the token budget dry (observed: Qwen filling `response` + /// with `"1234567890…"`). Making the name a literal tag prefix forces + /// the model to commit to a specific tool first, then fill only that + /// tool's arguments. It also removes the JSON whitespace wiggle room + /// around the `name`/`arguments` keys that open-source models tend to + /// exploit into long whitespace runs. /// - /// **Why structural tag over hand-rolled GBNF.** The envelope is a - /// JSON object whose shape depends on the tool's `parameters` - /// schema, which varies per tool. Emitting GBNF would require a - /// Swift-side JSON-schema-to-GBNF compiler — reinventing exactly - /// what xgrammar's `Grammar::FromJSONSchema` already does in C++. - /// Structural tag is xgrammar's first-class API for this - /// multi-format dispatch case; we assemble the dispatch shape in - /// Swift and let xgrammar compile the embedded JSON schema the - /// same way the plain `jsonSchema:` path does. + /// Accepting both wrapped and bare arms lets the model stay in its + /// trained distribution -- Qwen-family models overwhelmingly prefer + /// the wrapped form; the bare arm is a defensive fallback for models + /// trained on raw JSON. /// - /// **Why string literals, not special-token references.** The more - /// idiomatic structural-tag form for Qwen would use a - /// `TokenFormat` for `` / `` (Qwen encodes - /// them as single special tokens). That would require threading - /// the bound `GrammarTokenizer` through to `Grammar::FromStructuralTag` - /// for token-string resolution, which the shim entry point - /// (`xg_compile_structural_tag`) currently declines to do. The - /// plain-string form is equivalent at the byte level: xgrammar - /// matches the byte sequence `` against the vocab - /// mask, finds Qwen's `` special token (whose decoded - /// bytes are exactly that string), and accepts it. + /// **Why structural tag over hand-rolled GBNF.** Each tool's + /// `arguments` is a JSON object whose shape depends on the tool's + /// `parameters` schema. Emitting GBNF for it would require a + /// Swift-side JSON-schema-to-GBNF compiler -- reinventing what + /// xgrammar's `Grammar::FromJSONSchema` already does in C++. + /// Structural tag composes the fixed dispatch prefix with the + /// per-tool json_schema and lets xgrammar compile the embedded + /// schema the same way the plain `jsonSchema:` path does. /// /// Requires a non-empty tool list. static func encodeToolCallingGrammar( tools: [Transcript.ToolDefinition] ) throws -> String { - let envelope = try toolCallingEnvelopeObject(tools: tools) + guard !tools.isEmpty else { + throw SchemaConversionError.noTools + } - // `json_schema` entries must embed the schema as an inline - // JSON *object*, not a stringified schema — xgrammar's - // structural-tag parser rejects stringified schemas outright - // (see `StructuralTagParser::ParseJSONSchemaFormat`). The - // envelope is already an `[String: Any]`; pass the same - // reference into both `or.elements` arms so the emitted JSON - // round-trips identically on the wrapped and bare sides. - let jsonSchemaFormat: [String: Any] = [ - "type": "json_schema", - "json_schema": envelope, + let encoder = JSONEncoder() + // One tag per tool. `begin` fixes `{"name": "", "arguments": ` + // so the tool name is committed before the arguments schema opens; + // `content` is the tool's parameter schema; `end` closes the object. + let toolTags: [[String: Any]] = try tools.map { tool in + let paramsData = try encoder.encode(tool.parameters) + let paramsAny = try JSONSerialization.jsonObject(with: paramsData) + let nameData = try JSONSerialization.data( + withJSONObject: tool.name, options: [.fragmentsAllowed]) + let nameJSON = String(data: nameData, encoding: .utf8) ?? "\"\(tool.name)\"" + return [ + "type": "tag", + "begin": "{\"name\": \(nameJSON), \"arguments\": ", + "content": [ + "type": "json_schema", + "json_schema": paramsAny, + ], + "end": ["}"], + ] + } + let perToolOr: [String: Any] = [ + "type": "or", + "elements": toolTags, ] + let structuralTag: [String: Any] = [ "type": "structural_tag", "format": [ @@ -150,10 +182,10 @@ enum SchemaConverter { [ "type": "tag", "begin": "\n", - "content": jsonSchemaFormat, + "content": perToolOr, "end": ["\n"], ], - jsonSchemaFormat, + perToolOr, ] as [Any], ] as [String: Any], ] @@ -176,12 +208,36 @@ enum SchemaConverter { } let encoder = JSONEncoder() + // `GenerationSchema` serializes named sub-schemas (e.g. a nested + // `@Generable` type, or a named `DynamicGenerationSchema`) as + // root-level `$defs` plus root-anchored `"$ref": "#/$defs/..."` + // pointers. Embedding a tool's schema as a nested object under + // `oneOf[i].properties.arguments` buries its `$defs` inside + // `arguments` while the refs stay anchored to the document root — + // and xgrammar resolves JSON Pointers from the document root, so + // every ref dangles and grammar compilation hard-fails + // ("Cannot find field $defs in {\"oneOf\": ...", + // json_schema_converter.cc). Hoist each tool's `$defs` to the + // envelope root instead, namespacing keys per tool + // (`__`) so same-named defs across tools cannot collide. + var hoistedDefs: [String: Any] = [:] let oneOf: [[String: Any]] = try tools.map { tool in // Round-trip the tool's parameters through JSONSerialization so we // can embed it as a nested object in the envelope we assemble via // JSONSerialization.data(withJSONObject:). Cheap: schemas are small. let paramsData = try encoder.encode(tool.parameters) - let paramsAny = try JSONSerialization.jsonObject(with: paramsData) + var paramsAny = rewriteDefsRefs( + in: try JSONSerialization.jsonObject(with: paramsData), + toolName: tool.name + ) + if var paramsObj = paramsAny as? [String: Any] { + if let defs = paramsObj.removeValue(forKey: "$defs") as? [String: Any] { + for (key, value) in defs { + hoistedDefs["\(tool.name)__\(key)"] = value + } + } + paramsAny = paramsObj + } return [ "type": "object", "required": ["name", "arguments"], @@ -192,7 +248,40 @@ enum SchemaConverter { ], ] } - return ["oneOf": oneOf] + var envelope: [String: Any] = ["oneOf": oneOf] + if !hoistedDefs.isEmpty { + envelope["$defs"] = hoistedDefs + } + return envelope + } + + /// Rewrites every `"$ref": "#/$defs/"` in a parsed schema tree to + /// the per-tool namespaced key (`#/$defs/__`). + /// + /// The rewrite is structure-aware: only the string value directly under a + /// `$ref` key is touched, so other strings that merely mention the + /// pointer text — a `description`, `const`, `enum` entry, `default`, or + /// `pattern` containing "#/$defs/" — survive verbatim. Runs before the + /// `$defs` are hoisted out, and recurses through the whole tree because + /// refs can appear anywhere, including inside other `$defs` bodies. + private static func rewriteDefsRefs(in value: Any, toolName: String) -> Any { + switch value { + case let object as [String: Any]: + var result: [String: Any] = [:] + result.reserveCapacity(object.count) + for (key, nested) in object { + if key == "$ref", let ref = nested as? String, ref.hasPrefix("#/$defs/") { + result[key] = "#/$defs/\(toolName)__" + ref.dropFirst("#/$defs/".count) + } else { + result[key] = rewriteDefsRefs(in: nested, toolName: toolName) + } + } + return result + case let array as [Any]: + return array.map { rewriteDefsRefs(in: $0, toolName: toolName) } + default: + return value + } } enum SchemaConversionError: Error { diff --git a/Libraries/MLXFoundationModels/MLXLanguageModel.swift b/Libraries/MLXFoundationModels/MLXLanguageModel.swift index a663d6f5e..901aa95eb 100644 --- a/Libraries/MLXFoundationModels/MLXLanguageModel.swift +++ b/Libraries/MLXFoundationModels/MLXLanguageModel.swift @@ -502,9 +502,9 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // MARK: - LanguageModel Conformance /// MLX supports guided generation via xgrammar grammar-constrained - /// decoding (provided by the MLXGuidedGeneration library), tool - /// calling via the synthetic-final-answer envelope, and reasoning - /// (chain-of-thought) routing on the unconstrained generation path. + /// decoding (provided by the MLXGuidedGeneration library), native + /// `.allowed` tool routing, guided `.required` developer tool calls, + /// and reasoning (chain-of-thought) routing. /// /// Capabilities are declared explicitly by the caller at ``init(configuration:capabilities:configurationResolver:weightsLocation:load:)`` /// and stored verbatim. The caller includes @@ -665,6 +665,114 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// Executes inference requests for the model. public struct Executor: LanguageModelExecutor, Sendable { + // MARK: - Test observation hook + // + // The macOS 27 FoundationModels SDK made the generation-channel event + // and action types opaque: a consumer can no longer read back what was + // streamed. Tests need to read it, and the only place the content is + // available is here, right before it enters the channel. These emit + // helpers are the sole send sites for each event kind; each notifies an + // optional observer with a readable mirror. The observer is nil in + // shipping builds (only tests attach one via the task-local), so the + // arguments handed to `channel.send` are identical to before and + // behavior is unchanged. + + /// Readable, internal-only mirror of the events this executor streams + /// into the opaque FoundationModels channel. + enum GenerationEvent: Sendable { + enum Destination: Sendable { case response, reasoning } + case appendText(String, entryID: String?, destination: Destination) + case toolCall(id: String, name: String, arguments: String) + case updateMetadata([String: any Sendable & Codable & Equatable], entryID: String?) + case updateUsage( + input: LanguageModelExecutorGenerationChannel.Usage.Input, + output: LanguageModelExecutorGenerationChannel.Usage.Output, + entryID: String?) + } + + /// Attached only by tests (via `$generationObserver.withValue`); nil in + /// shipping. Task-local so it reaches child tasks that also emit (e.g. + /// the guided-generation text forwarder). + @TaskLocal static var generationObserver: (@Sendable (GenerationEvent) -> Void)? + + static func emit( + text: String, entryID: String?, destination: GenerationEvent.Destination, + into channel: LanguageModelExecutorGenerationChannel + ) async { + generationObserver?(.appendText(text, entryID: entryID, destination: destination)) + switch destination { + case .response: + await channel.send( + .response(entryID: entryID, action: .appendText(text, tokenCount: 1))) + case .reasoning: + await channel.send( + .reasoning(entryID: entryID, action: .appendText(text, tokenCount: 1))) + } + } + + static func emitMetadata( + _ values: [String: any Sendable & Codable & Equatable], entryID: String?, + into channel: LanguageModelExecutorGenerationChannel + ) async { + generationObserver?(.updateMetadata(values, entryID: entryID)) + await channel.send(.response(entryID: entryID, action: .updateMetadata(values))) + } + + static func emitUsage( + input: LanguageModelExecutorGenerationChannel.Usage.Input, + output: LanguageModelExecutorGenerationChannel.Usage.Output, + entryID: String?, + into channel: LanguageModelExecutorGenerationChannel + ) async { + generationObserver?(.updateUsage(input: input, output: output, entryID: entryID)) + + // TODO: papering over an FM-27 SDK symbol drift -- restore + // the channel usage send (the commented-out call at the end of this + // block) once the shipping dylib matches its own interface. + // + // Usage is intentionally NOT forwarded to the FoundationModels + // channel on this SDK. The FM-27 beta `.swiftinterface` declares + // Response.Action.updateUsage(input:output:metadata: = [:]) + // (three parameters), but the shipping FoundationModels dylib only + // exports the older two-parameter + // Response.Action.updateUsage(input:output:) + // Because our call relies on the `metadata:` default, the compiler + // resolves it to the three-parameter symbol, which does not exist + // at runtime. dyld cannot bind it: under chained-fixups linking + // (the arm64 default) the reference aborts the process the moment + // the image loads, and under lazy binding it faults through null + // (SIGSEGV at 0x0) the instant this send executes -- crashing every + // `respond()` path right after generation completes. + // + // A runtime `dlsym` guard cannot save this: the compiled reference + // to the missing symbol is enough to abort at launch regardless of + // any surrounding check. The only safe option is to not reference + // the symbol at all, so no `channel.send(.updateUsage(...))` here. + // + // Effect: the framework does not receive our per-response usage + // event, so consumer-visible usage for these responses may be + // absent or zero. Tests still observe usage through + // `generationObserver` above. When a later SDK ships a dylib that + // matches its interface, restore the send: + // await channel.send( + // .response( + // entryID: entryID, + // action: .updateUsage(input: input, output: output))) + } + + static func emitToolCall( + id: String, name: String, arguments: String, entryID: String, + into channel: LanguageModelExecutorGenerationChannel + ) async { + generationObserver?(.toolCall(id: id, name: name, arguments: arguments)) + await channel.send( + .toolCalls( + entryID: entryID, + action: .toolCall( + id: id, name: name, + action: .appendArguments(arguments, tokenCount: 1)))) + } + /// Default `maxTokens` when the caller doesn't set /// `GenerationOptions.maximumResponseTokens`. Applied uniformly /// across guided-JSON, tool-calling, and unconstrained generation @@ -697,46 +805,41 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { return Float(max(0, value)) } - /// Translate FoundationModels' `GenerationOptions.SamplingMode` into the - /// backend-local `MLXSamplingMode`, dropping the best-effort `seed` - /// (MLX's samplers expose no seed-injection hook). No mode set (`nil`) - /// and any future/unknown `Kind` both map to `nil` -- "use the provider - /// default" -- so an unrecognized case never traps and never reaches the - /// resolver. All value policy lives in `resolveSamplingParameters`; this - /// shim is a pure 1:1 case translation. - static func samplingMode( + /// Translate Foundation Models' `GenerationOptions.SamplingMode` into one + /// backend-local value that preserves both the sampling strategy and optional + /// `UInt64` seed. No mode set (`nil`) and any future/unknown `Kind` both map to + /// `nil`, selecting provider-default behavior without trapping or guessing. + static func samplingConfiguration( from samplingMode: GenerationOptions.SamplingMode? - ) -> MLXSamplingMode? { + ) -> MLXSamplingConfiguration? { guard let kind = samplingMode?.kind else { return nil } switch kind { case .greedy: - return .greedy - case .top(let k, _): - return .topK(k) - case .nucleus(let threshold, _): - return .nucleus(threshold) + return MLXSamplingConfiguration(mode: .greedy, seed: nil) + case .randomTopK(let k, let seed): + return MLXSamplingConfiguration(mode: .topK(k), seed: seed) + case .randomProbabilityThreshold(let threshold, let seed): + return MLXSamplingConfiguration(mode: .nucleus(threshold), seed: seed) @unknown default: return nil } } - /// Build the `GenerateParameters` for a generation pass, threading the - /// caller's temperature and sampling mode through the shared resolver so - /// every real-sampler path (unconstrained, reasoning, tool-call - /// reasoning) honors `samplingMode` identically. `maxTokens` is the - /// already-resolved budget -- callers keep their own default/budget - /// arithmetic, so this helper owns only temperature + sampling resolution. + /// Build `GenerateParameters` for a sampler-backed generation pass. The shared + /// resolver owns temperature and mode precedence; this helper preserves the + /// optional seed directly on the backend request. static func makeParameters( maxTokens: Int, requestedTemperature: Double?, - samplingMode: MLXSamplingMode? + samplingConfiguration: MLXSamplingConfiguration? ) -> GenerateParameters { - var params = GenerateParameters(maxTokens: maxTokens) + var parameters = GenerateParameters(maxTokens: maxTokens) resolveSamplingParameters( - mode: samplingMode, + mode: samplingConfiguration?.mode, clampedTemperature: clampedTemperature(requestedTemperature) - ).apply(to: ¶ms) - return params + ).apply(to: ¶meters) + parameters.seed = samplingConfiguration?.seed + return parameters } /// Map xgrammar errors to typed `LanguageModelError` cases where the @@ -862,6 +965,12 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { )) } + let toolCallingMode = ToolCallingModeResolution.resolve( + request.generationOptions.toolCallingMode) + let enabledToolDefinitions = try ToolCallingModeResolution.enabledToolDefinitions( + for: toolCallingMode, + from: request.enabledToolDefinitions) + let container = try await model.loadContainer() // Encode schema to JSON if present @@ -877,7 +986,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // Translate the SDK sampling mode once, here where generationOptions // is in scope; thread the bridge-local value down to every // real-sampler path so they honor it identically. - let requestedSamplingMode = Self.samplingMode( + let requestedSamplingConfiguration = Self.samplingConfiguration( from: request.generationOptions.samplingMode) // Per SKILL.md: response and tool-calls entries each need a fresh // UUID — they live in separate transcript entries. We preserve the @@ -895,13 +1004,9 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { do { // Send metadata first - await channel.send( - .response( - entryID: entryID, - action: .updateMetadata([ - "modelID": modelID, - "requestID": request.id.uuidString, - ]))) + await Self.emitMetadata( + ["modelID": modelID, "requestID": request.id.uuidString], + entryID: entryID, into: channel) // Generate tokens inside actor isolation. `messages` carries // non-Sendable `Chat.Message` instances (UserInput.Image and @@ -912,29 +1017,6 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { let userInput = UserInput(chat: messages) let input = try await context.processor.prepare(input: userInput) - // Single-turn tool-calling cap: if the transcript already - // contains prior tool-call or tool-output entries, this - // is a continuation round from `LanguageModelSession`'s - // auto-loop (it executed the tool and re-invoked us with - // the result appended). Our `TranscriptConverter` drops - // those entries, so re-entering the tool-calling branch - // would just make the model emit the same tool call - // again -- an infinite loop. Fall through to text - // generation so the session terminates cleanly after - // one round. - // - // Multi-turn tool calling -- where the model sees tool - // outputs in the transcript and continues with a - // data-aware response -- is not supported. - let isContinuationAfterToolCall = request.transcript.contains { entry in - switch entry { - case .instructions, .prompt, .response: return false - case .reasoning: return false - case .toolCalls, .toolOutput: return true - @unknown default: return true - } - } - // Resolve the per-instance configuration. Held strictly as // a local; it never lands in context.configuration or // Executor.Configuration, so two instances with the same id @@ -992,10 +1074,13 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // Reasoning is only consumed by the unconstrained path // (no tools, no schema). On the guided/tool paths the // grammar already constrains output, so suppression-prep - // would be wasted work. + // would be wasted work here. Continuation rounds run the + // tool path (below) like fresh turns: that path renders its + // own thinking state into the tool-aware prompt + // (`toolAwareContext`) -- thinking on with the think-then-call + // phase when reasoning is declared, forced off otherwise. let mayRunReasoningPath = - (request.enabledToolDefinitions.isEmpty - || isContinuationAfterToolCall) + enabledToolDefinitions.isEmpty && request.schema == nil // When .reasoning is OMITTED on the unconstrained path, @@ -1044,37 +1129,17 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // baseline `input` rendered above. let effectiveInput = suppressedInput ?? input - if !request.enabledToolDefinitions.isEmpty - && !isContinuationAfterToolCall - { - // Tool-calling path. Force the model to emit a JSON - // object matching one of the declared tools -- - // including a synthetic "final answer" tool whose - // arguments carry the free-text response. After - // generation, parse the output to route to either a - // toolCallDelta (real tool) or textDelta (final - // answer) event. - // - // Buffers the full output before emitting; streaming - // within the final-answer path (reparse-each-delta) is - // not yet implemented. - let finalAnswerDef = FinalAnswerTool.makeToolDefinition( - responseSchema: request.schema - ) - let allTools = - Array(request.enabledToolDefinitions) + [finalAnswerDef] - - // Re-tokenize using the model's native tool-aware chat + // Tool path, entered on every round while tools are enabled + // -- fresh turns and continuations alike. Allowed mode uses + // native generation so the model can answer or call a tool; + // required mode constrains generation to a real tool call. + if !enabledToolDefinitions.isEmpty { + // Re-render using the model's native tool-aware chat // template (Qwen/Llama/Phi/Gemma all ship one in their // tokenizer_config.json). This is what teaches the model // *what* tools exist and how to decide between them; the // grammar constraint below only enforces the *shape* of // whatever tool call it emits. - let toolSpecs = try ToolCallingConversions.makeToolSpecs( - from: allTools) - let tokenizerMessages = DefaultMessageGenerator().generate( - messages: messages) - // Think-then-call is gated to the enable_thinking // family (Qwen3/QwQ): their template both renders the tool // block AND honors `enable_thinking`. R1-style `.alwaysOn` @@ -1091,23 +1156,130 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { return cfg }() // Thread `enable_thinking` through the tool-aware template - // (3-arg form) so the prompt is both tool-aware and - // thinking-primed; nil on the single-phase path. - let reasoningContext = try thinkThenCallConfig.flatMap { - try $0.promptStrategy.additionalContext( - forThinkingEnabled: Self.thinkingEnabled( - for: request.contextOptions.reasoningLevel)) + // so the prompt's thinking state matches how we drive + // generation. For a toggleable model (`.templateFlag`, e.g. + // Qwen3 whose `enable_thinking` defaults ON) the effective + // value is: + // - reasoning declared: honor the requested level + // (default ON), and the think-then-call phase below + // lets the model reason before the grammar constrains it; + // - reasoning NOT declared: force thinking OFF, mirroring + // the unconstrained path's suppression (see + // `suppressedInput` above). + // Forcing OFF here is load-bearing for both modes. Native + // `.allowed` generation would otherwise surface undeclared + // reasoning through response/tool parsing. In `.required`, + // the grammar forces tool-call JSON from the first token, so + // a thinking-primed model cannot emit its `` block and + // greedy decoding of unconstrained developer tool string + // arguments can degenerate (Qwen: "1234567890..."). + // `.alwaysOn` models with reasoning undeclared were already + // rejected by the capability gate above; `.none`/no-config + // models take no context. + let toolAwareContext: [String: any Sendable]? + if case .templateFlag(let key, let defaultOn)? = + resolved.reasoningConfig?.promptStrategy + { + let enabled = + declaresReasoning + ? (Self.thinkingEnabled( + for: request.contextOptions.reasoningLevel) ?? defaultOn) + : false + toolAwareContext = [key: enabled] + } else { + toolAwareContext = nil + } + // Prepare through the model's UserInputProcessor (like the + // unconstrained and guided paths) instead of hand-building + // an LMInput from raw applyChatTemplate output: processors + // produce the token rank their model family requires (LLM + // processors emit [N]; VLM processors emit [1, N], and VLM + // `prepare` fatally aborts on 1-D input), and they carry + // image/video content through to the model. + if ToolCallingModeResolution.usesAllowedBehavior(toolCallingMode) { + let toolSpecs = try ToolCallingConversions.makeToolSpecs( + from: enabledToolDefinitions) + let toolAwareInput = try await context.processor.prepare( + input: UserInput( + chat: messages, + tools: toolSpecs, + additionalContext: toolAwareContext)) + let reasoning = thinkThenCallConfig.map { + ( + config: $0, + primedInside: Self.reasoningPrimedInside( + input: toolAwareInput, + config: $0, + tokenizer: context.tokenizer) + ) + } + let result = try await runAllowedToolGeneration( + input: toolAwareInput, + toolSpecs: toolSpecs, + reasoning: reasoning, + requestedMaxTokens: requestedMaxTokens, + requestedTemperature: request.generationOptions.temperature, + samplingConfiguration: requestedSamplingConfiguration, + reasoningEntryID: reasoningEntryID, + context: context, + channel: channel) + + if result.endedInsideReasoning { + await Self.emitMetadata( + ["incompleteOutput": true], entryID: entryID, into: channel) + } else if !result.toolCalls.isEmpty { + for call in result.toolCalls { + let argumentsData = try JSONEncoder().encode( + call.function.arguments) + let arguments = String( + decoding: argumentsData, as: UTF8.self) + await Self.emitToolCall( + id: call.id ?? UUID().uuidString, + name: call.function.name, + arguments: arguments, + entryID: toolCallsEntryID, + into: channel) + } + } else if let schemaJSON { + try await runSchemaGeneration( + schemaJSON: schemaJSON, + input: input, + modelID: modelID, + requestedMaxTokens: requestedMaxTokens, + entryID: entryID, + context: context, + channel: channel) + } else { + await Self.emit( + text: result.responseText, + entryID: entryID, + destination: .response, + into: channel) + } + if schemaJSON == nil || !result.toolCalls.isEmpty + || result.endedInsideReasoning + { + await emitAllowedUsage( + result, entryID: entryID, channel: channel) + } + Stream.gpu.synchronize() + return } - let toolAwareTokens = try context.tokenizer.applyChatTemplate( - messages: tokenizerMessages, - tools: toolSpecs, - additionalContext: reasoningContext - ) - let toolAwareInput = LMInput(tokens: MLXArray(toolAwareTokens)) + + // Required mode is the only mode that reaches guided + // tool generation, and it uses developer definitions only. + let requiredToolDefinitions = enabledToolDefinitions + let toolSpecs = try ToolCallingConversions.makeToolSpecs( + from: requiredToolDefinitions) + let toolAwareInput = try await context.processor.prepare( + input: UserInput( + chat: messages, + tools: toolSpecs, + additionalContext: toolAwareContext)) let toolCallingGrammar = try SchemaConverter.encodeToolCallingGrammar( - tools: allTools + tools: requiredToolDefinitions ) // The inner JSON envelope is still needed separately to // seed `CompletionReserve` -- the wrapper tokens @@ -1116,7 +1288,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // tokenized size adds noise rather than accuracy. let toolCallingEnvelopeJSON = try SchemaConverter.encodeToolCallingEnvelopeJSON( - tools: allTools + tools: requiredToolDefinitions ) let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( @@ -1169,7 +1341,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { primedInside: primedInside, maxTokens: maxTokens, requestedTemperature: request.generationOptions .temperature, - samplingMode: requestedSamplingMode, + samplingConfiguration: requestedSamplingConfiguration, reasoningEntryID: reasoningEntryID, responseEntryID: entryID, context: context, channel: channel) @@ -1179,12 +1351,8 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // ``). Don't prefill a truncated thought // into the grammar — signal and finish. Phase 1 // already synchronized the GPU on its way out. - await channel.send( - .response( - entryID: entryID, - action: .updateMetadata([ - "incompleteOutput": true - ]))) + await Self.emitMetadata( + ["incompleteOutput": true], entryID: entryID, into: channel) return } } @@ -1195,8 +1363,8 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { let phase2Input = reasoningTokenIDs.isEmpty ? toolAwareInput - : LMInput( - tokens: MLXArray(toolAwareTokens + reasoningTokenIDs)) + : Self.continuationInput( + from: toolAwareInput, appending: reasoningTokenIDs) // Shared budget (match the unconstrained path): the // envelope continues under the remaining budget, floored // at the completion reserve so it always has room to close @@ -1224,16 +1392,19 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { whitespaceTokenIDs: whitespaceTokenIDs ) { text in outputBuffer += text + GuidedGenerationDiagnosticSink.current?.recordEmit() return !Task.isCancelled } } catch GuidedGenerationError.incompleteOutput { incomplete = true } + try Task.checkCancellation() - try await emitToolCallingEvent( + GuidedGenerationDiagnosticSink.current?.recordBuffer( + outputBuffer, incompleteOutput: incomplete) + + await emitRequiredToolCallEvent( outputBuffer: outputBuffer, - userResponseSchema: request.schema, - entryID: entryID, toolCallsEntryID: toolCallsEntryID, channel: channel ) @@ -1244,149 +1415,36 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // clamped ≤ total. let reasoningCount = reasoningTokenIDs.count let totalOutput = generatedTokenCount + reasoningCount - await channel.send( - .response( - entryID: entryID, - action: .updateUsage( - input: .init( - totalTokenCount: toolAwareInput.text.tokens - .size, - cachedTokenCount: 0 - ), - output: .init( - totalTokenCount: totalOutput, - reasoningTokenCount: Swift.min( - reasoningCount, totalOutput) - ) - ) - )) + await Self.emitUsage( + input: .init( + totalTokenCount: toolAwareInput.text.tokens.size, + cachedTokenCount: 0), + output: .init( + totalTokenCount: totalOutput, + reasoningTokenCount: Swift.min(reasoningCount, totalOutput)), + entryID: entryID, into: channel) } if incomplete { - await channel.send( - .response( - entryID: entryID, - action: .updateMetadata(["incompleteOutput": true])) - ) + await Self.emitMetadata( + ["incompleteOutput": true], entryID: entryID, into: channel) } } else if let schemaJSON { - // Guided generation: stream text deltas as they arrive. - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( - modelID: modelID, - tokenizer: context.tokenizer - ) - - let constraint = try await MLXLanguageModel.makeConstraint( - modelID: modelID, - kind: .json, - source: schemaJSON, - tokenizer: xgTokenizer, - hostTokenizer: context.tokenizer, - fastForward: true - ) - // Bias and reserve computation: only when a token - // budget is set. Without a budget, the grammar mask - // and model's natural EOS tendency control termination. - let maxTokens = requestedMaxTokens ?? Self.defaultMaxTokens - let bias = await MLXLanguageModel.makeTokenizerBias( - modelID: modelID, - tokenizer: context.tokenizer - ) - let closingBias = bias.closing - let structuralReserve = CompletionReserve.estimate( + try await runSchemaGeneration( schemaJSON: schemaJSON, - tokenizer: context.tokenizer - ) - // The structural reserve is the bare minimum tokens for - // JSON skeleton (empty strings). Use the larger of 3x - // structural minimum or 25% of maxTokens, so closing - // bias activates early enough for the model to generate - // actual content in closing fields. - let completionReserve = Swift.max( - structuralReserve * 3, maxTokens / 4) - // Hard reserve: the point at which we force structural - // completion by penalizing non-closing tokens. Must be - // larger than the raw estimate because grammar-forced - // key names (FF tokens) and model-inserted whitespace - // cost more tokens than the compact minimal JSON string. - let hardReserve = structuralReserve * 8 - - let whitespaceBias = bias.whitespace - let whitespaceTokenIDs = bias.whitespaceTokenIDs - - // GuidedGenerationLoop.run's emit closure is synchronous (for - // performance -- it runs inside the tight MLX generation loop). - // channel.send is async. Bridge via an AsyncStream + concurrent - // forwarder so text deltas stream to the channel in order. - let (textStream, textContinuation) = AsyncStream - .makeStream() - async let forwarder: Void = { - for await text in textStream { - await channel.send( - .response( - entryID: entryID, - action: .appendText(text, tokenCount: 1) - )) - } - }() - - var incomplete = false - var generatedTokenCount: Int? - do { - generatedTokenCount = try GuidedGenerationLoop.run( - input: input, - context: context, - constraint: constraint, - maxTokens: maxTokens, - vocabSize: Int(xgTokenizer.vocabSize), - completionReserve: completionReserve, - hardReserve: hardReserve, - closingBias: closingBias, - whitespaceBias: whitespaceBias, - whitespaceTokenIDs: whitespaceTokenIDs - ) { text in - textContinuation.yield(text) - return !Task.isCancelled - } - } catch GuidedGenerationError.incompleteOutput { - // Grammar exhausted maxTokens before reaching a stop state. - // Text deltas already emitted are best-effort output. - incomplete = true - } - textContinuation.finish() - await forwarder - - if let generatedTokenCount { - await channel.send( - .response( - entryID: entryID, - action: .updateUsage( - input: .init( - totalTokenCount: input.text.tokens.size, - cachedTokenCount: 0 - ), - output: .init( - totalTokenCount: generatedTokenCount, - reasoningTokenCount: 0 - ) - ) - )) - } - - if incomplete { - await channel.send( - .response( - entryID: entryID, - action: .updateMetadata(["incompleteOutput": true])) - ) - } + input: input, + modelID: modelID, + requestedMaxTokens: requestedMaxTokens, + entryID: entryID, + context: context, + channel: channel) } else { try await runTextGeneration( reasoningSetup: reasoningSetup, fallbackInput: effectiveInput, requestedMaxTokens: requestedMaxTokens, requestedTemperature: request.generationOptions.temperature, - samplingMode: requestedSamplingMode, + samplingConfiguration: requestedSamplingConfiguration, responseEntryID: entryID, reasoningEntryID: reasoningEntryID, context: context, @@ -1414,13 +1472,220 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { } } + private struct AllowedToolGenerationResult { + var responseText = "" + var toolCalls: [MLXLMCommon.ToolCall] = [] + var completionInfo: GenerateCompletionInfo? + var reasoningTokenCount = 0 + var endedInsideReasoning = false + } + + private func runAllowedToolGeneration( + input: LMInput, + toolSpecs: [[String: any Sendable]], + reasoning: (config: ReasoningConfig, primedInside: Bool)?, + requestedMaxTokens: Int?, + requestedTemperature: Double?, + samplingConfiguration: MLXSamplingConfiguration?, + reasoningEntryID: String, + context: ModelContext, + channel: LanguageModelExecutorGenerationChannel + ) async throws -> AllowedToolGenerationResult { + let params = Self.makeParameters( + maxTokens: requestedMaxTokens ?? Self.defaultMaxTokens, + requestedTemperature: requestedTemperature, + samplingConfiguration: samplingConfiguration) + var router = AllowedToolOutputRouter( + format: context.configuration.toolCallFormat ?? .json, + tools: toolSpecs, + reasoning: reasoning) + var detokenizer = NaiveStreamingDetokenizer(tokenizer: context.tokenizer) + var result = AllowedToolGenerationResult() + let (stream, task) = try generateTokensTask( + input: input, parameters: params, context: context) + + do { + for await generation in stream { + try Task.checkCancellation() + switch generation { + case .token(let token): + if router.isInsideReasoning { + result.reasoningTokenCount += 1 + } + detokenizer.append(token: token) + if let chunk = detokenizer.next() { + let reasoningChunks = consumeAllowedEvents( + router.process(chunk), result: &result) + for text in reasoningChunks { + await Self.emit( + text: text, + entryID: reasoningEntryID, + destination: .reasoning, + into: channel) + try Task.checkCancellation() + } + } + case .info(let info): + result.completionInfo = info + } + } + } catch { + task.cancel() + await task.value + throw error + } + + await task.value + let finalReasoningChunks = consumeAllowedEvents( + router.finish(), result: &result) + for text in finalReasoningChunks { + await Self.emit( + text: text, + entryID: reasoningEntryID, + destination: .reasoning, + into: channel) + } + result.endedInsideReasoning = router.isInsideReasoning + return result + } + + private func consumeAllowedEvents( + _ events: [AllowedToolOutputRouter.Event], + result: inout AllowedToolGenerationResult + ) -> [String] { + var reasoningChunks: [String] = [] + for event in events { + switch event { + case .reasoning(let text): + reasoningChunks.append(text) + case .response(let text): + result.responseText += text + case .toolCall(let call): + result.toolCalls.append(call) + } + } + return reasoningChunks + } + + private func emitAllowedUsage( + _ result: AllowedToolGenerationResult, + entryID: String, + channel: LanguageModelExecutorGenerationChannel + ) async { + guard let info = result.completionInfo else { return } + await Self.emitUsage( + input: .init( + totalTokenCount: info.promptTokenCount, + cachedTokenCount: 0), + output: .init( + totalTokenCount: info.generationTokenCount, + reasoningTokenCount: min( + result.reasoningTokenCount, + info.generationTokenCount)), + entryID: entryID, + into: channel) + } + + private func runSchemaGeneration( + schemaJSON: String, + input: LMInput, + modelID: String, + requestedMaxTokens: Int?, + entryID: String, + context: ModelContext, + channel: LanguageModelExecutorGenerationChannel + ) async throws { + let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + modelID: modelID, + tokenizer: context.tokenizer) + let constraint = try await MLXLanguageModel.makeConstraint( + modelID: modelID, + kind: .json, + source: schemaJSON, + tokenizer: xgTokenizer, + hostTokenizer: context.tokenizer, + fastForward: true) + let maxTokens = requestedMaxTokens ?? Self.defaultMaxTokens + let bias = await MLXLanguageModel.makeTokenizerBias( + modelID: modelID, + tokenizer: context.tokenizer) + let structuralReserve = CompletionReserve.estimate( + schemaJSON: schemaJSON, + tokenizer: context.tokenizer) + let completionReserve = Swift.max(structuralReserve * 3, maxTokens / 4) + let hardReserve = structuralReserve * 8 + + let (textStream, textContinuation) = AsyncStream.makeStream() + async let forwarder: Void = { + for await text in textStream { + await Self.emit( + text: text, + entryID: entryID, + destination: .response, + into: channel) + } + }() + + var incomplete = false + var generatedTokenCount: Int? + do { + generatedTokenCount = try GuidedGenerationLoop.run( + input: input, + context: context, + constraint: constraint, + maxTokens: maxTokens, + vocabSize: Int(xgTokenizer.vocabSize), + completionReserve: completionReserve, + hardReserve: hardReserve, + closingBias: bias.closing, + whitespaceBias: bias.whitespace, + whitespaceTokenIDs: bias.whitespaceTokenIDs + ) { text in + textContinuation.yield(text) + GuidedGenerationDiagnosticSink.current?.recordEmit() + return !Task.isCancelled + } + } catch GuidedGenerationError.incompleteOutput { + incomplete = true + } + + let cancellationError: Error? + do { + try Task.checkCancellation() + cancellationError = nil + } catch { + cancellationError = error + } + textContinuation.finish() + await forwarder + if let cancellationError { + throw cancellationError + } + + if let generatedTokenCount { + await Self.emitUsage( + input: .init( + totalTokenCount: input.text.tokens.size, + cachedTokenCount: 0), + output: .init( + totalTokenCount: generatedTokenCount, + reasoningTokenCount: 0), + entryID: entryID, + into: channel) + } + if incomplete { + await Self.emitMetadata( + ["incompleteOutput": true], entryID: entryID, into: channel) + } + } + /// Unconstrained text generation. Used on the no-tools/no-schema /// path when the model has no reasoning config to route through. private func runUnconstrained( input: LMInput, requestedMaxTokens: Int?, requestedTemperature: Double?, - samplingMode: MLXSamplingMode?, + samplingConfiguration: MLXSamplingConfiguration?, entryID: String, context: ModelContext, channel: LanguageModelExecutorGenerationChannel @@ -1430,7 +1695,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { let params = Self.makeParameters( maxTokens: requestedMaxTokens ?? Self.defaultMaxTokens, requestedTemperature: requestedTemperature, - samplingMode: samplingMode + samplingConfiguration: samplingConfiguration ) for await generation in try generate( @@ -1441,31 +1706,19 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { try Task.checkCancellation() switch generation { case .chunk(let text): - await channel.send( - .response( - entryID: entryID, - action: .appendText(text, tokenCount: 1) - )) + await Self.emit( + text: text, entryID: entryID, destination: .response, into: channel) case .info(let info): // MLX-LM emits one .info event at end-of-generation with // authoritative scalar token counts (`promptTokenCount` // is the prompt; `generationTokenCount` is the // model-generated completion -- see Evaluate.swift's // `GenerateCompletionInfo` definition). - await channel.send( - .response( - entryID: entryID, - action: .updateUsage( - input: .init( - totalTokenCount: info.promptTokenCount, - cachedTokenCount: 0 - ), - output: .init( - totalTokenCount: info.generationTokenCount, - reasoningTokenCount: 0 - ) - ) - )) + await Self.emitUsage( + input: .init(totalTokenCount: info.promptTokenCount, cachedTokenCount: 0), + output: .init( + totalTokenCount: info.generationTokenCount, reasoningTokenCount: 0), + entryID: entryID, into: channel) case .toolCall(_): break } @@ -1479,7 +1732,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { fallbackInput: LMInput, requestedMaxTokens: Int?, requestedTemperature: Double?, - samplingMode: MLXSamplingMode?, + samplingConfiguration: MLXSamplingConfiguration?, responseEntryID: String, reasoningEntryID: String, context: ModelContext, @@ -1492,7 +1745,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { primedInside: reasoning.primedInside, requestedMaxTokens: requestedMaxTokens, requestedTemperature: requestedTemperature, - samplingMode: samplingMode, + samplingConfiguration: samplingConfiguration, responseEntryID: responseEntryID, reasoningEntryID: reasoningEntryID, context: context, @@ -1502,7 +1755,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { input: fallbackInput, requestedMaxTokens: requestedMaxTokens, requestedTemperature: requestedTemperature, - samplingMode: samplingMode, + samplingConfiguration: samplingConfiguration, entryID: responseEntryID, context: context, channel: channel) @@ -1523,7 +1776,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { primedInside: Bool, requestedMaxTokens: Int?, requestedTemperature: Double?, - samplingMode: MLXSamplingMode?, + samplingConfiguration: MLXSamplingConfiguration?, responseEntryID: String, reasoningEntryID: String, context: ModelContext, @@ -1532,7 +1785,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { let params = Self.makeParameters( maxTokens: requestedMaxTokens ?? Self.defaultMaxTokens, requestedTemperature: requestedTemperature, - samplingMode: samplingMode + samplingConfiguration: samplingConfiguration ) var emitter = ReasoningEventEmitter( @@ -1582,10 +1835,8 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // empty or partial answer for the model's chosen response — mirrors // the guided path's `incompleteOutput` convention. if emitter.isInsideReasoning { - await channel.send( - .response( - entryID: responseEntryID, - action: .updateMetadata(["incompleteOutput": true]))) + await Self.emitMetadata( + ["incompleteOutput": true], entryID: responseEntryID, into: channel) } if let info = completionInfo { @@ -1593,21 +1844,12 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // `.updateUsage` (the framework's aggregator replaces wholesale, // so we must not also rely on per-delta auto-summing). The // reasoning count is clamped to never exceed the total. - await channel.send( - .response( - entryID: responseEntryID, - action: .updateUsage( - input: .init( - totalTokenCount: info.promptTokenCount, - cachedTokenCount: 0 - ), - output: .init( - totalTokenCount: info.generationTokenCount, - reasoningTokenCount: min( - reasoningTokenCount, info.generationTokenCount) - ) - ) - )) + await Self.emitUsage( + input: .init(totalTokenCount: info.promptTokenCount, cachedTokenCount: 0), + output: .init( + totalTokenCount: info.generationTokenCount, + reasoningTokenCount: min(reasoningTokenCount, info.generationTokenCount)), + entryID: responseEntryID, into: channel) } } @@ -1620,15 +1862,11 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { ) async { switch segment { case .reasoning(let text): - await channel.send( - .reasoning( - entryID: reasoningEntryID, - action: .appendText(text, tokenCount: 1))) + await Self.emit( + text: text, entryID: reasoningEntryID, destination: .reasoning, into: channel) case .response(let text): - await channel.send( - .response( - entryID: responseEntryID, - action: .appendText(text, tokenCount: 1))) + await Self.emit( + text: text, entryID: responseEntryID, destination: .response, into: channel) } } @@ -1680,6 +1918,28 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// Decodes the rendered prompt's tail and asks whether it ends inside an /// open reasoning block (some model families prefill the opening /// delimiter). + /// Build the Phase-2 continuation input: the tool-aware prompt with the + /// completed reasoning token IDs appended along the sequence axis. + /// + /// The prompt tokens keep whatever rank the model's processor produced + /// ([N] from LLM processors, [1, N] from VLM processors — VLM `prepare` + /// requires the batched form), and processed image/video content is + /// carried through so a VLM's Phase-2 prefill still sees its pixels. + static func continuationInput( + from input: LMInput, appending tokenIDs: [Int] + ) -> LMInput { + let promptTokens = input.text.tokens + var appended = MLXArray(tokenIDs.map { Int32($0) }) + .asType(promptTokens.dtype) + if promptTokens.ndim == 2 { + appended = appended[.newAxis, 0...] + } + return LMInput( + text: .init(tokens: concatenated([promptTokens, appended], axis: -1)), + image: input.image, + video: input.video) + } + private static func reasoningPrimedInside( input: LMInput, config: ReasoningConfig, tokenizer: any Tokenizer ) -> Bool { @@ -1708,7 +1968,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { primedInside: Bool, maxTokens: Int, requestedTemperature: Double?, - samplingMode: MLXSamplingMode?, + samplingConfiguration: MLXSamplingConfiguration?, reasoningEntryID: String, responseEntryID: String, context: ModelContext, @@ -1717,7 +1977,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { let params = Self.makeParameters( maxTokens: maxTokens, requestedTemperature: requestedTemperature, - samplingMode: samplingMode + samplingConfiguration: samplingConfiguration ) var collector = ReasoningTokenCollector( config: config, primedInside: primedInside, tokenizer: context.tokenizer @@ -1736,6 +1996,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { reasoningEntryID: reasoningEntryID, channel: channel) } if collector.shouldStopAfterReasoning { + GuidedGenerationDiagnosticSink.current?.recordToolReasoningClose() closed = true break } @@ -1763,8 +2024,8 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { return (collector.reasoningTokenIDs, closed) } - /// Parses a tool-calling envelope JSON object and emits the - /// appropriate channel event. + /// Parses a required-mode tool-calling envelope JSON object and emits + /// its developer tool call. /// /// The output buffer is expected to be a JSON object matching the /// shape `{"name": , "arguments": }`. Grammars from @@ -1772,28 +2033,16 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// shape directly (bare JSON) or that shape wrapped in Qwen's /// `\n...\n` special-token delimiters -- /// `unwrapToolCallMarkers` below strips the wrapper if present. The - /// best-effort fallback only exists so that unexpected upstream - /// changes don't silently swallow output. + /// guided path emits a single `.toolCallDelta` with the arguments JSON + /// and a freshly minted toolCallID. /// - /// - If `name` is the synthetic final-answer tool: - /// - With no developer response schema: unwrap `arguments.response` - /// into a `.textDelta` event. - /// - With a developer response schema: re-serialize `arguments` - /// back to JSON text and emit as a single `.textDelta`. The - /// session's normal response-parsing path will decode the JSON - /// through the developer's `GenerationSchema`. - /// - If `name` is any real tool: emit a single `.toolCallDelta` - /// with the arguments JSON and a freshly minted toolCallID. - /// - /// `entryID` and `toolCallsEntryID` must be distinct: SKILL.md requires - /// `.response` and `.toolCalls` to live in separate transcript entries. - private func emitToolCallingEvent( + /// Required mode never degrades malformed or partial output into a + /// response event. + private func emitRequiredToolCallEvent( outputBuffer: String, - userResponseSchema: GenerationSchema?, - entryID: String, toolCallsEntryID: String, channel: LanguageModelExecutorGenerationChannel - ) async throws { + ) async { let unwrapped = Self.unwrapToolCallMarkers(outputBuffer) let data = Data(unwrapped.utf8) guard @@ -1801,52 +2050,27 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { as? [String: Any], let name = obj["name"] as? String else { - // Malformed output. The grammar should have prevented this; - // emit the raw buffer as text so failures surface loudly. - await channel.send( - .response( - entryID: entryID, - action: .appendText(outputBuffer, tokenCount: 1) - )) + GuidedGenerationDiagnosticSink.current?.recordParse( + parsedAsToolCall: false, parsedName: nil) return } - if name == FinalAnswerTool.toolName { - let text: String - if userResponseSchema == nil { - let args = obj["arguments"] as? [String: Any] - text = (args?["response"] as? String) ?? "" - } else if let args = obj["arguments"], - let argsData = try? JSONSerialization.data(withJSONObject: args), - let argsStr = String(data: argsData, encoding: .utf8) - { - text = argsStr - } else { - text = "" - } - await channel.send( - .response( - entryID: entryID, - action: .appendText(text, tokenCount: 1) - )) - } else { - guard - let args = obj["arguments"], - let argsData = try? JSONSerialization.data(withJSONObject: args), - let argsStr = String(data: argsData, encoding: .utf8) - else { - return - } - await channel.send( - .toolCalls( - entryID: toolCallsEntryID, - action: .toolCall( - id: UUID().uuidString, - name: name, - action: .appendArguments(argsStr, tokenCount: 1) - ) - )) + GuidedGenerationDiagnosticSink.current?.recordParse( + parsedAsToolCall: true, parsedName: name) + + guard + let arguments = obj["arguments"], + let argumentsData = try? JSONSerialization.data(withJSONObject: arguments), + let argumentsJSON = String(data: argumentsData, encoding: .utf8) + else { + return } + await Self.emitToolCall( + id: UUID().uuidString, + name: name, + arguments: argumentsJSON, + entryID: toolCallsEntryID, + into: channel) } /// Strips Qwen-style `\n...\n` wrapper markers diff --git a/Libraries/MLXFoundationModels/SamplingModeMapper.swift b/Libraries/MLXFoundationModels/SamplingModeMapper.swift index e55387433..9ba8cd90b 100644 --- a/Libraries/MLXFoundationModels/SamplingModeMapper.swift +++ b/Libraries/MLXFoundationModels/SamplingModeMapper.swift @@ -6,9 +6,8 @@ import MLXLMCommon /// Sampling-strategy selection for the adapter, resolved to the /// `GenerateParameters` fields MLX's sampler consumes. /// -/// The adapter translates the FoundationModels `GenerationOptions.SamplingMode` -/// into this enum at dispatch (dropping the best-effort `seed`, which MLX's -/// samplers cannot honor) and applies the result to `GenerateParameters` via +/// The adapter keeps the Foundation Models seed alongside this mode in +/// ``MLXSamplingConfiguration`` and applies mode policy through /// ``resolveSamplingParameters(mode:clampedTemperature:)``. public enum MLXSamplingMode: Sendable, Equatable { /// Deterministic decoding — always pick the most likely token. @@ -24,6 +23,13 @@ public enum MLXSamplingMode: Sendable, Equatable { case nucleus(Double) } +/// The adapter-local sampling request. Keeping mode and seed together prevents +/// call sites from applying one while accidentally dropping the other. +struct MLXSamplingConfiguration: Sendable, Equatable { + let mode: MLXSamplingMode + let seed: UInt64? +} + /// The sampling fields a resolved ``MLXSamplingMode`` contributes to /// `GenerateParameters`. A `nil` field means "leave the provider default in /// place." The resolver never emits a concrete temperature default, because that diff --git a/Libraries/MLXFoundationModels/ToolCalling/AllowedToolOutputRouter.swift b/Libraries/MLXFoundationModels/ToolCalling/AllowedToolOutputRouter.swift new file mode 100644 index 000000000..7fe185a3e --- /dev/null +++ b/Libraries/MLXFoundationModels/ToolCalling/AllowedToolOutputRouter.swift @@ -0,0 +1,89 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration +#if canImport(FoundationModels, _version: 2) + +import MLXLMCommon + +struct AllowedToolOutputRouter { + enum Event: Sendable, Equatable { + case reasoning(String) + case response(String) + case toolCall(MLXLMCommon.ToolCall) + } + + private var reasoningEmitter: ReasoningEventEmitter? + private let toolProcessor: ToolCallProcessor + private let allowedToolNames: Set + + init( + format: ToolCallFormat, + tools: [[String: any Sendable]], + reasoning: (config: ReasoningConfig, primedInside: Bool)? = nil + ) { + self.toolProcessor = ToolCallProcessor(format: format, tools: tools) + self.allowedToolNames = Set( + tools.compactMap { tool in + (tool["function"] as? [String: any Sendable])?["name"] as? String + }) + self.reasoningEmitter = reasoning.map { + ReasoningEventEmitter(config: $0.config, primedInside: $0.primedInside) + } + } + + var isInsideReasoning: Bool { + reasoningEmitter?.isInsideReasoning ?? false + } + + mutating func process(_ chunk: String) -> [Event] { + guard var emitter = reasoningEmitter else { + return processResponse(chunk) + } + + let segments = emitter.process(chunk) + reasoningEmitter = emitter + var events: [Event] = [] + for segment in segments { + switch segment { + case .reasoning(let text): + events.append(.reasoning(text)) + case .response(let text): + events.append(contentsOf: processResponse(text)) + } + } + return events + } + + mutating func finish() -> [Event] { + var events: [Event] = [] + if var emitter = reasoningEmitter { + for segment in emitter.finalize() { + switch segment { + case .reasoning(let text): events.append(.reasoning(text)) + case .response(let text): events.append(contentsOf: processResponse(text)) + } + } + reasoningEmitter = emitter + } + events.append(contentsOf: route(toolProcessor.processEOSOutputs())) + return events + } + + private mutating func processResponse(_ text: String) -> [Event] { + route(toolProcessor.processChunkOutputs(text)) + } + + private func route(_ outputs: [ToolCallProcessor.Output]) -> [Event] { + outputs.compactMap { output in + switch output { + case .response(let text): + .response(text) + case .toolCall(let call): + allowedToolNames.contains(call.function.name) ? .toolCall(call) : nil + } + } + } +} + +#endif +#endif diff --git a/Libraries/MLXFoundationModels/ToolCalling/FinalAnswerTool.swift b/Libraries/MLXFoundationModels/ToolCalling/FinalAnswerTool.swift deleted file mode 100644 index 3111d7e61..000000000 --- a/Libraries/MLXFoundationModels/ToolCalling/FinalAnswerTool.swift +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#if FoundationModelsIntegration -#if canImport(FoundationModels, _version: 2) - -import Foundation -import FoundationModels - -/// Synthetic tool used by MLX's tool-calling path to encode the model's -/// free-text response as a structured tool call. -/// -/// MLX constrains tool-calling generation to a JSON schema shaped as -/// `{oneOf: [{name: "T_i", arguments: }, …]}`. The -/// developer's real tools are the `T_1…T_N`; this synthetic tool is the -/// extra `T_{N+1}` whose arguments carry the text (or structured response) -/// the model wants to deliver directly to the user. -/// -/// When the model picks this tool at generation time, the executor does not -/// emit a `toolCallDelta` for it -- instead it extracts the `arguments` -/// payload and re-emits it as `textDelta` events, so consumers of the -/// channel see text in the same shape they would for a tools-free response. -@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) -enum FinalAnswerTool { - - /// Reserved tool name. Developers must not register a real tool with - /// this name; if they do, resolution silently keeps the synthetic - /// tool (no auto-renaming). - static let toolName = "mlx_final_answer" - - /// Human-readable description shown to the model alongside the real - /// tools' descriptions. - static let toolDescription = """ - Call this tool to respond directly to the user in natural language. \ - Use it when no other tool is needed, or once information gathered \ - from prior tool calls is sufficient to answer the user's request. - """ - - /// Wrapper schema used when the request has no developer-supplied - /// response schema. The tool's single argument `response` carries the - /// free-text response; the executor unwraps it into text deltas. - @Generable - struct StringResponse { - @Guide(description: "The natural-language response to return to the user.") - var response: String - } - - /// Builds the `Transcript.ToolDefinition` the model should see in its - /// prompt, alongside the developer's real tools. - /// - /// - Parameter responseSchema: The developer-provided response schema - /// for the current request, if any. - /// - `nil`: the synthetic tool uses the `StringResponse` wrapper, so - /// the tool's arguments are `{"response": ""}`. - /// - non-`nil`: the developer's schema is used verbatim as the - /// synthetic tool's `arguments` schema. Consumers then decode the - /// tool's arguments JSON through their own `GenerationSchema`. - static func makeToolDefinition( - responseSchema: GenerationSchema? - ) -> Transcript.ToolDefinition { - Transcript.ToolDefinition( - name: toolName, - description: toolDescription, - parameters: parameterSchema(for: responseSchema) - ) - } - - /// Selects the schema used for the synthetic tool's `arguments`. - static func parameterSchema( - for responseSchema: GenerationSchema? - ) -> GenerationSchema { - responseSchema ?? StringResponse.generationSchema - } -} - -#endif // canImport(FoundationModels) -#endif // FoundationModelsIntegration diff --git a/Libraries/MLXFoundationModels/ToolCalling/ToolCallingModeResolution.swift b/Libraries/MLXFoundationModels/ToolCalling/ToolCallingModeResolution.swift new file mode 100644 index 000000000..af9b3b09b --- /dev/null +++ b/Libraries/MLXFoundationModels/ToolCalling/ToolCallingModeResolution.swift @@ -0,0 +1,49 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration +#if canImport(FoundationModels, _version: 2) + +import FoundationModels + +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +enum ToolCallingModeResolution { + enum Error: Swift.Error, Equatable { + case requiredToolsMissing + } + + static func resolve( + _ mode: GenerationOptions.ToolCallingMode? + ) -> GenerationOptions.ToolCallingMode { + mode ?? .allowed + } + + static func usesAllowedBehavior( + _ mode: GenerationOptions.ToolCallingMode + ) -> Bool { + switch mode.kind { + case .allowed: + return true + case .required, .disallowed: + return false + @unknown default: + return true + } + } + + static func enabledToolDefinitions( + for mode: GenerationOptions.ToolCallingMode, + from definitions: [Transcript.ToolDefinition] + ) throws -> [Transcript.ToolDefinition] { + if usesAllowedBehavior(mode) { + return definitions + } + if mode.kind == .disallowed { + return [] + } + guard !definitions.isEmpty else { throw Error.requiredToolsMissing } + return definitions + } +} + +#endif +#endif diff --git a/Libraries/MLXFoundationModels/TranscriptConverter.swift b/Libraries/MLXFoundationModels/TranscriptConverter.swift index c18e3addf..805d7a6be 100644 --- a/Libraries/MLXFoundationModels/TranscriptConverter.swift +++ b/Libraries/MLXFoundationModels/TranscriptConverter.swift @@ -3,6 +3,7 @@ #if FoundationModelsIntegration #if canImport(FoundationModels, _version: 2) +import Foundation import FoundationModels import MLXLMCommon import os.log @@ -66,14 +67,76 @@ struct TranscriptConverter { logger.debug("Skipping reasoning entry (not replayed into chat history)") return nil + case .toolCalls(let toolCalls): + // Replay prior tool calls as an assistant message carrying the + // structured calls. The model's tool-aware chat template renders + // these into its native tool-call channel; DefaultMessageGenerator + // serializes each id/name/arguments (see ToolCallIdTests). Without + // this, a continuation round would re-issue the same call. + let calls = toolCalls.map { call -> MLXLMCommon.ToolCall in + let argumentsData = Data(call.arguments.jsonString.utf8) + let arguments: [String: JSONValue] + if let decoded = try? JSONDecoder().decode( + [String: JSONValue].self, from: argumentsData) + { + arguments = decoded + } else { + logger.warning( + "Failed to decode arguments for tool: \(call.toolName, privacy: .public)" + ) + arguments = [:] + } + return MLXLMCommon.ToolCall( + function: .init(name: call.toolName, arguments: arguments), + id: call.id) + } + guard !calls.isEmpty else { + logger.warning("Skipping toolCalls entry with no calls") + return nil + } + return Chat.Message.assistant("", toolCalls: calls) + + case .toolOutput(let output): + // Replay the tool result as a `tool` message correlated to its + // originating call by id. Text remains verbatim; structured + // GeneratedContent is serialized as JSON so the native chat + // template can expose it to the continuation model turn. + let content = extractToolOutputContent(from: output.segments) + return Chat.Message.tool(content, id: output.id) + default: - // Skip unsupported entry types (toolCalls, toolOutput, etc.) + // Skip unsupported entry types. Explicit `return nil` is a + // tripwire: a newly added SDK entry type surfaces here for review + // rather than being silently coerced into the wrong role. logger.debug("Skipping unsupported entry type") return nil } } } + /// Extracts supported tool-output content in transcript segment order. + /// + /// Foundation Models lowers `String` outputs to `.text` and + /// `GeneratedContent`/`@Generable` outputs to `.structure`. MLX chat + /// templates accept tool results as strings, so structured values retain + /// their JSON representation. Attachments and custom segments are deferred + /// until their media and prompt-representation contracts are implemented. + private static func extractToolOutputContent( + from segments: [Transcript.Segment] + ) -> String { + segments.compactMap { segment -> String? in + switch segment { + case .text(let textSegment): + return textSegment.content + case .structure(let structuredSegment): + return structuredSegment.content.jsonString + default: + logger.debug("Skipping unsupported tool-output segment") + return nil + } + }.joined(separator: "\n") + } + /// Extracts text content from transcript segments. /// /// Concatenates all text segments with newlines. diff --git a/Libraries/MLXGuidedGeneration/GuidedGenerationDiagnosticSink.swift b/Libraries/MLXGuidedGeneration/GuidedGenerationDiagnosticSink.swift new file mode 100644 index 000000000..6801cef0d --- /dev/null +++ b/Libraries/MLXGuidedGeneration/GuidedGenerationDiagnosticSink.swift @@ -0,0 +1,106 @@ +// Copyright © 2026 Apple Inc. + +import Foundation + +/// Test-only diagnostic capture for the tool-call grammar/tokenizer boundary. +/// +/// Off by default: production never binds `current`, so every recording site is +/// a nil-guarded no-op and there is no behavior change. A test binds an instance +/// with `GuidedGenerationDiagnosticSink.$current.withValue(_:)` around a real +/// generation call, drives the executor, then reads the accumulated signals. +/// +/// Reference type so writes made deep in the synchronous generation loop are +/// visible to the test afterward. `@unchecked Sendable` because it is mutated +/// only within the single generation task, never concurrently. +public final class GuidedGenerationDiagnosticSink: @unchecked Sendable { + + /// Task-local injection point. `nil` in production. + @TaskLocal public static var current: GuidedGenerationDiagnosticSink? + + /// Model-sampled token IDs, in generation order. Grammar-forced tokens are + /// recorded separately in `fastForwardTokenIDs`. + public private(set) var sampledTokenIDs: [Int] = [] + + /// Grammar-forced fast-forward token IDs, in generation order. + public private(set) var fastForwardTokenIDs: [Int] = [] + + /// Whether the grammar reached a stop state (vs. exhausting the budget). + public private(set) var grammarTerminated = false + + /// Total tokens generated (sampled + fast-forward), as the loop counts them. + public private(set) var generatedTokenCount = 0 + + /// The exact buffer handed to the tool-call parser. + public private(set) var finalBuffer: String? + + /// True when generation hit the token budget before the grammar stopped. + public private(set) var incompleteOutput = false + + /// Whether `finalBuffer` parsed as a valid tool call (an object with a + /// `name` field). `nil` until the parser runs. + public private(set) var parsedAsToolCall: Bool? + + /// The parsed tool name, when the buffer parsed as a tool call. + public private(set) var parsedName: String? + + /// Number of synchronous executor-side guided emit boundaries observed. + public private(set) var emitCount = 0 + + /// Number of required-tool reasoning close boundaries observed. + public private(set) var toolReasoningCloseCount = 0 + + private let cancelAfterEmitCount: Int? + private let cancelOnToolReasoningClose: Bool + + public init( + cancelAfterEmitCount: Int? = nil, + cancelOnToolReasoningClose: Bool = false + ) { + self.cancelAfterEmitCount = cancelAfterEmitCount + self.cancelOnToolReasoningClose = cancelOnToolReasoningClose + } + + // MARK: Recording (called from generation; no-ops when the sink is unbound) + + public func recordSampledToken(_ id: Int) { + sampledTokenIDs.append(id) + } + + public func recordFastForwardToken(_ id: Int) { + fastForwardTokenIDs.append(id) + } + + public func recordTermination(grammarTerminated: Bool, generatedTokenCount: Int) { + self.grammarTerminated = grammarTerminated + self.generatedTokenCount = generatedTokenCount + } + + public func recordBuffer(_ buffer: String, incompleteOutput: Bool) { + self.finalBuffer = buffer + self.incompleteOutput = incompleteOutput + } + + public func recordParse(parsedAsToolCall: Bool, parsedName: String?) { + self.parsedAsToolCall = parsedAsToolCall + self.parsedName = parsedName + } + + /// Records a guided loop's synchronous emit boundary and optionally + /// cancels its calling task. Cancellation is opt-in and used only by tests + /// that bind this sink through ``current``. + public func recordEmit() { + emitCount += 1 + if emitCount == cancelAfterEmitCount { + withUnsafeCurrentTask { $0?.cancel() } + } + } + + /// Records the boundary where Phase 1 has consumed the reasoning close + /// marker and optionally cancels its calling task. + public func recordToolReasoningClose() { + toolReasoningCloseCount += 1 + if cancelOnToolReasoningClose { + withUnsafeCurrentTask { $0?.cancel() } + } + } +} diff --git a/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift b/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift index 7d3b158d5..f530b6564 100644 --- a/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift +++ b/Libraries/MLXGuidedGeneration/GuidedGenerationLoop.swift @@ -99,6 +99,7 @@ public enum GuidedGenerationLoop { emit: (String) -> Bool ) throws -> Int { let model = context.model + let diagnosticSink = GuidedGenerationDiagnosticSink.current var cache = model.newCache(parameters: nil) var modelState: LMOutput.State? @@ -288,6 +289,7 @@ public enum GuidedGenerationLoop { // Commit to grammar let commitResult = try constraint.commitToken(Int32(token)) + diagnosticSink?.recordSampledToken(tokenId) // Yield the sampled token detokenizer.append(token: tokenId) if let text = detokenizer.next() { @@ -337,6 +339,7 @@ public enum GuidedGenerationLoop { shouldStopAfterFF = true break } + diagnosticSink?.recordFastForwardToken(Int(ffToken)) detokenizer.append(token: Int(ffToken)) if let text = detokenizer.next() { accumulatedText += text @@ -429,6 +432,9 @@ public enum GuidedGenerationLoop { logger.warning("[GuidedGen] xgrammar logs:\n\(logs)") } + diagnosticSink?.recordTermination( + grammarTerminated: grammarStopped, generatedTokenCount: tokenCount) + // If we exhausted maxTokens without the grammar reaching a stop state, // the output is structurally incomplete (e.g., truncated JSON). if !grammarStopped && tokenCount >= maxTokens { diff --git a/Libraries/MLXLLM/LLMModel.swift b/Libraries/MLXLLM/LLMModel.swift index 87aeea0a5..2b608818a 100644 --- a/Libraries/MLXLLM/LLMModel.swift +++ b/Libraries/MLXLLM/LLMModel.swift @@ -1,5 +1,6 @@ // Copyright © 2024 Apple Inc. +import Foundation import MLX import MLXLMCommon @@ -40,11 +41,15 @@ extension LLMModel { // prefill cannot be interrupted, so apps cannot stop GPU submissions // in time when entering the background. See ml-explore/mlx-swift-examples#230. try Task.checkCancellation() - let input = y[.newAxis, .."] ) + static public let hunyuan_mt_7b_4bit = ModelConfiguration( + id: "mlx-community/Hunyuan-MT-7B-4bit", + defaultPrompt: "Translate the following text into Chinese: Hello, how are you?" + ) + + static public let hunyuan_mt_7b_8bit = ModelConfiguration( + id: "mlx-community/Hunyuan-MT-7B-8bit", + defaultPrompt: "Translate the following text into Chinese: Hello, how are you?" + ) + + static public let hy_mt2_7b_4bit = ModelConfiguration( + id: "mlx-community/Hy-MT2-7B-4bit", + defaultPrompt: "Translate the following text into Chinese: Hello, how are you?" + ) + + static public let hy_mt2_7b_8bit = ModelConfiguration( + id: "mlx-community/Hy-MT2-7B-8bit", + defaultPrompt: "Translate the following text into Chinese: Hello, how are you?" + ) + static public let qwen205b4bit = ModelConfiguration( id: "mlx-community/Qwen1.5-0.5B-Chat-4bit", defaultPrompt: "why is the sky blue?", @@ -423,6 +445,10 @@ public class LLMRegistry: AbstractModelRegistry, @unchecked Sendable { gemma3n_E2B_it_lm_4bit, gemma4_e4b_it_4bit, gemma4_e2b_it_4bit, + hunyuan_mt_7b_4bit, + hunyuan_mt_7b_8bit, + hy_mt2_7b_4bit, + hy_mt2_7b_8bit, granite3_3_2b_4bit, granite_4_0_h_tiny_4bit_dwq, llama3_1_8B_4bit, @@ -572,8 +598,10 @@ public final class LLMModelFactory: GenericModelFactory { configurationURL.lastPathComponent, configuration.name, error) } - // Load generation_config.json for EOS token IDs and recommended sampling params. - var eosTokenIds = Set(baseConfig.eosTokenIds?.values ?? []) + // Load EOS token IDs from config.json (nested text configs included), with + // optional override from generation_config.json — which also carries the + // recommended sampling params we surface via ModelContext. + var eosTokenIds = baseConfig.effectiveEOSTokenIds let generationConfigURL = modelDirectory.appending(component: "generation_config.json") let parsedGenerationConfig: GenerationConfigFile? = if let generationData = try? Data(contentsOf: generationConfigURL) { diff --git a/Libraries/MLXLLM/Models/DeepseekV2.swift b/Libraries/MLXLLM/Models/DeepseekV2.swift new file mode 100644 index 000000000..41200b7e0 --- /dev/null +++ b/Libraries/MLXLLM/Models/DeepseekV2.swift @@ -0,0 +1,406 @@ +// port of https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/deepseek_v2.py + +import Foundation +import MLX +import MLXLMCommon +import MLXNN + +public struct DeepseekV2Configuration: Codable, Sendable { + var vocabSize: Int = 102400 + var hiddenSize: Int = 4096 + var intermediateSize: Int = 11008 + var moeIntermediateSize: Int = 1407 + var numHiddenLayers: Int = 30 + var numAttentionHeads: Int = 32 + var numKeyValueHeads: Int = 32 + var nSharedExperts: Int? + var nRoutedExperts: Int? + var routedScalingFactor: Float = 1.0 + var kvLoraRank: Int = 512 + // Optional: DeepSeek-V2-Lite sets `q_lora_rank: null` (direct q_proj, no + // low-rank q compression); full DeepSeek-V2 sets it to 1536. nil ⇒ q_proj. + var qLoraRank: Int? + var qkRopeHeadDim: Int = 64 + var vHeadDim: Int = 128 + var qkNopeHeadDim: Int = 128 + var topkMethod: String = "greedy" + var nGroup: Int? + var topkGroup: Int? + var numExpertsPerTok: Int? + var moeLayerFreq: Int = 1 + var firstKDenseReplace: Int = 0 + var maxPositionEmbeddings: Int = 2048 + var rmsNormEps: Float = 1e-6 + var ropeTheta: Float = 10000.0 + var ropeScaling: [String: StringOrNumber]? + var attentionBias: Bool = false + + enum CodingKeys: String, CodingKey { + case vocabSize = "vocab_size" + case hiddenSize = "hidden_size" + case intermediateSize = "intermediate_size" + case moeIntermediateSize = "moe_intermediate_size" + case numHiddenLayers = "num_hidden_layers" + case numAttentionHeads = "num_attention_heads" + case numKeyValueHeads = "num_key_value_heads" + case nSharedExperts = "n_shared_experts" + case nRoutedExperts = "n_routed_experts" + case routedScalingFactor = "routed_scaling_factor" + case kvLoraRank = "kv_lora_rank" + case qLoraRank = "q_lora_rank" + case qkRopeHeadDim = "qk_rope_head_dim" + case vHeadDim = "v_head_dim" + case qkNopeHeadDim = "qk_nope_head_dim" + case topkMethod = "topk_method" + case nGroup = "n_group" + case topkGroup = "topk_group" + case numExpertsPerTok = "num_experts_per_tok" + case moeLayerFreq = "moe_layer_freq" + case firstKDenseReplace = "first_k_dense_replace" + case maxPositionEmbeddings = "max_position_embeddings" + case rmsNormEps = "rms_norm_eps" + case ropeTheta = "rope_theta" + case ropeScaling = "rope_scaling" + case attentionBias = "attention_bias" + } + + public init(from decoder: Decoder) throws { + let c = try decoder.container(keyedBy: CodingKeys.self) + self.vocabSize = try c.decodeIfPresent(Int.self, forKey: .vocabSize) ?? 102400 + self.hiddenSize = try c.decodeIfPresent(Int.self, forKey: .hiddenSize) ?? 4096 + self.intermediateSize = try c.decodeIfPresent(Int.self, forKey: .intermediateSize) ?? 11008 + self.moeIntermediateSize = + try c.decodeIfPresent(Int.self, forKey: .moeIntermediateSize) ?? 1407 + self.numHiddenLayers = try c.decodeIfPresent(Int.self, forKey: .numHiddenLayers) ?? 30 + self.numAttentionHeads = try c.decodeIfPresent(Int.self, forKey: .numAttentionHeads) ?? 32 + self.numKeyValueHeads = try c.decodeIfPresent(Int.self, forKey: .numKeyValueHeads) ?? 32 + self.nSharedExperts = try c.decodeIfPresent(Int.self, forKey: .nSharedExperts) + self.nRoutedExperts = try c.decodeIfPresent(Int.self, forKey: .nRoutedExperts) + self.routedScalingFactor = + try c.decodeIfPresent(Float.self, forKey: .routedScalingFactor) ?? 1.0 + self.kvLoraRank = try c.decodeIfPresent(Int.self, forKey: .kvLoraRank) ?? 512 + self.qLoraRank = try c.decodeIfPresent(Int.self, forKey: .qLoraRank) + self.qkRopeHeadDim = try c.decodeIfPresent(Int.self, forKey: .qkRopeHeadDim) ?? 64 + self.vHeadDim = try c.decodeIfPresent(Int.self, forKey: .vHeadDim) ?? 128 + self.qkNopeHeadDim = try c.decodeIfPresent(Int.self, forKey: .qkNopeHeadDim) ?? 128 + self.topkMethod = try c.decodeIfPresent(String.self, forKey: .topkMethod) ?? "greedy" + self.nGroup = try c.decodeIfPresent(Int.self, forKey: .nGroup) + self.topkGroup = try c.decodeIfPresent(Int.self, forKey: .topkGroup) + self.numExpertsPerTok = try c.decodeIfPresent(Int.self, forKey: .numExpertsPerTok) + self.moeLayerFreq = try c.decodeIfPresent(Int.self, forKey: .moeLayerFreq) ?? 1 + self.firstKDenseReplace = try c.decodeIfPresent(Int.self, forKey: .firstKDenseReplace) ?? 0 + self.maxPositionEmbeddings = + try c.decodeIfPresent(Int.self, forKey: .maxPositionEmbeddings) ?? 2048 + self.rmsNormEps = try c.decodeIfPresent(Float.self, forKey: .rmsNormEps) ?? 1e-6 + self.ropeTheta = try c.decodeIfPresent(Float.self, forKey: .ropeTheta) ?? 10000.0 + self.ropeScaling = try c.decodeIfPresent( + [String: StringOrNumber].self, forKey: .ropeScaling) + self.attentionBias = try c.decodeIfPresent(Bool.self, forKey: .attentionBias) ?? false + } +} + +class DeepseekV2Attention: Module { + let config: DeepseekV2Configuration + let numHeads: Int + let qLoraRank: Int? + let qkRopeHeadDim: Int + let kvLoraRank: Int + let vHeadDim: Int + let qkNopeHeadDim: Int + let qHeadDim: Int + var scale: Float + + let rope: RoPELayer + @ModuleInfo(key: "q_proj") var qProj: Linear? + @ModuleInfo(key: "q_a_proj") var qAProj: Linear? + @ModuleInfo(key: "q_a_layernorm") var qALayerNorm: RMSNorm? + @ModuleInfo(key: "q_b_proj") var qBProj: Linear? + @ModuleInfo(key: "o_proj") var oProj: Linear + @ModuleInfo(key: "kv_a_proj_with_mqa") var kvAProjWithMqa: Linear + @ModuleInfo(key: "kv_a_layernorm") var kvALayerNorm: RMSNorm + @ModuleInfo(key: "kv_b_proj") var kvBProj: Linear + + init(config: DeepseekV2Configuration) { + self.config = config + self.numHeads = config.numAttentionHeads + self.qLoraRank = (config.qLoraRank ?? 0) > 0 ? config.qLoraRank : nil + self.qkRopeHeadDim = config.qkRopeHeadDim + self.kvLoraRank = config.kvLoraRank + self.vHeadDim = config.vHeadDim + self.qkNopeHeadDim = config.qkNopeHeadDim + self.qHeadDim = config.qkNopeHeadDim + config.qkRopeHeadDim + self.scale = pow(Float(qHeadDim), -0.5) + + if let qLoraRank = self.qLoraRank { + self._qAProj.wrappedValue = Linear( + config.hiddenSize, qLoraRank, bias: config.attentionBias) + self._qALayerNorm.wrappedValue = RMSNorm(dimensions: qLoraRank, eps: 1e-6) + self._qBProj.wrappedValue = Linear(qLoraRank, numHeads * qHeadDim, bias: false) + } else { + self._qProj.wrappedValue = Linear( + config.hiddenSize, numHeads * qHeadDim, bias: false) + } + + self._kvAProjWithMqa.wrappedValue = Linear( + config.hiddenSize, kvLoraRank + qkRopeHeadDim, bias: config.attentionBias) + self._kvALayerNorm.wrappedValue = RMSNorm(dimensions: kvLoraRank, eps: 1e-6) + self._kvBProj.wrappedValue = Linear( + kvLoraRank, numHeads * (qHeadDim - qkRopeHeadDim + vHeadDim), bias: false) + self._oProj.wrappedValue = Linear( + numHeads * vHeadDim, config.hiddenSize, bias: config.attentionBias) + + // YaRN mscale applied to the attention scale, mirroring deepseek_v2.py. + if let ropeScaling = config.ropeScaling { + let mScaleAllDim = ropeScaling["mscale_all_dim"]?.asFloat() ?? 0.0 + let scalingFactor = ropeScaling["factor"]?.asFloat() ?? 1.0 + if mScaleAllDim != 0, scalingFactor > 1 { + let s = 0.1 * mScaleAllDim * log(scalingFactor) + 1.0 + self.scale = self.scale * s * s + } + } + + self.rope = initializeRope( + dims: qkRopeHeadDim, base: config.ropeTheta, traditional: true, + scalingConfig: config.ropeScaling, maxPositionEmbeddings: config.maxPositionEmbeddings) + } + + func callAsFunction( + _ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache? + ) -> MLXArray { + let (B, L, _) = (x.dim(0), x.dim(1), x.dim(2)) + + var q: MLXArray + if qLoraRank == nil { + q = qProj!(x) + } else { + q = qBProj!(qALayerNorm!(qAProj!(x))) + } + q = q.reshaped(B, L, numHeads, qHeadDim).transposed(0, 2, 1, 3) + let splitQ = split(q, indices: [qkNopeHeadDim], axis: -1) + let qNope = splitQ[0] + var qPe = splitQ[1] + + var compressedKv = kvAProjWithMqa(x) + let splitKv = split(compressedKv, indices: [kvLoraRank], axis: -1) + compressedKv = splitKv[0] + var kPe = splitKv[1] + kPe = kPe.reshaped(B, L, 1, qkRopeHeadDim).transposed(0, 2, 1, 3) + + var kv = kvBProj(kvALayerNorm(compressedKv)) + kv = kv.reshaped(B, L, numHeads, -1).transposed(0, 2, 1, 3) + let splitKv2 = split(kv, indices: [qkNopeHeadDim], axis: -1) + let kNope = splitKv2[0] + let values = splitKv2[1] + + let offset = cache?.ropeOffset + qPe = applyRotaryPosition(rope, to: qPe, offset: offset) + kPe = applyRotaryPosition(rope, to: kPe, offset: offset) + kPe = repeated(kPe, count: numHeads, axis: 1) + + let keys = concatenated([kNope, kPe], axis: -1) + let queries = concatenated([qNope, qPe], axis: -1) + + let output = attentionWithCacheUpdate( + queries: queries, keys: keys, values: values, + cache: cache, scale: scale, mask: mask + ) + .transposed(0, 2, 1, 3) + .reshaped(B, L, -1) + + return oProj(output) + } +} + +class DeepseekV2MLP: Module, UnaryLayer { + @ModuleInfo(key: "gate_proj") var gateProj: Linear + @ModuleInfo(key: "up_proj") var upProj: Linear + @ModuleInfo(key: "down_proj") var downProj: Linear + + init(config: DeepseekV2Configuration, hiddenSize: Int? = nil, intermediateSize: Int? = nil) { + let h = hiddenSize ?? config.hiddenSize + let i = intermediateSize ?? config.intermediateSize + self._gateProj.wrappedValue = Linear(h, i, bias: false) + self._upProj.wrappedValue = Linear(h, i, bias: false) + self._downProj.wrappedValue = Linear(i, h, bias: false) + } + + func callAsFunction(_ x: MLXArray) -> MLXArray { + downProj(silu(gateProj(x)) * upProj(x)) + } +} + +class DeepseekV2MoEGate: Module { + let topK: Int + let nRoutedExperts: Int + let routedScalingFactor: Float + let topkMethod: String + let nGroup: Int + let topkGroup: Int + + var weight: MLXArray + + init(config: DeepseekV2Configuration) { + self.topK = config.numExpertsPerTok ?? 1 + self.nRoutedExperts = config.nRoutedExperts ?? 1 + self.routedScalingFactor = config.routedScalingFactor + self.topkMethod = config.topkMethod + self.nGroup = config.nGroup ?? 1 + self.topkGroup = config.topkGroup ?? 1 + self.weight = zeros([nRoutedExperts, config.hiddenSize]) + } + + func callAsFunction(_ x: MLXArray) -> (MLXArray, MLXArray) { + let (bsz, seqLen, _) = (x.dim(0), x.dim(1), x.dim(2)) + + let gates = x.matmul(weight.T) + var scores = softmax(gates, axis: -1, precise: true) + + if topkMethod == "group_limited_greedy" { + // Mask out all but the top `topkGroup` expert groups (by the group's + // max score), then top-k over the surviving experts. + scores = scores.reshaped(bsz, seqLen, nGroup, -1) + let groupScores = MLX.max(scores, axis: -1, keepDims: true) + let k = nGroup - topkGroup + var groupIdx = argPartition(groupScores, kth: k - 1, axis: -2)[.ellipsis, .. MLXArray { + let (inds, scores) = gate(x) + var y = switchMLP(x, inds) + y = weightedExpertSum(y, scores) + if let shared = sharedExperts { + y = y + shared(x) + } + return y + } +} + +class DeepseekV2DecoderLayer: Module { + @ModuleInfo(key: "self_attn") var selfAttn: DeepseekV2Attention + var mlp: UnaryLayer + @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm + @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm + + init(config: DeepseekV2Configuration, layerIdx: Int) { + self._selfAttn.wrappedValue = DeepseekV2Attention(config: config) + + if config.nRoutedExperts != nil, + layerIdx >= config.firstKDenseReplace, + layerIdx % config.moeLayerFreq == 0 + { + self.mlp = DeepseekV2MoE(config: config) + } else { + self.mlp = DeepseekV2MLP(config: config) + } + + self._inputLayerNorm.wrappedValue = RMSNorm( + dimensions: config.hiddenSize, eps: config.rmsNormEps) + self._postAttentionLayerNorm.wrappedValue = RMSNorm( + dimensions: config.hiddenSize, eps: config.rmsNormEps) + } + + func callAsFunction( + _ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache? + ) -> MLXArray { + let r = selfAttn(inputLayerNorm(x), mask: mask, cache: cache) + let h = x + r + return h + mlp(postAttentionLayerNorm(h)) + } +} + +public class DeepseekV2ModelInner: Module { + @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding + fileprivate let layers: [DeepseekV2DecoderLayer] + @ModuleInfo(key: "norm") var norm: RMSNorm + + init(config: DeepseekV2Configuration) { + precondition(config.vocabSize > 0) + self._embedTokens.wrappedValue = Embedding( + embeddingCount: config.vocabSize, dimensions: config.hiddenSize) + self.layers = (0 ..< config.numHiddenLayers).map { + DeepseekV2DecoderLayer(config: config, layerIdx: $0) + } + self._norm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps) + } + + func callAsFunction(_ x: MLXArray, cache: [KVCache]?) -> MLXArray { + var h = embedTokens(x) + let mask = createAttentionMask(h: h, cache: cache?.first) + for (i, layer) in layers.enumerated() { + h = layer(h, mask: mask, cache: cache?[i]) + } + return norm(h) + } +} + +public class DeepseekV2Model: Module, LLMModel, KVCacheDimensionProvider, LoRAModel { + public var kvHeads: [Int] + let config: DeepseekV2Configuration + public let model: DeepseekV2ModelInner + @ModuleInfo(key: "lm_head") var lmHead: Linear + + public init(_ config: DeepseekV2Configuration) { + self.config = config + self.kvHeads = Array(repeating: config.numKeyValueHeads, count: config.numHiddenLayers) + self.model = DeepseekV2ModelInner(config: config) + self._lmHead.wrappedValue = Linear(config.hiddenSize, config.vocabSize, bias: false) + } + + public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray { + lmHead(model(inputs, cache: cache)) + } + + public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + var newWeights = weights + for l in 0 ..< config.numHiddenLayers { + let prefix = "model.layers.\(l)" + for projName in ["gate_proj", "down_proj", "up_proj"] { + for key in ["weight", "scales", "biases"] { + let firstKey = "\(prefix).mlp.experts.0.\(projName).\(key)" + if weights[firstKey] != nil { + let joined = (0 ..< (config.nRoutedExperts ?? 1)).map { + newWeights.removeValue( + forKey: "\(prefix).mlp.experts.\($0).\(projName).\(key)")! + } + newWeights["\(prefix).mlp.switch_mlp.\(projName).\(key)"] = stacked(joined) + } + } + } + } + return newWeights.filter { key, _ in !key.contains("rotary_emb.inv_freq") } + } + + public var loraLayers: [Module] { + model.layers + } +} diff --git a/Libraries/MLXLLM/Models/DeepseekV3.swift b/Libraries/MLXLLM/Models/DeepseekV3.swift index 8e0109046..23f3814b9 100644 --- a/Libraries/MLXLLM/Models/DeepseekV3.swift +++ b/Libraries/MLXLLM/Models/DeepseekV3.swift @@ -195,21 +195,14 @@ class DeepseekV3Attention: Module { kv = kv.reshaped(B, L, self.numHeads, -1).transposed(0, 2, 1, 3) let splitKv = split(kv, indices: [self.qkNopeHeadDim], axis: -1) - var (kNope, values) = (splitKv[0], splitKv[1]) + let (kNope, values) = (splitKv[0], splitKv[1]) let offset = cache?.ropeOffset qPe = applyRotaryPosition(rope, to: qPe, offset: offset) kPe = applyRotaryPosition(rope, to: kPe, offset: offset) kPe = repeated(kPe, count: numHeads, axis: 1) - var keys: MLXArray - if let cache = cache { - (keys, values) = cache.update( - keys: concatenated([kNope, kPe], axis: -1), values: values) - } else { - keys = concatenated([kNope, kPe], axis: -1) - } - + let keys = concatenated([kNope, kPe], axis: -1) let queries = concatenated([qNope, qPe], axis: -1) let output = attentionWithCacheUpdate( @@ -415,7 +408,7 @@ public class DeepseekV3ModelInner: Module { } public class DeepseekV3Model: Module, LLMModel, KVCacheDimensionProvider, LoRAModel { - public var kvHeads: [Int] = [] + public var kvHeads: [Int] var args: DeepseekV3Configuration public var model: DeepseekV3ModelInner @@ -423,6 +416,7 @@ public class DeepseekV3Model: Module, LLMModel, KVCacheDimensionProvider, LoRAMo public init(_ args: DeepseekV3Configuration) { self.args = args + self.kvHeads = Array(repeating: args.numKeyValueHeads, count: args.numHiddenLayers) self.model = DeepseekV3ModelInner(config: args) self._lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabSize, bias: false) } diff --git a/Libraries/MLXLLM/Models/Gemma3Text.swift b/Libraries/MLXLLM/Models/Gemma3Text.swift index 1cc94b059..2283267c0 100644 --- a/Libraries/MLXLLM/Models/Gemma3Text.swift +++ b/Libraries/MLXLLM/Models/Gemma3Text.swift @@ -14,8 +14,8 @@ import MLXNN public struct Gemma3TextConfiguration: Codable { let modelType: String - let hiddenSize: Int - let hiddenLayers: Int + @_spi(GemmaEncoder) public let hiddenSize: Int + @_spi(GemmaEncoder) public let hiddenLayers: Int let intermediateSize: Int let attentionHeads: Int let headDim: Int @@ -232,7 +232,13 @@ class Gemma3MLP: Module { } } -class Gemma3TransformerBlock: Module { +/// A single Gemma 3 transformer layer. +/// +/// Exposed at `@_spi(GemmaEncoder)` scope so opted-in client code can drive the +/// layer stack directly — e.g. encoder-style taps that collect every layer's +/// hidden state — without this becoming advertised public API. Construction +/// remains internal to ``Gemma3Model``. +@_spi(GemmaEncoder) public class Gemma3TransformerBlock: Module { @ModuleInfo(key: "self_attn") var selfAttention: Gemma3Attention @ModuleInfo var mlp: Gemma3MLP @ModuleInfo(key: "input_layernorm") var inputLayerNorm: Gemma.RMSNorm @@ -265,7 +271,7 @@ class Gemma3TransformerBlock: Module { super.init() } - func callAsFunction( + @_spi(GemmaEncoder) public func callAsFunction( _ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache? = nil @@ -283,11 +289,18 @@ class Gemma3TransformerBlock: Module { } public class Gemma3Model: Module { - @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding - @ModuleInfo var layers: [Gemma3TransformerBlock] + /// Token embedding table. + /// + /// Exposed at `@_spi(GemmaEncoder)` scope. Callers must scale its output by + /// `sqrt(hiddenSize)` (computed in bfloat16) to match Gemma 3 semantics, as + /// this type's own `callAsFunction(_:mask:cache:)` does. + @ModuleInfo(key: "embed_tokens") @_spi(GemmaEncoder) public var embedTokens: Embedding + /// The transformer layer stack, exposed at `@_spi(GemmaEncoder)` scope for + /// encoder-style client taps. + @ModuleInfo @_spi(GemmaEncoder) public var layers: [Gemma3TransformerBlock] @ModuleInfo var norm: Gemma.RMSNorm - let config: Gemma3TextConfiguration + @_spi(GemmaEncoder) public let config: Gemma3TextConfiguration init(_ config: Gemma3TextConfiguration) { self.config = config diff --git a/Libraries/MLXLLM/Models/Gemma4Text.swift b/Libraries/MLXLLM/Models/Gemma4Text.swift index 4f199e1eb..bf94058b4 100644 --- a/Libraries/MLXLLM/Models/Gemma4Text.swift +++ b/Libraries/MLXLLM/Models/Gemma4Text.swift @@ -65,6 +65,11 @@ public struct Gemma4TextConfiguration: Codable, Sendable { var attentionKeqV: Bool = false var finalLogitSoftcapping: Float = 30.0 var useDoubleWideMlp: Bool = true + // MoE block (E-series: enable_moe_block=true) + var enableMoEBlock: Bool = false + var numExperts: Int? + var topKExperts: Int? + var moeIntermediateSize: Int? var layerTypes: [String] = [] var tieWordEmbeddings: Bool = true @@ -98,6 +103,10 @@ public struct Gemma4TextConfiguration: Codable, Sendable { case attentionKeqV = "attention_k_eq_v" case finalLogitSoftcapping = "final_logit_softcapping" case useDoubleWideMlp = "use_double_wide_mlp" + case enableMoEBlock = "enable_moe_block" + case numExperts = "num_experts" + case topKExperts = "top_k_experts" + case moeIntermediateSize = "moe_intermediate_size" case layerTypes = "layer_types" case tieWordEmbeddings = "tie_word_embeddings" case ropeParameters = "rope_parameters" @@ -142,6 +151,12 @@ public struct Gemma4TextConfiguration: Codable, Sendable { try container.decodeIfPresent(Float.self, forKey: .finalLogitSoftcapping) ?? 30.0 self.useDoubleWideMlp = try container.decodeIfPresent(Bool.self, forKey: .useDoubleWideMlp) ?? true + self.enableMoEBlock = + try container.decodeIfPresent(Bool.self, forKey: .enableMoEBlock) ?? false + self.numExperts = try container.decodeIfPresent(Int.self, forKey: .numExperts) + self.topKExperts = try container.decodeIfPresent(Int.self, forKey: .topKExperts) + self.moeIntermediateSize = + try container.decodeIfPresent(Int.self, forKey: .moeIntermediateSize) if let decoded = try container.decodeIfPresent([String].self, forKey: .layerTypes) { self.layerTypes = decoded } else { @@ -412,6 +427,85 @@ private class Gemma4MLP: Module { } } +// MARK: - MoE Router + Experts (E-series) + +// MoE block ported from MLXVLM/Models/Gemma4.swift (ml-explore PRs #180, #228). yooz-engine. + +private class Gemma4TextRouter: Module { + let topKExperts: Int + let hiddenSize: Int + let rmsNormEps: Float + private let rootSize: Float + + @ModuleInfo(key: "proj") var proj: Linear + @ParameterInfo(key: "scale") var scale: MLXArray + @ParameterInfo(key: "per_expert_scale") var perExpertScale: MLXArray + + init(_ config: Gemma4TextConfiguration) { + guard let numExperts = config.numExperts, let topKExperts = config.topKExperts else { + fatalError("Gemma4 MoE router requires numExperts and topKExperts in config") + } + self.topKExperts = topKExperts + self.hiddenSize = config.hiddenSize + self.rmsNormEps = config.rmsNormEps + self.rootSize = pow(Float(config.hiddenSize), -0.5) + + self._proj.wrappedValue = Linear(config.hiddenSize, numExperts, bias: false) + self._scale.wrappedValue = MLXArray.ones([config.hiddenSize]) + self._perExpertScale.wrappedValue = MLXArray.ones([numExperts]) + super.init() + } + + func callAsFunction(_ x: MLXArray) -> (MLXArray, MLXArray) { + let normed = MLXFast.rmsNorm( + x, weight: (scale * rootSize).asType(x.dtype), eps: rmsNormEps) + let scores = proj(normed) + let topKIndices = MLX.argPartition(scores, kth: -topKExperts, axis: -1)[ + .ellipsis, (-topKExperts)..., + ] + var topKWeights = MLX.takeAlong(scores, topKIndices, axis: -1) + topKWeights = MLX.softmax(topKWeights, axis: -1) + topKWeights = topKWeights * perExpertScale[topKIndices].asType(topKWeights.dtype) + return (topKIndices, topKWeights) + } +} + +private class Gemma4TextExperts: Module { + @ModuleInfo(key: "switch_glu") var switchGLU: SwitchGLU + + init(_ config: Gemma4TextConfiguration) { + guard let numExperts = config.numExperts, + let moeIntermediateSize = config.moeIntermediateSize + else { + fatalError("Gemma4 MoE experts require numExperts and moeIntermediateSize in config") + } + self._switchGLU.wrappedValue = SwitchGLU( + inputDims: config.hiddenSize, + hiddenDims: moeIntermediateSize, + numExperts: numExperts, + activation: geluApproximate, + bias: false + ) + super.init() + } + + func callAsFunction( + _ x: MLXArray, topKIndices: MLXArray, topKWeights: MLXArray + ) -> MLXArray { + let batch = x.dim(0) + let length = x.dim(1) + let hidden = x.dim(2) + let topK = topKIndices.dim(-1) + + let expertOutput = switchGLU( + x.reshaped(batch * length, hidden), + topKIndices.reshaped(batch * length, topK) + ) + let weights = topKWeights.reshaped(batch * length, topK).asType(expertOutput.dtype) + return weightedExpertSum(expertOutput, weights).reshaped(batch, length, hidden) + } +} + // MARK: - Decoder Layer private class Gemma4DecoderLayer: Module { @@ -419,6 +513,7 @@ private class Gemma4DecoderLayer: Module { let layerIdx: Int let layerType: String let hiddenSizePerLayerInput: Int + let enableMoE: Bool @ModuleInfo(key: "self_attn") var selfAttn: Gemma4Attention @ModuleInfo var mlp: Gemma4MLP @@ -427,6 +522,13 @@ private class Gemma4DecoderLayer: Module { @ModuleInfo(key: "pre_feedforward_layernorm") var preFeedforwardLayernorm: RMSNorm @ModuleInfo(key: "post_feedforward_layernorm") var postFeedforwardLayernorm: RMSNorm + // MoE block (E-series): router, experts, and their extra norms + @ModuleInfo(key: "router") var router: Gemma4TextRouter? + @ModuleInfo(key: "experts") var experts: Gemma4TextExperts? + @ModuleInfo(key: "post_feedforward_layernorm_1") var postFeedforwardLayernorm1: RMSNorm? + @ModuleInfo(key: "post_feedforward_layernorm_2") var postFeedforwardLayernorm2: RMSNorm? + @ModuleInfo(key: "pre_feedforward_layernorm_2") var preFeedforwardLayernorm2: RMSNorm? + // Per-layer input (PLE) gating @ModuleInfo(key: "per_layer_input_gate") var perLayerInputGate: Linear? @ModuleInfo(key: "per_layer_projection") var perLayerProjection: Linear? @@ -448,6 +550,7 @@ private class Gemma4DecoderLayer: Module { self.layerIdx = layerIdx self.layerType = config.layerTypes[layerIdx] self.hiddenSizePerLayerInput = config.hiddenSizePerLayerInput + self.enableMoE = config.enableMoEBlock self._selfAttn.wrappedValue = Gemma4Attention(config, layerIdx: layerIdx) self._mlp.wrappedValue = Gemma4MLP(config, layerIdx: layerIdx) @@ -461,6 +564,17 @@ private class Gemma4DecoderLayer: Module { self._postFeedforwardLayernorm.wrappedValue = RMSNorm( dimensions: config.hiddenSize, eps: config.rmsNormEps) + if config.enableMoEBlock { + self._router.wrappedValue = Gemma4TextRouter(config) + self._experts.wrappedValue = Gemma4TextExperts(config) + self._postFeedforwardLayernorm1.wrappedValue = RMSNorm( + dimensions: config.hiddenSize, eps: config.rmsNormEps) + self._postFeedforwardLayernorm2.wrappedValue = RMSNorm( + dimensions: config.hiddenSize, eps: config.rmsNormEps) + self._preFeedforwardLayernorm2.wrappedValue = RMSNorm( + dimensions: config.hiddenSize, eps: config.rmsNormEps) + } + if hiddenSizePerLayerInput > 0 { self._perLayerInputGate.wrappedValue = Linear( config.hiddenSize, hiddenSizePerLayerInput, bias: false) @@ -492,10 +606,33 @@ private class Gemma4DecoderLayer: Module { var out = _addRMSNorm(residual, attnOut, postAttentionLayernorm.weight) let residual2 = out - out = preFeedforwardLayernorm(out) - out = mlp(out) - // Fused: residual + RMSNorm(out) * weight - out = _addRMSNorm(residual2, out, postFeedforwardLayernorm.weight) + + if enableMoE, + let router, + let experts, + let pfln1 = postFeedforwardLayernorm1, + let pfln2 = postFeedforwardLayernorm2, + let pfln2pre = preFeedforwardLayernorm2 + { + // Dense branch + var dense = preFeedforwardLayernorm(out) + dense = mlp(dense) + dense = pfln1(dense) + + // Sparse branch + let (topKIndices, topKWeights) = router(out) + var sparse = pfln2pre(out) + sparse = experts(sparse, topKIndices: topKIndices, topKWeights: topKWeights) + sparse = pfln2(sparse) + + // Combine then apply shared post-ff norm + out = _addRMSNorm(residual2, dense + sparse, postFeedforwardLayernorm.weight) + } else { + out = preFeedforwardLayernorm(out) + out = mlp(out) + // Fused: residual + RMSNorm(out) * weight + out = _addRMSNorm(residual2, out, postFeedforwardLayernorm.weight) + } // PLE gating if let gate = perLayerInputGate, @@ -714,9 +851,13 @@ public class Gemma4TextModel: Module, LLMModel, KVCacheDimensionProvider { } public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + // MoE expert weight remapping ported from MLXVLM/Models/Gemma4.swift. yooz-engine. + // HuggingFace stores expert weights as fused gate_up_proj; SwitchGLU expects + // separate gate_proj and up_proj, each shaped [numExperts, hiddenDims, inputDims]. let firstKvSharedLayerIdx = config.numHiddenLayers - config.numKvSharedLayers var sanitized = [String: MLXArray]() - for (k, v) in weights { + for (key, value) in weights { + let k = key // Skip vision/audio/rotary weights if k.contains("self_attn.rotary_emb") || k.contains("input_max") @@ -740,7 +881,37 @@ public class Gemma4TextModel: Module, LLMModel, KVCacheDimensionProvider { { continue } - sanitized[k] = v + + // Remap .experts.down_proj -> .experts.switch_glu.down_proj.weight + if k.hasSuffix(".experts.down_proj") { + sanitized[ + k.replacingOccurrences( + of: ".experts.down_proj", + with: ".experts.switch_glu.down_proj.weight" + ) + ] = value + continue + } + + // Remap .experts.gate_up_proj -> split into gate_proj + up_proj + if k.hasSuffix(".experts.gate_up_proj") { + let mid = value.dim(-2) / 2 + sanitized[ + k.replacingOccurrences( + of: ".experts.gate_up_proj", + with: ".experts.switch_glu.gate_proj.weight" + ) + ] = value[.ellipsis, .. MLXArray { + let (B, L) = (x.dim(0), x.dim(1)) + + var queries = wq(x) + var keys = wk(x) + var values = wv(x) + + queries = queries.reshaped(B, L, args.attentionHeads, -1).transposed(0, 2, 1, 3) + keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3) + values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3) + + let offset = cache?.ropeOffset + queries = applyRotaryPosition(rope, to: queries, offset: offset) + keys = applyRotaryPosition(rope, to: keys, offset: offset) + + if let queryNorm, let keyNorm { + queries = queryNorm(queries) + keys = keyNorm(keys) + } + + let output = attentionWithCacheUpdate( + queries: queries, + keys: keys, + values: values, + cache: cache, + scale: scale, + mask: mask + ) + .transposed(0, 2, 1, 3) + .reshaped(B, L, -1) + + return wo(output) + } +} + +class HunyuanMLP: Module, UnaryLayer { + @ModuleInfo(key: "gate_proj") var gate: Linear + @ModuleInfo(key: "down_proj") var down: Linear + @ModuleInfo(key: "up_proj") var up: Linear + + public init(dimensions: Int, hiddenDimensions: Int) { + _gate.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false) + _down.wrappedValue = Linear(hiddenDimensions, dimensions, bias: false) + _up.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false) + } + + public func callAsFunction(_ x: MLXArray) -> MLXArray { + down(silu(gate(x)) * up(x)) + } +} + +class HunyuanTransformerBlock: Module { + @ModuleInfo(key: "self_attn") var attention: HunyuanAttention + let mlp: HunyuanMLP + + @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm + @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm + + public init(_ args: HunyuanConfiguration) { + _attention.wrappedValue = HunyuanAttention(args) + self.mlp = HunyuanMLP( + dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize) + _inputLayerNorm.wrappedValue = RMSNorm( + dimensions: args.hiddenSize, eps: args.rmsNormEps) + _postAttentionLayerNorm.wrappedValue = RMSNorm( + dimensions: args.hiddenSize, eps: args.rmsNormEps) + } + + public func callAsFunction( + _ x: MLXArray, mask: MLXFast.ScaledDotProductAttentionMaskMode, cache: KVCache? + ) -> MLXArray { + var r = attention(inputLayerNorm(x), mask: mask, cache: cache) + let h = x + r + r = mlp(postAttentionLayerNorm(h)) + let out = h + r + return out + } +} + +public class HunyuanModelInner: Module { + @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding + + fileprivate let layers: [HunyuanTransformerBlock] + let norm: RMSNorm + + public init(_ args: HunyuanConfiguration) { + precondition(args.vocabularySize > 0) + + _embedTokens.wrappedValue = Embedding( + embeddingCount: args.vocabularySize, dimensions: args.hiddenSize) + + self.layers = (0 ..< args.hiddenLayers).map { _ in + HunyuanTransformerBlock(args) + } + self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps) + } + + public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray { + var h = embedTokens(inputs) + + let mask = createAttentionMask(h: h, cache: cache?.first) + + for (i, layer) in layers.enumerated() { + h = layer(h, mask: mask, cache: cache?[i]) + } + + return norm(h) + } +} + +public class HunyuanModel: Module, LLMModel, KVCacheDimensionProvider { + public let vocabularySize: Int + public let kvHeads: [Int] + + public let model: HunyuanModelInner + let configuration: HunyuanConfiguration + + @ModuleInfo(key: "lm_head") var lmHead: Linear? + + public init(_ args: HunyuanConfiguration) { + self.configuration = args + self.vocabularySize = args.vocabularySize + self.kvHeads = (0 ..< args.hiddenLayers).map { _ in args.kvHeads } + self.model = HunyuanModelInner(args) + + if !args.tieWordEmbeddings { + _lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabularySize, bias: false) + } + } + + public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray { + var out = model(inputs, cache: cache) + if let lmHead { + out = lmHead(out) + } else { + out = model.embedTokens.asLinear(out) + } + return out + } + + public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + var weights = weights + + if configuration.tieWordEmbeddings { + weights["lm_head.weight"] = nil + } + + return weights + } +} + +public struct HunyuanConfiguration: Codable, Sendable { + public var hiddenSize: Int + public var hiddenLayers: Int + public var intermediateSize: Int + public var attentionHeads: Int + public var kvHeads: Int + public var rmsNormEps: Float + public var vocabularySize: Int + public var headDim: Int + public var ropeTheta: Float = 10000 + public var attentionBias: Bool = false + public var useQkNorm: Bool = true + public var ropeScaling: [String: StringOrNumber]? = nil + public var tieWordEmbeddings: Bool = false + public var maxPositionEmbeddings: Int = 32768 + + enum CodingKeys: String, CodingKey { + case hiddenSize = "hidden_size" + case hiddenLayers = "num_hidden_layers" + case intermediateSize = "intermediate_size" + case attentionHeads = "num_attention_heads" + case kvHeads = "num_key_value_heads" + case rmsNormEps = "rms_norm_eps" + case vocabularySize = "vocab_size" + case headDim = "head_dim" + case ropeTheta = "rope_theta" + case attentionBias = "attention_bias" + case useQkNorm = "use_qk_norm" + case ropeScaling = "rope_scaling" + case tieWordEmbeddings = "tie_word_embeddings" + case maxPositionEmbeddings = "max_position_embeddings" + } + + /// Some Hunyuan configs (e.g. Hy-MT2-7B) spell `head_dim` as `attention_head_dim`. + private enum AltCodingKeys: String, CodingKey { + case attentionHeadDim = "attention_head_dim" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + + self.hiddenSize = try container.decode(Int.self, forKey: .hiddenSize) + self.hiddenLayers = try container.decode(Int.self, forKey: .hiddenLayers) + self.intermediateSize = try container.decode(Int.self, forKey: .intermediateSize) + self.attentionHeads = try container.decode(Int.self, forKey: .attentionHeads) + self.kvHeads = try container.decode(Int.self, forKey: .kvHeads) + self.rmsNormEps = try container.decode(Float.self, forKey: .rmsNormEps) + self.vocabularySize = try container.decode(Int.self, forKey: .vocabularySize) + let alt = try decoder.container(keyedBy: AltCodingKeys.self) + self.headDim = + try container.decodeIfPresent(Int.self, forKey: .headDim) + ?? alt.decodeIfPresent(Int.self, forKey: .attentionHeadDim) + ?? (hiddenSize / attentionHeads) + self.ropeTheta = + try container.decodeIfPresent(Float.self, forKey: .ropeTheta) ?? 10000 + self.attentionBias = + try container.decodeIfPresent(Bool.self, forKey: .attentionBias) ?? false + self.useQkNorm = + try container.decodeIfPresent(Bool.self, forKey: .useQkNorm) ?? true + self.ropeScaling = try container.decodeIfPresent( + [String: StringOrNumber].self, forKey: .ropeScaling) + self.tieWordEmbeddings = + try container.decodeIfPresent(Bool.self, forKey: .tieWordEmbeddings) ?? false + self.maxPositionEmbeddings = + try container.decodeIfPresent(Int.self, forKey: .maxPositionEmbeddings) ?? 32768 + } +} + +// MARK: - LoRA + +extension HunyuanModel: LoRAModel { + public var loraLayers: [Module] { + model.layers + } +} diff --git a/Libraries/MLXLMCommon/AttentionUtils.swift b/Libraries/MLXLMCommon/AttentionUtils.swift index 5b4e6c76b..2829a818a 100644 --- a/Libraries/MLXLMCommon/AttentionUtils.swift +++ b/Libraries/MLXLMCommon/AttentionUtils.swift @@ -51,7 +51,25 @@ public func attentionWithCacheUpdate( mask: mask ) } - if let quantizedKVCache = cache as? QuantizedKVCacheProtocol { + if let turboCache = cache as? TurboQuantKVCache { + let L = queries.dim(2) + if L > 1 && !turboCache.isCompressed { + // Prefill (L>1) on a raw cache: plain update + standard SDPA, // zero overhead; compression is deferred to the first decode step. + let (cachedKeys, cachedValues) = turboCache.update(keys: keys, values: values) + return MLXFast.scaledDotProductAttention( + queries: queries, keys: cachedKeys, values: cachedValues, + scale: scale, mask: mask + ) + } + // Decode (L=1) or any call once the cache is compressed (speculative + // verify chunks, multi-turn re-prefill): the compressed path. The raw + // update() path is invalid after compression, its raw buffers are + // gone. First decode call triggers compressRawCache(). + return turboCache.compressedAttention( + queries: queries, keys: keys, values: values, + scale: scale, mask: mask + ) + } else if let quantizedKVCache = cache as? QuantizedKVCacheProtocol { let (quantizedKeys, quantizedValues) = quantizedKVCache.updateQuantized( keys: keys, values: values) return quantizedScaledDotProductAttention( diff --git a/Libraries/MLXLMCommon/BaseConfiguration.swift b/Libraries/MLXLMCommon/BaseConfiguration.swift index 1a57fab61..9d1dd3a8c 100644 --- a/Libraries/MLXLMCommon/BaseConfiguration.swift +++ b/Libraries/MLXLMCommon/BaseConfiguration.swift @@ -185,12 +185,28 @@ public struct BaseConfiguration: Codable, Sendable { } } + private struct TextConfiguration: Codable, Sendable { + var eosTokenIds: IntOrIntArray? + + enum CodingKeys: String, CodingKey { + case eosTokenIds = "eos_token_id" + } + } + /// Internal storage for quantization details extracted from `config.json`. var quantizationContainer: QuantizationContainer? + /// Text-model metadata nested by composite model configurations. + private var textConfiguration: TextConfiguration? + /// EOS token IDs from config.json. Can be a single Int or an array of Ints. public var eosTokenIds: IntOrIntArray? + /// EOS token IDs declared at either the model root or in `text_config`. + public var effectiveEOSTokenIds: Set { + Set(eosTokenIds?.values ?? textConfiguration?.eosTokenIds?.values ?? []) + } + /// The default quantization settings. @available(*, deprecated, message: "Please use perLayerQuantization instead") public var quantization: Quantization? { @@ -205,6 +221,7 @@ public struct BaseConfiguration: Codable, Sendable { enum CodingKeys: String, CodingKey { case modelType = "model_type" case quantizationContainer = "quantization" + case textConfiguration = "text_config" case eosTokenIds = "eos_token_id" } } diff --git a/Libraries/MLXLMCommon/Documentation.docc/Documentation.md b/Libraries/MLXLMCommon/Documentation.docc/Documentation.md index c9dd22b6b..f1d480d1f 100644 --- a/Libraries/MLXLMCommon/Documentation.docc/Documentation.md +++ b/Libraries/MLXLMCommon/Documentation.docc/Documentation.md @@ -7,6 +7,7 @@ Common language model code. - - - +- ## Other MLX Libraries Packages diff --git a/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md b/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md new file mode 100644 index 000000000..5338e932b --- /dev/null +++ b/Libraries/MLXLMCommon/Documentation.docc/kv-cache-quantization.md @@ -0,0 +1,109 @@ +# KV Cache Quantization + +Reduce the memory footprint of long-context generation by compressing the +key/value cache. + +## Overview + +At long context lengths the KV cache, not the weights, dominates memory. Two +mechanisms are available through ``GenerateParameters``: + +- **Affine quantization**: `kvBits`, `kvGroupSize`, and `quantizedKVStart` + quantize both K and V with MLX's affine scheme. The named schemes + `"affine4"` and `"affine8"` are shorthands for the same path. +- **TurboQuant schemes via `kvScheme`**: vectors are rotated with a + Walsh-Hadamard transform and quantized against a Lloyd-Max codebook. + TurboQuant schemes are asymmetric. Keys and values can use different + precision because attention quality is far more sensitive to key error + (softmax amplifies it) than to value error (linear averaging smooths it). + +```swift +var parameters = GenerateParameters() +parameters.kvScheme = "turbo8v3" // recommended default +``` + +Prefill is unaffected: the cache stores raw fp16 during prompt processing, +compresses on the first decode step, and encodes each new token incrementally +afterwards. Compressed caches round-trip through prompt-cache save/restore. + +## Scheme reference + +Scheme names read `turbov`; `0` means keys stay fp16. + +| scheme | keys | values | KV compression | character | +|---|---|---|---|---| +| `affine8` | 8-bit affine | 8-bit affine | 1.88x | near-lossless on most models, full decode speed | +| `affine4` | 4-bit affine | 4-bit affine | 3.56x | collapses on some families; validate first | +| `turbo0v4` | fp16 | 4-bit turbo | 1.58x | safest start; beats affine8 quality on most models tested | +| `turbo0v3` | fp16 | 3-bit turbo | 1.66x | light value compression | +| `turbo0v2` | fp16 | 2-bit turbo | 1.58x† | aggressive value compression | +| `turbo8v4` | 8-bit affine | 4-bit turbo | 2.51x | conservative asymmetric | +| `turbo8v3` | 8-bit affine | 3-bit turbo | 2.75x | recommended default | +| `turbo8v2` | 8-bit affine | 2-bit turbo | 2.32x† | memory-bound long context | +| `turbo4`, `turbo3`, `turbo2` | turbo | turbo | up to 3.4x† | maximum compression; key sensitivity varies strongly by family | + +† boundary-layer protection auto-engages for fragile schemes (turbo-quantized +keys or 2-bit values): the first and last two attention layers use 8-bit +affine instead of the turbo scheme. + +## Choosing a scheme + +Start asymmetric and light, verify output quality on your model, then +compress further: + +1. `turbo0v4`: keys untouched. If this is not faithful, the model is + unusually quantization-sensitive; stop here. +2. `turbo8v4`: 8-bit keys are near-lossless at half the key bytes. +3. `turbo8v3`: the sweet spot for most dense models. +4. `turbo8v2`: memory-bound long context, after validating step 3. +5. Symmetric `turbo4`, `turbo3`, `turbo2`: maximum compression. Key + compression is where models break. Per-dimension key calibration + (computed automatically from the prefill cache at compression time) + equalizes post-rotation key variance and recovers most of the gap on + sensitive families; sensitivity still varies, so validate on your model + before deploying. + +## Family sensitivity (measured) + +WikiText-2 decode-time KL divergence against an fp16 cache with identical +weights, so the numbers isolate cache compression from weight quantization: + +- Mistral-class models tolerate even symmetric `turbo4` (KLD 0.040 at 2.8x + on Mistral-7B). +- qk-norm families (Qwen3) and small models are the most sensitive to key + compression. Per-dimension key calibration recovers most of it: turbo4 + KLD 2.65 to 0.15 on Qwen3-1.7B, 2.76 to 0.036 on Phi-4-mini, 0.62 to + 0.060 on Qwen2.5-7B, all at about 3.3x. The asymmetric family stays + healthy with or without calibration. +- Qwen2.5-class models carry large key outliers in their first and last + layers; attention scores are computed in f32 for this reason, and affine8 + is notably weaker on this family (KLD 0.041) than the turbo value schemes + (turbo0v4 KLD 0.005). +- Phi-family models are unusually affine-friendly (affine8 KLD 0.0004). +- Symmetric key compression improves with model size (KLD 2.7 at 1.7B, 0.6 + at 7B on the same family). + +## Performance + +TurboQuant decode runs through JIT-compiled Metal kernels; kernels are +compiled once per model shape, not per call. Single-token decode is a +SIMD-parallel flash kernel (packed, raw fp16, or 8-bit affine key scoring +with inline value dequantization). Measured on Qwen3-1.7B (M5 Max, fp16 +150 tok/s): turbo8v3 114, turbo4 122, turbo0v4 102. Prefill stays raw +fp16, so prefill throughput is unaffected. Use TurboQuant when memory is +the constraint (long contexts, larger models per machine); use affine8 +when exact fp16-parity decode matters more than footprint. + +## Limitations + +- Rotating and sliding-window cache layers (most Gemma layers) are not + converted; their storage is already bounded by the window size and their + eviction layout does not fit the sequential-append compression path. On + mixed models the global (non-rotating) layers, where cache memory + actually grows, do compress, and a one-time notice lists the layers that + kept fp16 rotating caches. Hybrid recurrent layers are likewise left + untouched. +- Unrecognized scheme strings are ignored. +- For memory estimation with wired limits see ; effective + bytes per element follow from the table above (for example `turbo8v3` is + about 0.73 bytes per K/V element pair average against 4 for fp16). diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index a0f920082..df34922e1 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -74,8 +74,21 @@ public struct GenerateParameters: Sendable { public var quantizedKVStart: Int /// KV cache compression scheme. Overrides kvBits when set. - /// Built-in: "affine4", "affine8" (equivalent to kvBits 4/8). - /// Extensible for custom schemes (e.g. WHT-based compression). + /// + /// Built-in affine: "affine4", "affine8" (equivalent to kvBits 4/8). + /// + /// TurboQuant schemes (`turbov`, keys × values). Start light, + /// verify quality on your model, then ratchet V: + /// - "turbo0v4" FP16 K + 4-bit V, the safest start + /// - "turbo0v3"/"turbo0v2" FP16 K, more aggressive V + /// - "turbo8v4" 8-bit affine K + 4-bit V, conservative asymmetric + /// - "turbo8v3" recommended default (near-lossless K, compressed V) + /// - "turbo8v2" aggressive V (boundary-layer protection auto-engages) + /// - "turbo4"/"turbo4v2"/"turbo3"/"turbo2" turbo-quantized keys: + /// maximum compression; K sensitivity varies by model family, so + /// validate on your model (asym is the recommended starting point). + /// + /// Unrecognized schemes are ignored. See `resolveTurboScheme`. public var kvScheme: String? /// Sampling temperature @@ -746,17 +759,35 @@ public struct TokenIterator: TokenIteratorProtocol { return nil } - // save current value -- this will be returned - let previousY = y - - // compute the next state and async eval the next token - let token = step(previous: previousY) - y = .init(tokens: token) - asyncEval(token) + // Drain autoreleased MLX temporaries every step: a full model forward + // produces hundreds of autoreleased wrapper objects, and the token + // loop never returns to a pool boundary on its own; without this, + // long generations grow host memory without bound. + return autoreleasepool { + // save current value -- this will be returned + let previousY = y - tokenCount += 1 + // compute the next state and async eval the next token + let token = step(previous: previousY) + y = .init(tokens: token) + // Evaluate the cache state together with the token: caches update + // through functional ops (concatenation, slice assignment), and an + // unevaluated chain of those updates keeps every prior step's + // intermediates alive. Python mlx-lm settles cache state the same + // way in its generation loop. + asyncEval([token] + cache.flatMap { $0.state }) + + tokenCount += 1 + + // Periodically return freed buffers that cannot be reused (odd or + // monotonically growing sizes accumulate in the pool otherwise). + // Matches mlx-lm's clear cadence. + if tokenCount % 256 == 0 { + MLX.Memory.clearCache() + } - return previousY.tokens.item(Int.self) + return previousY.tokens.item(Int.self) + } } } @@ -861,7 +892,8 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { cache: &cache, kvBits: parameters.kvBits, kvGroupSize: parameters.kvGroupSize, - quantizedKVStart: parameters.quantizedKVStart + quantizedKVStart: parameters.quantizedKVStart, + kvScheme: parameters.kvScheme ) } @@ -1018,10 +1050,12 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { return token } - // Run a new speculation round + // Run a new speculation round. Pooled: a round runs draft + verify + // forwards whose autoreleased MLX temporaries otherwise accumulate + // across the whole generation. pendingTokens.removeAll(keepingCapacity: true) pendingIndex = 0 - speculateRound() + autoreleasepool { speculateRound() } if pendingTokens.isEmpty { return nil @@ -1168,7 +1202,7 @@ private func runSynchronousGenerationLoop( var iterator = iterator var stopReason: GenerateStopReason? - while let token = iterator.next() { + while let token = autoreleasepool(invoking: { iterator.next() }) { // Compute the timing for the prompt. if promptTime == 0 { let now = Date.timeIntervalSinceReferenceDate @@ -1864,12 +1898,12 @@ private func generateLoopTask( // Check cancellation BEFORE iterator.next(): next() calls asyncEval() to // pipeline the next GPU evaluation, so checking after it (the previous // `while let token = iterator.next()` form) allowed one extra asyncEval to be - // submitted post-cancellation — which faults if the app has backgrounded + // submitted post-cancellation, which faults if the app has backgrounded // (kIOGPUCommandBufferCallbackErrorBackgroundExecutionNotPermitted). The // post-loop block below assigns `.cancelled`; Stream().synchronize() still // settles any in-flight evaluation at the end of the task body. tokenLoop: while !Task.isCancelled { - guard let token = iterator.next() else { break } + guard let token = autoreleasepool(invoking: { iterator.next() }) else { break } if promptTime == 0 { let now = Date.timeIntervalSinceReferenceDate diff --git a/Libraries/MLXLMCommon/KVCache.swift b/Libraries/MLXLMCommon/KVCache.swift index c4fa11fa2..6c56e92ca 100644 --- a/Libraries/MLXLMCommon/KVCache.swift +++ b/Libraries/MLXLMCommon/KVCache.swift @@ -174,6 +174,10 @@ open class BaseKVCache: KVCache { public var offset: Int = 0 public var maxSize: Int? { nil } + /// RoPE offset for this cache. `open` so subclasses can return a non-scalar + /// offset (e.g. a batched cache's per-row `.batch(...)`). + open var ropeOffset: RoPEOffset { .scalar(offset) } + public func innerState() -> [MLXArray] { [] } open func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) { @@ -797,7 +801,7 @@ public class RotatingKVCache: BaseKVCache, CustomDebugStringConvertible { } } -private func resolvedKVQuantizationGroupSize( +func resolvedKVQuantizationGroupSize( requested: Int, keyHeadDim: Int, valueHeadDim: Int @@ -1577,6 +1581,7 @@ private func cacheClassName(_ cache: KVCache) -> String { case is ArraysCache: return "ArraysCache" case is RotatingKVCache: return "RotatingKVCache" case is QuantizedKVCache: return "QuantizedKVCache" + case is TurboQuantKVCache: return "TurboQuantKVCache" case is KVCacheSimple: return "KVCache" case is CacheList: return "CacheList" default: return "KVCache" @@ -1730,6 +1735,22 @@ private func restoreCacheFromMetaState( cache.restoreFromMetaState(state: state, savedMetaState: metaState) return cache + case "TurboQuantKVCache": + guard metaState.count >= 5, + let bits = Int(metaState[1]), + let keyBits = Int(metaState[2]), + let valueBits = Int(metaState[3]), + let seed = UInt64(metaState[4]) + else { + throw KVCacheError( + message: "Invalid TurboQuantKVCache metaState") + } + let cache = TurboQuantKVCache( + bits: bits, keyBits: keyBits, valueBits: valueBits, seed: seed) + cache.state = state + cache.metaState = metaState + return cache + case "CacheList": return try CacheList.fromState(state: state, metaState: metaState) @@ -2004,8 +2025,9 @@ public func resolveAffineScheme(_ scheme: String?) -> (bits: Int, groupSize: Int /// - kvGroupSize: Group size for quantization /// - quantizedKVStart: Token count threshold to begin quantizing /// - kvScheme: Scheme selector; overrides kvBits when it names a built-in -/// affine scheme ("affine4", "affine8"). Unrecognized schemes are left to -/// custom cache implementations and do not quantize here. +/// affine scheme ("affine4", "affine8") or a TurboQuant scheme +/// ("turbo4", "turbo4v2", ...). Unrecognized schemes are left to custom +/// cache implementations and do not quantize here. public func maybeQuantizeKVCache( cache: inout [KVCache], kvBits: Int?, @@ -2013,6 +2035,18 @@ public func maybeQuantizeKVCache( quantizedKVStart: Int = 0, kvScheme: String? = nil ) { + // TurboQuant schemes convert eligible layers to TurboQuantKVCache + // (handled in TurboQuantKVCache.swift to keep this file scheme-agnostic). + if let scheme = kvScheme, let turbo = resolveTurboScheme(scheme) { + maybeTurboQuantizeKVCache( + cache: &cache, + keyBits: turbo.keyBits, + valueBits: turbo.valueBits, + quantizedKVStart: quantizedKVStart + ) + return + } + // Resolve effective bits: kvScheme overrides kvBits. let effectiveBits: Int let effectiveGroupSize: Int diff --git a/Libraries/MLXLMCommon/Load.swift b/Libraries/MLXLMCommon/Load.swift index bae3b5ef5..f5d9cdab5 100644 --- a/Libraries/MLXLMCommon/Load.swift +++ b/Libraries/MLXLMCommon/Load.swift @@ -61,6 +61,7 @@ public func loadWeights( weights = model.sanitize(weights: weights, metadata: metadata) // quantize if needed + var builtMXFP8Layer = false if quantization != nil || perLayerQuantization != nil { quantize( model: model, @@ -82,6 +83,7 @@ public func loadWeights( // Use the explicit initializer with biases:nil; weight/scales are // overwritten by model.update(parameters:) with the real checkpoint data. if mode == .mxfp8, let linear = layer as? Linear { + builtMXFP8Layer = true let w = linear.weight let dummyScales = MLXArray.zeros([w.dim(0), max(1, w.dim(1) / groupSize)]) return QuantizedLinear( @@ -96,12 +98,15 @@ public func loadWeights( // apply the loaded weights // - // Use .noUnusedKeys only (not .all) because pre-quantized mxfp8 models have a - // different weight layout than the affine QuantizedLinear created by quantize(model:) - // above (shapes are reconciled via _updateInternal once the real data is applied). - // .noUnusedKeys still catches any unexpected extra keys in the checkpoint. + // Strict verification (.all) by default — a checkpoint whose packed shapes don't + // match the quantized modules must fail loudly, not decode as garbage (upstream + // regression #395). Relax to .noUnusedKeys ONLY when an mxfp8 layer was built: + // pre-quantized mxfp8 models carry a different weight layout than the affine + // QuantizedLinear placeholder above, and their shapes are reconciled via + // _updateInternal once the real data is applied. let parameters = ModuleParameters.unflattened(weights) - try model.update(parameters: parameters, verify: [.noUnusedKeys]) + try model.update( + parameters: parameters, verify: builtMXFP8Layer ? [.noUnusedKeys] : [.all]) eval(model) } diff --git a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift b/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift index b990a2cdb..b0a75fdae 100644 --- a/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift +++ b/Libraries/MLXLMCommon/MTPSpeculativeTokenIterator.swift @@ -423,7 +423,7 @@ public struct MTPSpeculativeTokenIterator: TokenIteratorProtocol { // Run a new speculation round (may transition to passthrough). pendingTokens.removeAll(keepingCapacity: true) pendingIndex = 0 - speculateRound() + autoreleasepool { speculateRound() } if pendingTokens.isEmpty { // speculateRound chose passthrough -- fall through. diff --git a/Libraries/MLXLMCommon/RoPEUtils.swift b/Libraries/MLXLMCommon/RoPEUtils.swift index 1609b0a15..871fa8608 100644 --- a/Libraries/MLXLMCommon/RoPEUtils.swift +++ b/Libraries/MLXLMCommon/RoPEUtils.swift @@ -405,6 +405,40 @@ public class YarnRoPE: Module, OffsetLayer, ArrayOffsetLayer { } +/// RoPE whose base is rescaled once by `alpha^(dim/(dim-2))` (Hunyuan's static NTK scaling). +public class DynamicNTKAlphaRoPE: Module, OffsetLayer, ArrayOffsetLayer { + let dimensions: Int + let traditional: Bool + private let _freqs: MLXArray + + public init( + dimensions: Int, + base: Float = 10000, + scalingAlpha: Float = 1.0, + traditional: Bool = false + ) { + precondition(dimensions % 2 == 0, "Dimensions must be even") + self.dimensions = dimensions + self.traditional = traditional + + let adjustedBase = base * pow(scalingAlpha, Float(dimensions) / Float(dimensions - 2)) + self._freqs = pow( + adjustedBase, + MLXArray(stride(from: 0, to: dimensions, by: 2)).asType(.float32) / Float(dimensions)) + } + + public func callAsFunction(_ x: MLXArray, offset: Int = 0) -> MLXArray { + MLXFast.RoPE( + x, dimensions: dimensions, traditional: traditional, base: nil, scale: 1.0, + offset: offset, freqs: _freqs) + } + + public func callAsFunction(_ x: MLXArray, offset: MLXArray) -> MLXArray { + MLXFast.RoPE( + x, dimensions: dimensions, traditional: traditional, base: nil, scale: 1.0, + offset: offset, freqs: _freqs) + } +} public typealias RoPELayer = OffsetLayer & ArrayOffsetLayer public func initializeRope( diff --git a/Libraries/MLXLMCommon/Tool/Parsers/JSONToolCallParser.swift b/Libraries/MLXLMCommon/Tool/Parsers/JSONToolCallParser.swift index bb924507c..9bc9cf669 100644 --- a/Libraries/MLXLMCommon/Tool/Parsers/JSONToolCallParser.swift +++ b/Libraries/MLXLMCommon/Tool/Parsers/JSONToolCallParser.swift @@ -7,6 +7,7 @@ import Foundation public struct JSONToolCallParser: ToolCallParser, Sendable { public let startTag: String? public let endTag: String? + private let jsonObjectScanner = JSONLeadingObjectScanner(startCharacter: "{") public init(startTag: String, endTag: String) { self.startTag = startTag @@ -29,10 +30,8 @@ public struct JSONToolCallParser: ToolCallParser, Sendable { let jsonStr = text.trimmingCharacters(in: .whitespacesAndNewlines) - guard - let data = jsonStr.data(using: .utf8), - let toolCall = parseToolCall(from: data) - else { return nil } + let toolCall = parseToolCall(from: jsonStr) ?? parseRedundantOuterBraces(from: jsonStr) + guard let toolCall else { return nil } // If tool schemas are provided, only accept calls to declared tools. if let tools, !tools.isEmpty { @@ -53,6 +52,28 @@ public struct JSONToolCallParser: ToolCallParser, Sendable { return toolCall } + /// Some Qwen chat templates emit an EOS-delimited JSON call with a + /// redundant leading brace and two redundant closing braces. + /// Recover only that exact shape and only when the enclosed prefix is one + /// complete, valid tool-call object followed solely by those braces. + private func parseRedundantOuterBraces(from text: String) -> ToolCall? { + guard text.hasPrefix("{{") else { return nil } + let withoutLeadingBrace = String(text.dropFirst()) + guard let split = jsonObjectScanner.splitLeadingObject(from: withoutLeadingBrace) else { + return nil + } + let trailing = split.trailing.trimmingCharacters(in: .whitespacesAndNewlines) + guard trailing == "}}" else { + return nil + } + return parseToolCall(from: split.object) + } + + private func parseToolCall(from text: String) -> ToolCall? { + guard let data = text.data(using: .utf8) else { return nil } + return parseToolCall(from: data) + } + private func parseToolCall(from data: Data) -> ToolCall? { guard var jsonObject = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { diff --git a/Libraries/MLXLMCommon/Tool/ToolCallProcessor.swift b/Libraries/MLXLMCommon/Tool/ToolCallProcessor.swift index ae98d55c1..64efd5215 100644 --- a/Libraries/MLXLMCommon/Tool/ToolCallProcessor.swift +++ b/Libraries/MLXLMCommon/Tool/ToolCallProcessor.swift @@ -24,6 +24,12 @@ import Foundation /// ``` public class ToolCallProcessor { + /// An ordered item emitted while processing generated output. + public enum Output: Sendable, Equatable { + case response(String) + case toolCall(ToolCall) + } + // MARK: - Properties private let format: ToolCallFormat @@ -35,6 +41,8 @@ public class ToolCallProcessor { private var state = State.normal private var toolCallBuffer = "" private var emittedToolCallIDs: Set = [] + private var orderedOutputQueue: [Output] = [] + private var orderedOutputEnabled = false /// The tool calls extracted during processing. public var toolCalls: [ToolCall] = [] @@ -91,11 +99,26 @@ public class ToolCallProcessor { return processTaggedChunk(chunk) } - /// Drain queued tool calls in parse order. + /// Processes a generated chunk and removes its output in source order. /// - /// This is useful for streaming consumers that need deterministic FIFO dispatch - /// of all complete calls discovered while processing a token chunk. - func drainToolCalls() -> [ToolCall] { + /// Tool protocol syntax that does not parse as a call is not emitted as a + /// response. Use this streaming operation when tool calls and response text + /// must retain their relative order. Do not mix this API with `processChunk`, + /// `processEOS`, or `drainToolCalls()` on the same processor instance. + public func processChunkOutputs(_ chunk: String) -> [Output] { + orderedOutputEnabled = true + let outputCount = orderedOutputQueue.count + let visible = processChunk(chunk) + if orderedOutputQueue.count == outputCount, let visible { + recordResponse(sanitizingProtocol: visible) + } + _ = drainToolCalls() + return drainOrderedOutputs() + } + + /// Removes and returns every parsed call in parse order. + /// A second call returns an empty array until more chunks are processed. + public func drainToolCalls() -> [ToolCall] { guard !toolCalls.isEmpty else { return [] } let drained = toolCalls toolCalls.removeAll(keepingCapacity: true) @@ -145,12 +168,36 @@ public class ToolCallProcessor { return returnBufferedText && parsedCalls.isEmpty ? buffered : nil } + /// Finishes processing and removes residual output in source order. + /// + /// This preserves non-tool text following EOS-delimited calls. Do not mix + /// this API with the legacy processing and draining APIs. + public func processEOSOutputs() -> [Output] { + orderedOutputEnabled = true + if format == .mistral, let outputs = processMistralEOSOutputs() { + orderedOutputQueue.removeAll(keepingCapacity: true) + return outputs + } + if format == .lfm2, let outputs = processLFM2EOSOutputs() { + orderedOutputQueue.removeAll(keepingCapacity: true) + return outputs + } + + let outputCount = orderedOutputQueue.count + let visible = processEOS(returnBufferedText: true) + if orderedOutputQueue.count == outputCount, let visible { + recordEOSResidual(visible) + } + _ = drainToolCalls() + return drainOrderedOutputs() + } + // MARK: - Private Methods /// Process chunk for inline formats (no wrapper tags). /// - /// Uses brace counting to detect when output looks like a JSON tool call. - /// While braces are unbalanced the content is buffered (returns `nil`) + /// Uses quote-aware JSON object scanning to detect when output looks like a JSON tool call. + /// While the object is incomplete the content is buffered (returns `nil`) /// so partial JSON is never leaked to the UI. private func processInlineChunk(_ chunk: String) -> String? { switch state { @@ -163,25 +210,30 @@ public class ToolCallProcessor { state = .collectingToolCall if let toolCall = parser.parse(content: toolCallBuffer, tools: tools) { + recordResponse(leading.replacingOccurrences(of: "<|python_tag|>", with: "")) appendToolCall(toolCall) toolCallBuffer = "" state = .normal return leading.isEmpty ? nil : leading } - // Still collecting — check if braces are balanced (would mean parse + // Still collecting — check if the first JSON object is complete (would mean parse // failed on complete JSON, so it's not a tool call) - if jsonBracesBalanced(toolCallBuffer) { + if jsonObjectScanner.splitLeadingObject(from: toolCallBuffer) != nil { state = .normal let buffer = toolCallBuffer toolCallBuffer = "" - return leading + buffer + let response = leading + buffer + recordResponse(sanitizingProtocol: response) + return response } + recordResponse(sanitizingProtocol: leading) return leading.isEmpty ? nil : leading } // No brace seen — pass through as regular text + recordResponse(sanitizingProtocol: chunk) return chunk case .potentialToolCall, .collectingToolCall, .collectingJSONToolCall: @@ -194,11 +246,12 @@ public class ToolCallProcessor { return nil } - // If braces are balanced but parse failed, this isn't a tool call — flush - if jsonBracesBalanced(toolCallBuffer) { + // If the object is complete but parse failed, this isn't a tool call — flush + if jsonObjectScanner.splitLeadingObject(from: toolCallBuffer) != nil { state = .normal let buffer = toolCallBuffer toolCallBuffer = "" + recordResponse(sanitizingProtocol: buffer) return buffer } @@ -207,13 +260,187 @@ public class ToolCallProcessor { } } - /// Check whether open/close braces are balanced in the string. - private func jsonBracesBalanced(_ text: String) -> Bool { + private func appendResponse(_ text: String, to outputs: inout [Output]) { + guard !text.isEmpty else { return } + outputs.append(.response(text)) + } + + private func recordResponse(_ text: String) { + guard orderedOutputEnabled, !text.isEmpty else { return } + orderedOutputQueue.append(.response(text)) + } + + private func recordResponse(sanitizingProtocol text: String) { + recordResponse(stripProtocolSpans(from: text)) + } + + private func recordEOSResidual(_ text: String) { + recordResponse(sanitizeEOSResidual(text)) + } + + private func drainOrderedOutputs() -> [Output] { + let outputs = orderedOutputQueue + orderedOutputQueue.removeAll(keepingCapacity: true) + return outputs + } + + private func stripProtocolSpans(from text: String) -> String { + var result = text + let tags = + [parser.startTag, parser.endTag].compactMap { $0 } + + (format == .llama3 ? ["<|python_tag|>"] : []) + + for tag in tags { + while let range = result.range(of: tag) { + if tag == parser.startTag, + let endTag = parser.endTag, + let end = result.range(of: endTag, range: range.upperBound ..< result.endIndex) + { + result.removeSubrange(range.lowerBound ..< end.upperBound) + } else { + result.removeSubrange(range) + } + } + + guard let first = tag.first else { continue } + var index = result.startIndex + while index < result.endIndex { + guard result[index] == first else { + index = result.index(after: index) + continue + } + let suffix = result[index...] + let matchCount = zip(suffix, tag).prefix { $0 == $1 }.count + guard matchCount >= nearCompleteMatchLength(for: tag) else { + index = result.index(after: index) + continue + } + let markerEnd = + suffix.firstIndex(of: ">") + ?? suffix.firstIndex(of: "]") + let removalEnd = markerEnd.map { result.index(after: $0) } ?? result.endIndex + result.removeSubrange(index ..< removalEnd) + } + } + return result + } + + private func sanitizeEOSResidual(_ text: String) -> String { + guard let startTag = parser.startTag else { + return stripProtocolSpans(from: text) + } + + var searchStart = text.startIndex + while let startRange = text.range(of: startTag, range: searchStart ..< text.endIndex) { + guard + let endTag = parser.endTag, + let endRange = text.range( + of: endTag, range: startRange.upperBound ..< text.endIndex) + else { + return stripProtocolSpans(from: String(text[.. Int { + max(tag.count - 2, 1) + } + + private func processMistralEOSOutputs() -> [Output]? { + guard + state == .collectingToolCall || state == .potentialToolCall + || state == .collectingJSONToolCall, + !toolCallBuffer.isEmpty + else { return nil } + + let startTag = "[TOOL_CALLS]" + let argsTag = "[ARGS]" + var remaining = toolCallBuffer + var outputs: [Output] = [] + + while remaining.hasPrefix(startTag) { + guard let argsRange = remaining.range(of: argsTag) else { break } + let arguments = String(remaining[argsRange.upperBound...]) + guard let split = jsonObjectScanner.splitLeadingObject(from: arguments) else { break } + + let callText = String(remaining[.. [Output]? { + guard + state == .collectingToolCall || state == .potentialToolCall + || state == .collectingJSONToolCall, + !toolCallBuffer.isEmpty, + let startTag = parser.startTag + else { return nil } + + var remaining = toolCallBuffer + var outputs: [Output] = [] + + while let startRange = remaining.range(of: startTag) { + let responsePrefix = String(remaining[.. String.Index? { var depth = 0 - for ch in text { - if ch == "{" { depth += 1 } else if ch == "}" { depth -= 1 } + var stringQuote: Character? + var escaped = false + + for index in text.indices[start...] { + let character = text[index] + if let quote = stringQuote { + if escaped { + escaped = false + } else if character == "\\" { + escaped = true + } else if character == quote { + stringQuote = nil + } + continue + } + switch character { + case "\"", "'": stringQuote = character + case "[": depth += 1 + case "]": + depth -= 1 + if depth == 0 { return index } + default: break + } } - return depth == 0 + return nil } /// Process chunk for tagged formats. @@ -229,11 +456,13 @@ public class ToolCallProcessor { ? taggedStartMode(in: chunk, startChar: startChar) : .none guard startMode != .none || state != .normal else { + recordResponse(chunk) return chunk } toolCallBuffer += chunk var leadingToken: String? + var leadingTokenWasRecorded = false switch state { case .normal: @@ -270,8 +499,12 @@ public class ToolCallProcessor { if partialMatch(buffer: toolCallBuffer, tag: startTag) { if toolCallBuffer.starts(with: startTag) { state = .collectingToolCall + recordResponse(leadingToken ?? "") + leadingTokenWasRecorded = true fallthrough } else { + recordResponse(leadingToken ?? "") + leadingTokenWasRecorded = true return nil } } else { @@ -279,7 +512,9 @@ public class ToolCallProcessor { state = .normal let buffer = toolCallBuffer toolCallBuffer = "" - return (leadingToken ?? "") + buffer + let response = (leadingToken ?? "") + buffer + recordResponse(sanitizingProtocol: response) + return response } case .collectingToolCall: @@ -296,6 +531,9 @@ public class ToolCallProcessor { // Parse the tool call using the parser. if let toolCall = parser.parse(content: bufferedToolCall, tools: tools) { + if !leadingTokenWasRecorded { + recordResponse(leadingToken ?? "") + } appendToolCall(toolCall) state = .normal toolCallBuffer = "" @@ -309,12 +547,16 @@ public class ToolCallProcessor { // Otherwise, return trailing text if non-empty. let trailingText = trailingToken?.isEmpty ?? true ? nil : trailingToken + if let trailingText { recordResponse(trailingText) } return combine(leadingToken, trailingText) } // Preserve unparsed tagged payload as plain text, then continue scanning. state = .normal toolCallBuffer = "" + if !leadingTokenWasRecorded { + recordResponse(leadingToken ?? "") + } if let trailingToken, tokenCouldContainToolStart(trailingToken, startChar: startChar) { @@ -323,6 +565,7 @@ public class ToolCallProcessor { combine(bufferedToolCall, processChunk(trailingToken)) ) } + if let trailingToken { recordResponse(trailingToken) } return combine(leadingToken, combine(bufferedToolCall, trailingToken)) } @@ -347,7 +590,9 @@ public class ToolCallProcessor { state = .normal let buffered = toolCallBuffer toolCallBuffer = "" - return combine(leadingToken, buffered) + let response = (leadingToken ?? "") + buffered + recordResponse(sanitizingProtocol: response) + return response } switch jsonObjectScanner.evaluatePrefix(in: toolCallBuffer) { @@ -357,15 +602,19 @@ public class ToolCallProcessor { toolCallBuffer = "" // vLLM-style recovery: if a tagged tool call exists later, retry tagged parsing. if buffered.contains(startTag) { + recordResponse(leadingToken ?? "") return combine(leadingToken, processChunk(buffered)) } - return combine(leadingToken, buffered) + let response = (leadingToken ?? "") + buffered + recordResponse(sanitizingProtocol: response) + return response case .needsMore, .validObject: break } guard let split = jsonObjectScanner.splitLeadingObject(from: toolCallBuffer) else { // Continue buffering until a complete top-level JSON object is available. + recordResponse(leadingToken ?? "") return leadingToken?.isEmpty ?? true ? nil : leadingToken } @@ -373,6 +622,7 @@ public class ToolCallProcessor { let trailingToken = split.trailing if let toolCall = parser.parse(content: jsonCandidate, tools: tools) { + recordResponse(leadingToken ?? "") appendToolCall(toolCall) state = .normal @@ -386,6 +636,7 @@ public class ToolCallProcessor { return combine(leadingToken, processChunk(trailingToken)) } + recordResponse(trailingToken) return combine(leadingToken, trailingToken) } @@ -394,9 +645,12 @@ public class ToolCallProcessor { state = .normal toolCallBuffer = "" if tokenCouldContainToolStart(trailingToken, startChar: startChar) { + recordResponse((leadingToken ?? "") + jsonCandidate) return combine(leadingToken, combine(jsonCandidate, processChunk(trailingToken))) } - return combine(leadingToken, combine(jsonCandidate, trailingToken)) + let response = (leadingToken ?? "") + jsonCandidate + trailingToken + recordResponse(sanitizingProtocol: response) + return response } private func taggedStartMode( @@ -447,7 +701,11 @@ public class ToolCallProcessor { } private func appendToolCall(_ call: ToolCall) { - toolCalls.append(normalizedToolCall(call)) + let normalized = normalizedToolCall(call) + toolCalls.append(normalized) + if orderedOutputEnabled { + orderedOutputQueue.append(.toolCall(normalized)) + } } private func normalizedToolCall(_ call: ToolCall) -> ToolCall { @@ -501,7 +759,7 @@ public class ToolCallProcessor { } } -private struct JSONLeadingObjectScanner { +struct JSONLeadingObjectScanner { enum PrefixState { case needsMore case validObject diff --git a/Libraries/MLXLMCommon/TurboQuantKVCache.swift b/Libraries/MLXLMCommon/TurboQuantKVCache.swift new file mode 100644 index 000000000..9f8ff9af7 --- /dev/null +++ b/Libraries/MLXLMCommon/TurboQuantKVCache.swift @@ -0,0 +1,1765 @@ +// Copyright © 2026 Apple Inc. +// +// Implements TurboQuant Algorithm 1 (MSE-optimal, arXiv:2504.19874) for KV cache: +// rotation Π + optimal Lloyd-Max scalar codebook quantization on Beta distribution. +// +// Both keys and values use Algorithm 1 (MSE-only at b bits). QJL (Algorithm 2) +// is omitted, Tom Turney's research shows no quality benefit on Apple Silicon, +// and at 4-bit the MSE bias is negligible (paper Section 3.2: bias = 2/π, +// diminishing with bit-width). +// +// Enhancements beyond paper: +// - Norm extraction/restoration: paper assumes ||x||=1; we store norms for arbitrary vectors +// - Norm correction: store ||x|| / ||ỹ|| for dense rotation path (WHT skips, orthogonal preserves norms) +// - WHT rotation option: O(d log d) butterfly in Metal kernel for power-of-2 dims +// - Two-phase architecture: raw prefill → batch compress → compressed decode +// - Pre-rotated queries: q' = Π·q computed once, reused for all cached keys +// +// References: +// - TurboQuant: https://arxiv.org/abs/2504.19874 +// - QJL: https://arxiv.org/abs/2406.03482 +// - PolarQuant: https://arxiv.org/abs/2502.02617 + +import Foundation +import MLX +import MLXNN + +extension DType { + fileprivate var bytesPerElement: Int { + switch self { + case .bfloat16, .float16: return 2 + case .float32: return 4 + case .int32, .uint32: return 4 + case .int16, .uint16: return 2 + case .int8, .uint8: return 1 + default: return 4 + } + } +} + +// MARK: - Codebook Generation + +/// Optimal Lloyd-Max codebook centroids for Beta-distributed coordinates. +/// +/// After random orthogonal rotation, each coordinate of a unit-sphere vector +/// follows Beta distribution f_X(x) ∝ (1-x²)^((d-3)/2) on [-1,1]. +/// For large d, this converges to N(0, 1/d). +enum TurboQuantCodebook { + + // MARK: - Pre-computed Centroids + + /// Pre-computed Lloyd-Max centroids for common (dim, bits) pairs. + /// Generated offline via 100-iteration weighted k-means on 32K-point Beta PDF grid. + /// Avoids ~50ms runtime codebook generation per codec. + /// + /// Additional dims (80, 96) are lazily generated on first access to support + /// Qwen3-4B (head_dim=80) and similar models without startup cost for unused dims. + nonisolated(unsafe) private static var precomputed: [Int: [Int: [Float]]] = [ + 64: [ + 2: [-0.18745463, -0.05649366, 0.05649367, 0.18745449], + 3: [ + -0.26375133, -0.16599470, -0.09368263, -0.03040462, 0.03040464, 0.09368261, + 0.16599482, 0.26375186, + ], + 4: [ + -0.32913971, -0.25096416, -0.19681059, -0.15295772, -0.11478586, -0.08000945, + -0.04726735, -0.01563822, 0.01563822, 0.04723797, 0.07994876, 0.11472529, + 0.15289739, 0.19675052, 0.25090477, 0.32908401, + ], + ], + 128: [ + 2: [-0.13302007, -0.03998107, 0.03998102, 0.13302033], + 3: [ + -0.18828832, -0.11801215, -0.06648001, -0.02156330, 0.02156329, 0.06648005, + 0.11801218, 0.18828897, + ], + 4: [ + -0.23639172, -0.17934021, -0.14023653, -0.10881814, -0.08157559, -0.05678632, + -0.03350975, -0.01108178, 0.01108178, 0.03350975, 0.05678631, 0.08157560, + 0.10881804, 0.14023650, 0.17934017, 0.23639278, + ], + ], + 256: [ + 2: [-0.09420358, -0.02827190, 0.02827190, 0.09420330], + 3: [ + -0.13371243, -0.08361249, -0.04704370, -0.01524900, 0.01524901, 0.04704368, + 0.08361248, 0.13371260, + ], + 4: [ + -0.16852295, -0.12754069, -0.09961203, -0.07719406, -0.05781249, -0.04021866, + -0.02370371, -0.00783269, 0.00783269, 0.02370371, 0.04021868, 0.05781246, + 0.07719407, 0.09961203, 0.12754090, 0.16852276, + ], + ], + ] + + /// Lock for thread-safe lazy population of precomputed centroids. + private static let centroidLock = NSLock() + + /// Dims that should be lazily pre-populated (non-power-of-2 dims used by real models). + /// These fall back to dense rotation path since WHT requires power-of-2, but still + /// benefit from cached centroids to avoid ~50ms runtime k-means per codec init. + /// + /// - 80: Qwen3-4B (head_dim=80) + /// - 96: Various smaller models + private static let lazyDims: [Int] = [80, 96] + private static let lazyBits: [Int] = [2, 3, 4] + + /// Ensure centroids for a given dim are populated. Thread-safe, generates once. + private static func ensureCentroidsPopulated(dim: Int) { + centroidLock.lock() + let exists = precomputed[dim] != nil + centroidLock.unlock() + guard !exists else { return } + + // Generate all bit-widths for this dim + var dimTable: [Int: [Float]] = [:] + for bits in lazyBits { + dimTable[bits] = generateCentroids(dim: dim, bits: bits) + } + + centroidLock.lock() + // Double-check after lock (another thread may have populated) + if precomputed[dim] == nil { + precomputed[dim] = dimTable + } + centroidLock.unlock() + } + + // MARK: - Public API + + /// Codebook centroids for (dim, bits). Uses pre-computed table for common configs, + /// lazily generates and caches for known model dims (80, 96), falls back to + /// runtime generation for truly uncommon ones. + /// + static func codebook(dim: Int, bits: Int) -> MLXArray { + if let dimTable = precomputed[dim], let centroids = dimTable[bits] { + return MLXArray(centroids) + } + // Lazy populate for known model dims (Qwen3-4B dim=80, etc.) + if lazyDims.contains(dim) { + ensureCentroidsPopulated(dim: dim) + if let dimTable = precomputed[dim], let centroids = dimTable[bits] { + return MLXArray(centroids) + } + } + let centroids = generateCentroids(dim: dim, bits: bits) + return MLXArray(centroids) + } + + /// Codebook boundaries (midpoints between adjacent centroids). + static func boundaries(dim: Int, bits: Int) -> MLXArray { + let centroids: [Float] + if let dimTable = precomputed[dim], let cached = dimTable[bits] { + centroids = cached + } else if lazyDims.contains(dim) { + ensureCentroidsPopulated(dim: dim) + if let dimTable = precomputed[dim], let cached = dimTable[bits] { + centroids = cached + } else { + centroids = generateCentroids(dim: dim, bits: bits) + } + } else { + centroids = generateCentroids(dim: dim, bits: bits) + } + var bounds = [Float]() + for i in 0 ..< centroids.count - 1 { + bounds.append((centroids[i] + centroids[i + 1]) / 2.0) + } + return MLXArray(bounds) + } + + /// Generate codebook centroids via weighted k-means on Beta distribution. + /// Used as fallback for uncommon (dim, bits) pairs not in the pre-computed table. + static func generateCentroids(dim: Int, bits: Int) -> [Float] { + let levels = 1 << bits + let gridSize = 32768 + let sigma = 1.0 / sqrt(Float(dim)) + + // Generate grid points and PDF weights + var grid = [Float](repeating: 0, count: gridSize) + var weights = [Float](repeating: 0, count: gridSize) + for i in 0 ..< gridSize { + let x = -1.0 + 2.0 * Float(i) / Float(gridSize - 1) + grid[i] = x + // Beta PDF ∝ (1 - x²)^((d-3)/2), approximated by Gaussian for large d + let exponent = Float(dim - 3) / 2.0 + let w = pow(max(1.0 - x * x, 1e-30), exponent) + weights[i] = w + } + + // Initialize centroids via quantiles + let totalW = weights.reduce(0, +) + var centroids = [Float](repeating: 0, count: levels) + var cumW: Float = 0 + var ci = 0 + for i in 0 ..< gridSize { + cumW += weights[i] + let target = (Float(ci) + 0.5) / Float(levels) * totalW + if cumW >= target && ci < levels { + centroids[ci] = grid[i] + ci += 1 + } + } + // Fill remaining + while ci < levels { + centroids[ci] = centroids[ci - 1] + sigma + ci += 1 + } + + // K-means iterations + for _ in 0 ..< 100 { + var sums = [Float](repeating: 0, count: levels) + var counts = [Float](repeating: 0, count: levels) + for i in 0 ..< gridSize { + var bestJ = 0 + var bestDist = Float.infinity + for j in 0 ..< levels { + let d = abs(grid[i] - centroids[j]) + if d < bestDist { + bestDist = d + bestJ = j + } + } + sums[bestJ] += grid[i] * weights[i] + counts[bestJ] += weights[i] + } + for j in 0 ..< levels { + if counts[j] > 0 { centroids[j] = sums[j] / counts[j] } + } + } + + return centroids.sorted() + } +} + +// MARK: - Rotation Matrix + +/// Random orthogonal rotation matrix generation. +/// +/// TurboQuant Algorithm 1 line 2: Π ∈ ℝ^(d×d) via QR decomposition +/// on random Gaussian matrix. Sign-corrected for determinism. +enum TurboQuantRotation { + + /// Generate a deterministic random orthogonal rotation matrix (dense, d×d). + /// Uses QR decomposition on CPU (not yet GPU-supported in MLX). + static func rotationMatrix(dim: Int, seed: UInt64) -> MLXArray { + let key = MLXRandom.key(seed) + let gaussian = MLXRandom.normal([dim, dim], key: key) + + // QR on CPU (MLX GPU QR not supported yet). MLX's QR runs in f32 + // regardless of input dtype, so orthogonality is good to ~1e-3, // fine for the dense fallback path, whose matched-norm scales absorb + // the residual. Real models (pow2 head dims) take the exact WHT path. + let (q, r) = MLXLinalg.qr(gaussian, stream: .cpu) + let diagR = r.diagonal(stream: .cpu) + let signs = sign(diagR, stream: .cpu) + let result = q * expandedDimensions(signs, axis: 0) + eval(result) + return result + } + + /// Generate a Hadamard matrix of size dim × dim via recursive Kronecker product. + /// Requires dim to be a power of 2. The resulting matrix H satisfies H·H = dim·I. + static func hadamardMatrix(dim: Int) -> MLXArray { + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + // Build recursively: H_1 = [[1]], H_2n = [[H_n, H_n], [H_n, -H_n]] + var h: [[Float]] = [[1.0]] + var size = 1 + while size < dim { + var newH = [[Float]](repeating: [Float](repeating: 0, count: size * 2), count: size * 2) + for i in 0 ..< size { + for j in 0 ..< size { + newH[i][j] = h[i][j] + newH[i][j + size] = h[i][j] + newH[i + size][j] = h[i][j] + newH[i + size][j + size] = -h[i][j] + } + } + h = newH + size *= 2 + } + let flat = h.flatMap { $0 } + let result = MLXArray(flat, [dim, dim]) + eval(result) + return result + } + + /// Generate WHT sign vector: random ±1 per dimension, length d. + /// Used with Walsh-Hadamard Transform for O(d log d) rotation. + static func whtSigns(dim: Int, seed: UInt64) -> MLXArray { + let key = MLXRandom.key(seed) + // Random bits → ±1 + // Generate random ±1 signs using uniform random + let uniform = MLXRandom.uniform(low: 0, high: 1, [dim], key: key) + let signs = MLX.where( + uniform .> MLXArray(Float(0.5)), MLXArray(Float(1.0)), MLXArray(Float(-1.0))) + eval(signs) + return signs + } + + /// Apply WHT butterfly on the last dimension of x. Shape-preserving. + /// Computes unnormalized Walsh-Hadamard transform: H * x along last dim. + private static func whtButterfly(_ x: MLXArray) -> MLXArray { + let dim = x.dim(-1) + let logDim = Int(log2(Double(dim))) + let origShape = x.shape + // Flatten leading dims: [N, dim] + let N = origShape.dropLast().reduce(1, *) + var y = x.reshaped([N, dim]) + + for s in 0 ..< logDim { + let halfBlock = 1 << s + let blockSize = halfBlock << 1 + let numBlocks = dim / blockSize + // Reshape to [N, numBlocks, 2, halfBlock] + y = y.reshaped([N, numBlocks, blockSize]) + let a = y[0..., 0..., .. MLXArray { + let dim = x.dim(-1) + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + let signed = x * signs + let transformed = whtButterfly(signed) + return transformed * MLXArray(Float(1.0 / sqrt(Float(dim)))) + } + + /// Apply SRHT inverse rotation: x = diag(signs) * H * y / sqrt(dim) + /// WHT is self-inverse up to scale. Inverse of (H·D/√d) is (D·H/√d). + static func fwhtInverse(_ y: MLXArray, signs: MLXArray) -> MLXArray { + let dim = y.dim(-1) + precondition(dim > 0 && (dim & (dim - 1)) == 0, "dim must be power of 2") + let transformed = whtButterfly(y) + return transformed * MLXArray(Float(1.0 / sqrt(Float(dim)))) * signs + } +} + +// MARK: - Bit Packing + +/// Efficient bit packing/unpacking for codebook indices. +enum TurboQuantPacking { + + /// Number of uint32 words needed to pack `count` values at `bits` each. + static func packedWidth(count: Int, bits: Int) -> Int { + (count * bits + 31) / 32 + } + + /// Pack b-bit indices into uint32 words. + /// Input: [rows, count] as uint32 (values 0..2^bits-1) + /// Output: [rows, packedWidth] as uint32 + static func packLowBit(_ indices: MLXArray, bits: Int) -> MLXArray { + let count = indices.dim(-1) + let batchShape = Array(indices.shape.dropLast()) + let rows = batchShape.reduce(1, *) + let flat = indices.reshaped([rows, count]) + let pw = packedWidth(count: count, bits: bits) + let mask = UInt32((1 << bits) - 1) + + var wordArrays = [MLXArray]() + for w in 0 ..< pw { + var word = MLXArray.zeros([rows], dtype: .uint32) + for d in 0 ..< count { + let bitOffset = d * bits + let wordIdx = bitOffset / 32 + let offset = bitOffset % 32 + let spill = offset + bits - 32 + + if wordIdx == w { + let shifted = + (flat[0..., d].asType(.uint32) & MLXArray(mask)) << MLXArray(UInt32(offset)) + word = word | shifted + } + if spill > 0 && wordIdx + 1 == w { + let shifted = + (flat[0..., d].asType(.uint32) & MLXArray(mask)) + >> MLXArray(UInt32(bits - spill)) + word = word | shifted + } + } + wordArrays.append(expandedDimensions(word, axis: -1)) + } + let packed = concatenated(wordArrays, axis: -1) // [rows, pw] + return packed.reshaped(batchShape + [pw]) + } + + /// Unpack b-bit indices from uint32 words. + /// Input: [rows, packedWidth] as uint32 + /// Output: [rows, count] as uint32 + static func unpackLowBit(_ packed: MLXArray, bits: Int, count: Int) -> MLXArray { + let shape = packed.shape + let batchShape = Array(shape.dropLast()) + let rows = batchShape.reduce(1, *) + let flat = packed.reshaped([rows, -1]) + let mask = UInt32((1 << bits) - 1) + + var dimArrays = [MLXArray]() + for d in 0 ..< count { + let bitOffset = d * bits + let wordIdx = bitOffset / 32 + let offset = bitOffset % 32 + let spill = offset + bits - 32 + + var value = (flat[0..., wordIdx] >> MLXArray(UInt32(offset))) & MLXArray(mask) + if spill > 0 { + let high = + (flat[0..., wordIdx + 1] << MLXArray(UInt32(bits - spill))) & MLXArray(mask) + value = value | high + } + dimArrays.append(expandedDimensions(value, axis: -1)) + } + let unpacked = concatenated(dimArrays, axis: -1) // [rows, count] + return unpacked.reshaped(batchShape + [count]) + } +} + +// MARK: - MSE Codec (TurboQuant Algorithm 1) + +/// State for MSE-quantized vectors. +struct MSECodecState { + var norms: MLXArray // [B, H, T], original vector L2 norms + var packedIndices: MLXArray // [B, H, T, PackedWidth], packed codebook indices + var tokenCount: Int + let dim: Int + let bits: Int +} + +/// MSE-optimal codec per TurboQuant Algorithm 1. +/// +/// QUANT: y ← Π·x, idx_j ← argmin|y_j - c_k| +/// DEQUANT: ỹ_j ← c_{idx_j}, x̃ ← Π^T · ỹ +class MSECodec { + let dim: Int + let bits: Int + let seed: UInt64 + + /// Codebook centroids [2^bits] + let codebook: MLXArray + /// Codebook boundaries for fast quantization [2^bits - 1] + let boundaries: MLXArray + + /// Whether to use WHT (power-of-2 dim) or dense rotation + let useWHT: Bool + /// WHT sign vector [dim], for O(d log d) Metal encode kernel (power-of-2 dims only) + let whtSigns: MLXArray? + /// Dense rotation matrix Π [dim, dim], used for decode/query rotation (single matmul, fast) + let rotation: MLXArray + /// Π^T, for forward rotation + let rotationT: MLXArray + + init(dim: Int, bits: Int, seed: UInt64 = 42) { + self.dim = dim + self.bits = bits + self.seed = seed + self.codebook = TurboQuantCodebook.codebook(dim: dim, bits: bits) + self.boundaries = TurboQuantCodebook.boundaries(dim: dim, bits: bits) + + // Use WHT for power-of-2 dims (O(d log d) Metal encode kernel) + let isPowerOf2 = dim > 0 && (dim & (dim - 1)) == 0 + self.useWHT = isPowerOf2 && dim <= 1024 + if useWHT { + let signs = TurboQuantRotation.whtSigns(dim: dim, seed: seed) + self.whtSigns = signs + // Build dense WHT rotation matrix for decode/query path (single matmul is faster + // than FWHT butterfly via MLX ops due to graph overhead) + let hadamard = TurboQuantRotation.hadamardMatrix(dim: dim) + let signsDiag = expandedDimensions(signs, axis: 0) + let whtRot = hadamard * signsDiag / MLXArray(Float(sqrt(Float(dim)))) + eval(whtRot) + self.rotation = whtRot + self.rotationT = whtRot.transposed() + } else { + self.whtSigns = nil + self.rotation = TurboQuantRotation.rotationMatrix(dim: dim, seed: seed) + self.rotationT = self.rotation.transposed() + } + } + + /// Encode vectors (Algorithm 1 QUANT) with optional norm correction and optional + /// per-dimension key calibration. + /// Input: [B, H, T, D] + /// Returns MSECodecState with norms and packed indices. + /// + /// WHT path: stores raw norms directly. WHT is an orthogonal transform that preserves + /// norms, so reconstruction_norm ≈ original_norm (within floating point error). + /// Skipping norm correction saves one codebook lookup, one norm computation, and one + /// division per encoded vector. + /// + /// Dense rotation path: stores `original_norm / reconstruction_norm` (norm correction). + /// This compensates for quantization error in the non-orthogonal rotation case. + /// + /// `scale`, when non-nil, is `[dim]` and divides the rotated vector elementwise + /// before quantization: quantize(rotate(k) / scale). This is TurboQuant key + /// calibration (equalize post-rotation per-dimension variance on + /// quantization-sensitive families, see TurboQuantKVCache.computeKeyCalibrationScale). + /// Dividing by scale breaks the orthogonal norm-preservation the WHT raw-norm + /// shortcut relies on, so norm correction always runs when scale is supplied, + /// even on the WHT path. + /// Rotation front half of `encode`: unit-normalize and rotate. + /// + /// Exposed separately so the calibrated fused-kernel dispatch quantizes + /// exactly these values. Boundary quantization is sensitive to the fp + /// reduce order of the rotation; real caches contain degenerate rows + /// (attention sinks) whose rotated coordinates sit near quantization + /// boundaries, so an in-kernel rotation with a different reduce order + /// flips many dimensions of exactly the rows softmax weights most. + /// Sharing these ops makes kernel and MLX encode indices identical. + func rotatedUnit(_ vectors: MLXArray) -> (rotated: MLXArray, norms: MLXArray) { + // Extract norms and normalize (paper assumes unit sphere; we store norms separately) + let norms = sqrt((vectors * vectors).sum(axis: -1)) + let safeNorms = maximum(norms, MLXArray(Float(1e-8))) + let unit = vectors / expandedDimensions(safeNorms, axis: -1) + + // Rotate: y ← Π · x (Algorithm 1 line 5) + return (matmul(unit, rotationT), norms) + } + + func encode(_ vectors: MLXArray, scale: MLXArray? = nil) -> MSECodecState { + let (rotated, norms) = rotatedUnit(vectors) + let calibrated = scale.map { rotated / $0 } ?? rotated + + // Quantize via boundary comparison (fast, no broadcast) + let indices = boundaryQuantize(calibrated) + + let storedNorms: MLXArray + if let scale { + // Calibrated path: undo the /scale divide before measuring the + // reconstruction, per-vector norm is computed AFTER scaling, which + // matches what the flash kernels actually reconstruct (query is + // pre-multiplied by scale, so the dot product resolves in the true + // rotated space, see TurboQuantKVCache.compressedAttention). + let reconstructed = codebook[indices] * scale + let reconNormSq = (reconstructed * reconstructed).sum(axis: -1) + let reconNorms = sqrt(maximum(reconNormSq, MLXArray(Float(1e-16)))) + storedNorms = norms / reconNorms + } else if useWHT { + // WHT fast path: orthogonal transform preserves norms, skip correction + storedNorms = norms + } else { + // Dense rotation path: norm correction compensates for quantization error + let reconstructed = codebook[indices] // [B,H,T,D], quantized approximation in rotated space + let reconNormSq = (reconstructed * reconstructed).sum(axis: -1) + let reconNorms = sqrt(maximum(reconNormSq, MLXArray(Float(1e-16)))) + storedNorms = norms / reconNorms // original_norm / reconstruction_norm + } + + // Pack indices + let packed = TurboQuantPacking.packLowBit(indices, bits: bits) + + return MSECodecState( + norms: storedNorms, + packedIndices: packed, + tokenCount: vectors.dim(2), + dim: dim, + bits: bits + ) + } + + /// Decode from state (Algorithm 1 DEQUANT). + /// `scale`, when non-nil, undoes the calibration divide applied at + /// `encode(_:scale:)`: multiply the codebook lookup elementwise by scale + /// before inverse rotation. Must match the scale passed to encode. + /// Returns: [B, H, T, D] + func decode(_ state: MSECodecState, scale: MLXArray? = nil) -> MLXArray { + // Unpack indices + let indices = TurboQuantPacking.unpackLowBit(state.packedIndices, bits: bits, count: dim) + + // Codebook lookup: ỹ_j ← c_{idx_j} (Algorithm 1 line 9) + var approx = codebook[indices] + if let scale { + approx = approx * scale + } + + // Inverse rotate: x̃ ← Π^T · ỹ (Algorithm 1 line 10) + let unrotated = matmul(approx, rotation) + + // Rescale by stored norms + return expandedDimensions(state.norms, axis: -1) * unrotated + } + + /// Decode in rotated space (skip inverse rotation). + /// Returns centroid values scaled by norm, still in Π-rotated coordinate space. + /// Used with pre-rotated queries for dequant-first SDPA. + func decodeRotated(_ state: MSECodecState) -> MLXArray { + let indices = TurboQuantPacking.unpackLowBit(state.packedIndices, bits: bits, count: dim) + let approx = codebook[indices] + return expandedDimensions(state.norms, axis: -1) * approx + } + + /// Pre-rotate queries for compressed-domain scoring. + /// q' ← Π · q (once per query, reused for all cached keys) + func prepareQueries(_ queries: MLXArray) -> MLXArray { + return matmul(queries, rotationT) + } + + /// Fast quantization via boundary comparison instead of argmin broadcast. + /// boundaries = sorted midpoints between adjacent centroids. + /// Returns uint32 indices in [0, 2^bits - 1]. + func boundaryQuantize(_ rotated: MLXArray) -> MLXArray { + // For each coordinate, count how many boundaries it exceeds + // This gives the codebook index directly + let ndim = rotated.ndim + let expanded = expandedDimensions(rotated, axis: -1) // [..., D, 1] + // Reshape boundaries to broadcast: [1, 1, ..., 1, numBoundaries] + var bShape = [Int](repeating: 1, count: ndim + 1) + bShape[ndim] = boundaries.count + let b = boundaries.reshaped(bShape) + let greater = (expanded .> b).asType(.uint32) // compare against all boundaries + let indices = greater.sum(axis: -1) // count exceeded = index + return indices.asType(.uint32) + } +} + +// MARK: - TurboQuantKVCache + +/// KV cache using TurboQuant compression with two-phase architecture: +/// +/// **Phase 1, Prefill** (L>1): Store raw K/V like KVCacheSimple. Zero overhead. +/// **Transition**: On first decode call, compress entire raw cache in one batch. +/// **Phase 2, Decode** (L=1): Encode 1 new token. Metal kernel scores against +/// all compressed tokens. Zero dequantization. +/// +/// Both keys and values: Algorithm 1 (MSE at b bits, no QJL) +public class TurboQuantKVCache: BaseKVCache { + + public let bits: Int // Legacy: used when keyBits == valueBits + public let keyBits: Int // Bit-width for key compression (0 = raw FP16, no compression) + public let valueBits: Int // Bit-width for value compression (can be lower, V compression is nearly free) + private let seed: UInt64 + + /// Affine-K mode: keys stored as 8-bit affine-quantized (upstream QuantizedKVCache + /// layout, groupSize 64) while values are TurboQuant compressed. The efficient K per + /// the asymmetric-kv paper: near-lossless like FP16 keys at half the bytes. + /// Enabled when keyBits == 8. + public let affineKeyMode: Bool + + /// Affine-K quantization group size. Defaults to 64, but callers that + /// already know the head dimension (e.g. `maybeTurboQuantizeKVCache`) + /// can pass a value pre-resolved via `resolvedKVQuantizationGroupSize` + /// so kernels and quantization agree. Resolved (and possibly adjusted) + /// against the actual head dimension the first time an affine-K encode + /// runs, see `resolveAffineKeyGroupSize(headDim:)`. + public private(set) var keyGroupSize: Int + + // Affine-K storage (wq/scales/biases triplet, grown in steps) + private var affKeyW: MLXArray? + private var affKeyScales: MLXArray? + private var affKeyBiases: MLXArray? + + /// Raw-K mode: keys stay at FP16 (uncompressed) while only values are TurboQuant compressed. + /// This is the single biggest quality finding from TurboQuant+, K precision dominates + /// quality via softmax amplification, V compression is nearly free (linear averaging). + /// Enabled when keyBits == 0. + public let rawKeyMode: Bool + + // Codecs (lazy init) + private var keyMSECodec: MSECodec? // keyBits for keys + private var valueMSECodec: MSECodec? // valueBits for values + + // Phase 1: Raw K/V storage (like KVCacheSimple), used during prefill + private var rawKeys: MLXArray? // [B, H, allocSteps, D] + private var rawValues: MLXArray? // [B, H, allocSteps, D] + private var rawAllocSteps = 0 + + // Phase 2: Compressed storage, used during decode + // MSE-only: packed indices + norms (no QJL, simpler, same quality) + private var keyPackedMSE: MLXArray? + private var keyNorms: MLXArray? + private var valPackedMSE: MLXArray? + private var valNorms: MLXArray? + private var compressedAllocSteps = 0 + + /// Per-dimension key calibration scale s ∈ ℝ^dim, computed once at + /// compressRawCache time from the full prefill K cache and reused for + /// every subsequent encodeNewToken call. `nil` means identity (no + /// calibration): the fast fused Metal encode kernels stay in use for + /// keys. Only the standard (both-K-and-V-quantized) symmetric schemes + /// ever set this, rawKeyMode/affineKeyMode keys are untouched. + private var keyCalibScale: MLXArray? + + /// Whether we've transitioned from raw → compressed + public private(set) var isCompressed = false + + private let step = 256 + + public init( + bits: Int = 4, keyBits: Int? = nil, valueBits: Int? = nil, seed: UInt64 = 42, + keyGroupSize: Int = 64 + ) { + self.bits = bits + self.keyBits = keyBits ?? bits + self.valueBits = valueBits ?? bits + self.rawKeyMode = (keyBits ?? bits) == 0 + self.affineKeyMode = (keyBits ?? bits) == 8 + self.seed = seed + self.keyGroupSize = keyGroupSize + super.init() + } + + /// Resolves (and, if needed, adjusts) `keyGroupSize` against the actual + /// key head dimension the first time an affine-K encode runs. Callers + /// that pre-resolve the group size (`maybeTurboQuantizeKVCache`) hit a + /// no-op here; a directly constructed cache with an incompatible head + /// dimension gets an actionable crash instead of an unhelpful failure + /// deep inside MLX's `quantized(...)`. + private func resolveAffineKeyGroupSize(headDim: Int) { + guard + let resolved = resolvedKVQuantizationGroupSize( + requested: keyGroupSize, keyHeadDim: headDim, valueHeadDim: headDim) + else { + fatalError( + "TurboQuant affine key quantization requires head dimensions divisible by one of the supported group sizes (32, 64, 128). Requested group size: \(keyGroupSize). Key head dim: \(headDim)." + ) + } + keyGroupSize = resolved + } + + override public var isTrimmable: Bool { true } + + // MARK: - Shared Codec Cache + + /// Shared codec cache: all layers with the same (dim, bits, seed) reuse the same codec. + /// Eliminates 56 redundant [128,128] rotation matrices (~7 MB) across 28 layers. + private static let codecLock = NSLock() + nonisolated(unsafe) private static var sharedCodecs: [String: MSECodec] = [:] + + private static func getOrCreateCodec(dim: Int, bits: Int, seed: UInt64) -> MSECodec { + let key = "\(dim)_\(bits)_\(seed)" + codecLock.lock() + if let cached = sharedCodecs[key] { + codecLock.unlock() + return cached + } + codecLock.unlock() + let codec = MSECodec(dim: dim, bits: bits, seed: seed) + codecLock.lock() + sharedCodecs[key] = codec + codecLock.unlock() + return codec + } + + /// Initialize codecs if needed. Uses shared cache to avoid duplicating rotation matrices. + /// In rawKeyMode, key codec is nil, keys stay at FP16, no rotation/quantization needed. + private func ensureCodecs(headDim: Int) { + guard valueMSECodec == nil else { return } + if !rawKeyMode && !affineKeyMode { + keyMSECodec = Self.getOrCreateCodec(dim: headDim, bits: keyBits, seed: seed) + } + valueMSECodec = Self.getOrCreateCodec(dim: headDim, bits: valueBits, seed: seed + 1) + } + + /// Dispatch to WHT or dense fused encode kernel based on codec configuration. + private func fusedEncodeDispatch( + input: MLXArray, codec: MSECodec, headDim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + if codec.useWHT, let signs = codec.whtSigns { + return TurboQuantKernelOps.fusedEncodeWHT( + input: input, whtSigns: signs, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } else { + return TurboQuantKernelOps.fusedEncode( + input: input, rotation: codec.rotation, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } + } + + /// Calibrated key encode (per-dimension key calibration). The rotation + /// runs as MLX ops through `MSECodec.rotatedUnit`, the same ops + /// `MSECodec.encode(_:scale:)` uses, so the kernel quantizes + /// bit-identical rotated values; only the boundary compare, packing, and + /// norm correction run in the fused kernel. Rotating in-kernel instead + /// flips quantization ties on degenerate rows (attention sinks) whose + /// coordinates sit near boundaries, which measurably distorts attention. + private func fusedEncodeDispatchScaled( + input: MLXArray, scale: MLXArray, codec: MSECodec, headDim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let (rotated, norms) = codec.rotatedUnit(input.asType(.float32)) + return TurboQuantKernelOps.fusedQuantizePackScaled( + rotated: rotated, rawNorms: norms, scale: scale, + boundaries: codec.boundaries, codebook: codec.codebook, + bits: codec.bits, dim: headDim + ) + } + + // MARK: - Phase 1: Raw Prefill + + /// Prefill update: store raw K/V, return raw. Zero encoding overhead. + /// Uses KVCacheSimple-style allocation with concatenated growth. + override public func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray) { + let previous = self.offset + + let reset = + if let currentKeys = self.rawKeys, (previous + keys.dim(2)) > currentKeys.dim(2) { + true + } else { + self.rawKeys == nil + } + if reset { + let B = keys.dim(0) + let kvHeads = keys.dim(1) + let kHeadDim = keys.dim(3) + let vHeadDim = values.dim(3) + + let nSteps = (step + keys.dim(2) - 1) / step + let kShape = [B, kvHeads, nSteps * step, kHeadDim] + let vShape = [B, kvHeads, nSteps * step, vHeadDim] + let newK = MLXArray.zeros(kShape, dtype: keys.dtype) + let newV = MLXArray.zeros(vShape, dtype: values.dtype) + + if var currentKeys = self.rawKeys, var currentValues = self.rawValues { + if previous % step != 0 { + currentKeys = currentKeys[.ellipsis, .. MLXArray? { + let dim = allKeys.dim(-1) + let sampleCount = allKeys.dim(0) * allKeys.dim(1) * allKeys.dim(2) + guard sampleCount >= 32 else { return nil } + + let keysF32 = allKeys.asType(.float32) + let norms = sqrt((keysF32 * keysF32).sum(axis: -1)) + let safeNorms = maximum(norms, MLXArray(Float(1e-8))) + let unit = keysF32 / expandedDimensions(safeNorms, axis: -1) + let rotated = matmul(unit, codec.rotationT).reshaped([-1, dim]) + + let perDimStd = std(rotated, axis: 0) + let meanStd = maximum(perDimStd.mean(), MLXArray(Float(1e-8))) + let clamped = clip(perDimStd / meanStd, min: Float(0.25), max: Float(4.0)) + eval(clamped) + + let maxDeviation = (clamped - MLXArray(Float(1.0))).abs().max().item(Float.self) + guard maxDeviation >= 0.1 else { return nil } + return clamped + } + + // MARK: - Transition: Compress Raw Cache + + /// Compress the entire raw K/V cache into packed format in one batch. + /// Called once when transitioning from prefill to decode. + /// + /// In rawKeyMode: only compress values. Keys stay as raw FP16 in rawKeys buffer. + /// This is the highest-quality TurboQuant+ mode, K precision dominates quality. + private func compressRawCache() { + guard !isCompressed, let rk = rawKeys, let rv = rawValues, offset > 0 else { return } + let allKeys = rk[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newKW = MLXArray.zeros([B, H, newAlloc, kw.dim(-1)], dtype: kw.dtype) + let newKS = MLXArray.zeros([B, H, newAlloc, ks.dim(-1)], dtype: ks.dtype) + let newKB = MLXArray.zeros([B, H, newAlloc, kb.dim(-1)], dtype: kb.dtype) + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newKW[.ellipsis, .. rawAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newRK = MLXArray.zeros([B, H, newAlloc, headDim], dtype: keys.dtype) + if prev > 0, let rk = rawKeys { + newRK[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newVP[.ellipsis, .. compressedAllocSteps { + let newAlloc = ((prev + numSteps + step - 1) / step) * step + let newKP = MLXArray.zeros([B, H, newAlloc, kpw], dtype: .uint32) + let newKN = MLXArray.zeros([B, H, newAlloc]) + let newVP = MLXArray.zeros([B, H, newAlloc, vpw], dtype: .uint32) + let newVN = MLXArray.zeros([B, H, newAlloc]) + if prev > 0 { + newKP[.ellipsis, ..1 (prefill chunks): falls back to separate score → softmax → value kernels + /// since causal masking across multiple query positions requires the full score matrix. + /// + /// In rawKeyMode: uses standard matmul for Q*K scoring (raw FP16 keys, no rotation), + /// then compressed-domain Metal kernel for Attn*V (TurboQuant compressed values). + /// TurboFlash is NOT used, it assumes both K and V are packed. + public func compressedAttention( + queries: MLXArray, + keys newKeys: MLXArray, + values newValues: MLXArray, + scale: Float, + mask: MLXFast.ScaledDotProductAttentionMaskMode = .none + ) -> MLXArray { + let headDim = newKeys.dim(-1) + let B = queries.dim(0) + let nQHeads = queries.dim(1) + let nKVHeads = newKeys.dim(1) + let L = queries.dim(2) + let nRepeats = nQHeads / nKVHeads + + // Transition: compress raw cache on first decode call + if !isCompressed { + compressRawCache() + } + + // Phase A: Encode new token + encodeNewToken(keys: newKeys, values: newValues) + + guard let valueMSECodec else { + return queries + } + + let tokenCount = offset + + // Shared V slicing (used by all paths) + let flatValPacked = valPackedMSE![0..., 0..., .. 1 { + q = q.reshaped([B, nKVHeads, nRepeats, L, headDim]) + kwS = expandedDimensions(kwS, axis: -3) + ksS = expandedDimensions(ksS, axis: -3) + kbS = expandedDimensions(kbS, axis: -3) + } + scores = quantizedMM( + q, kwS, scales: ksS, biases: kbS, + transpose: true, groupSize: keyGroupSize, bits: 8) + if nRepeats > 1 { + scores = scores.reshaped([B, nQHeads, L, tokenCount]) + } + } else { + // Q*K scoring: standard matmul with raw FP16 keys + guard let rk = rawKeys else { return queries } + let allKeys = rk[.ellipsis, .. 1 { + let expanded = expandedDimensions(allKeys, axis: 2) + let tiledKeys = MLX.tiled(expanded, repetitions: [1, 1, nRepeats, 1, 1]) + expandedKeys = tiledKeys.reshaped([B, nQHeads, tokenCount, headDim]) + } else { + expandedKeys = allKeys + } + + // scores = Q * K^T * scale → [B, nQHeads, L, T], f32: the f16 + // dot product overflows on outlier-heavy families (Qwen2.5 + // layer-0/last K amax ~200-400 → q·k beyond f16 max). + scores = + matmul( + queries.asType(.float32), + expandedKeys.asType(.float32).transposed(0, 1, 3, 2)) * MLXArray(scale) + } + + // Mask + softmax + switch mask { + case .array(let maskArray): + if maskArray.dtype == .bool { + scores = MLX.where(maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + case .causal: + // Build causal mask manually + let queryOffset = tokenCount - L + let causalMask = MLXArray.tri(L, m: tokenCount, k: queryOffset, type: Bool.self) + let expandedMask = expandedDimensions( + expandedDimensions(causalMask, axis: 0), axis: 0) + scores = MLX.where(expandedMask, scores, MLXArray(-Float.greatestFiniteMagnitude)) + case .arrays(let maskArrays): + if let maskArray = maskArrays.first { + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + } + case .none: break + } + + let attnWeights = softmax(scores, axis: -1) + + // Attn*V: compressed-domain Metal kernel for weighted sum of TurboQuant values + let flatWeights = attnWeights.reshaped([B * nQHeads * L, tokenCount]) + let rotatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: flatWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valueMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.valueBits, dim: headDim, queryChunkLength: L + ) + + // Inverse value rotation (V was encoded in rotated space) + output = matmul( + rotatedOutput.reshaped([B, nQHeads, L, headDim]), + valueMSECodec.rotation + ) + } else { + // ═══ Standard TurboQuant path: both K and V compressed ═══ + guard let keyMSECodec else { return queries } + + // Pre-rotate query for compressed-domain K scoring. When key + // calibration is active, fold s into the query side (elementwise, + // post-rotation): score = rotate(q)*s · rotate(k)/s = rotate(q)·rotate(k). + // The flash/score kernels stay untouched, they only ever see a + // plain dot product between rotatedQueries and the packed K codebook. + var qRot = keyMSECodec.prepareQueries(queries) * MLXArray(scale) + if let keyCalibScale { + qRot = qRot * keyCalibScale + } + let flatQ = qRot.reshaped([B * nQHeads * L, headDim]) + + // K slicing + let flatKeyPacked = keyPackedMSE![0..., 0..., ..1) + let queryOffset = tokenCount - L + output = TurboQuantKernelOps.turboFlashAttentionCausal( + rotatedQueries: flatQ, + keyPacked: flatKeyPacked, keyNorms: flatKeyNorms, + keyCodebook: keyMSECodec.codebook, + valPacked: flatValPacked, valNorms: flatValNorms, + valCodebook: valueMSECodec.codebook, + tokenCount: tokenCount, repeatCount: nRepeats, + keyBits: self.keyBits, valueBits: self.valueBits, dim: headDim, + queryChunkLength: L, queryOffset: queryOffset, + valRotation: valRotation + ).reshaped([B, nQHeads, L, headDim]) + } else { + // Separated path (L>1, non-causal masks) + var scores = TurboQuantKernelOps.mseScore( + rotatedQueries: flatQ, packed: flatKeyPacked, norms: flatKeyNorms, + codebook: keyMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.keyBits, dim: headDim, queryChunkLength: L + ).reshaped([B, nQHeads, L, tokenCount]) + + // Mask + softmax + switch mask { + case .array(let maskArray): + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + case .arrays(let maskArrays): + if let maskArray = maskArrays.first { + if maskArray.dtype == .bool { + scores = MLX.where( + maskArray, scores, MLXArray(-Float.greatestFiniteMagnitude)) + } else { + scores = scores + maskArray + } + } + case .none: break + default: break + } + + let attnWeights = softmax(scores, axis: -1) + + // Metal value kernel + let flatWeights = attnWeights.reshaped([B * nQHeads * L, tokenCount]) + let rotatedOutput = TurboQuantKernelOps.mseWeightedSum( + weights: flatWeights, packed: flatValPacked, norms: flatValNorms, + codebook: valueMSECodec.codebook, tokenCount: tokenCount, + repeatCount: nRepeats, bits: self.valueBits, dim: headDim, queryChunkLength: L + ) + + // Inverse rotation + output = matmul( + rotatedOutput.reshaped([B, nQHeads, L, headDim]), + valueMSECodec.rotation + ) + } + } + + // Kernels emit f32; return in the activation dtype so a turbo layer + // does not promote the whole downstream stream (and every later + // layer's KV cache) to f32. + return output.dtype == queries.dtype ? output : output.asType(queries.dtype) + } + + // MARK: - Memory Reporting + + /// Actual memory footprint: compressed storage (packed indices + norms) for K and V, + /// plus any raw FP16 buffers if still in prefill phase or rawKeyMode. + /// Does NOT include codec overhead (rotation matrices, codebooks) which is shared across layers. + /// In rawKeyMode: rawKeys is always present (FP16 keys), no keyPackedMSE/keyNorms. + public var memoryBytes: Int { + var total = 0 + // Raw FP16 buffers (always present in rawKeyMode for keys, or during prefill) + if let rk = rawKeys { total += rk.shape.reduce(1, *) * rk.dtype.bytesPerElement } + if let rv = rawValues { total += rv.shape.reduce(1, *) * rv.dtype.bytesPerElement } + // Compressed storage (K only present when NOT rawKeyMode) + if let kw = affKeyW { total += kw.shape.reduce(1, *) * kw.dtype.bytesPerElement } + if let ks = affKeyScales { total += ks.shape.reduce(1, *) * ks.dtype.bytesPerElement } + if let kb = affKeyBiases { total += kb.shape.reduce(1, *) * kb.dtype.bytesPerElement } + if let kp = keyPackedMSE { total += kp.shape.reduce(1, *) * kp.dtype.bytesPerElement } + if let kn = keyNorms { total += kn.shape.reduce(1, *) * kn.dtype.bytesPerElement } + if let vp = valPackedMSE { total += vp.shape.reduce(1, *) * vp.dtype.bytesPerElement } + if let vn = valNorms { total += vn.shape.reduce(1, *) * vn.dtype.bytesPerElement } + if let kcs = keyCalibScale { total += kcs.shape.reduce(1, *) * kcs.dtype.bytesPerElement } + return total + } + + // MARK: - State / Trim + + override public var state: [MLXArray] { + get { + if isCompressed { + if affineKeyMode { + // Affine-K compressed: [kW, kScales, kBiases, valPacked, valNorms] + guard let kw = affKeyW, let ks = affKeyScales, let kb = affKeyBiases, + let vpm = valPackedMSE, let vn = valNorms, + offset > 0 + else { return [] } + return [ + kw[0..., 0..., .. 0 + else { return [] } + return [ + rk[0..., 0..., .. 0 + else { return [] } + var arrays = [ + kpm[0..., 0..., .. 0 else { return [] } + return [rk[0..., 0..., ..= 5, + let o = Int(newValue[0]) + else { return } + offset = o + } + } + + @discardableResult + override public func trim(_ n: Int) -> Int { + guard n > 0, offset > 0 else { return 0 } + let trimCount = min(n, offset) + offset -= trimCount + if offset == 0 { + rawKeys = nil + rawValues = nil + rawAllocSteps = 0 + keyPackedMSE = nil + keyNorms = nil + affKeyW = nil + affKeyScales = nil + affKeyBiases = nil + valPackedMSE = nil + valNorms = nil + keyCalibScale = nil + compressedAllocSteps = 0 + isCompressed = false + } + return trimCount + } +} + +// MARK: - kvScheme routing + +/// Resolve a TurboQuant kvScheme string to (keyBits, valueBits). +/// `keyBits == 0` means raw FP16 keys (only values are compressed). +/// Returns nil for non-TurboQuant schemes. +public func resolveTurboScheme(_ scheme: String?) -> (keyBits: Int, valueBits: Int)? { + switch scheme { + // Raw-key schemes: keys stay FP16, values are group-quantized. The + // teacher-forced PPL/KLD gate shows these are near-lossless (see the + // PR validation table); they are the recommended configurations. + case "turbo0v4": return (0, 4) + case "turbo0v3": return (0, 3) + case "turbo0v2": return (0, 2) + // Affine-K asymmetric family (keyBits == 8 → 8-bit affine keys, turbo + // values): the recommended ladder from the asymmetric-kv paper. q8-K is + // near-lossless at half the key bytes; f16-K buys nothing further. + case "turbo8v4": return (8, 4) + case "turbo8v3": return (8, 3) + case "turbo8v2": return (8, 2) + // Key-quantizing (symmetric) schemes: maximum compression. K sensitivity + // varies by model family, validate on your model; the asym family above + // is the recommended starting point. + case "turbo4": return (4, 4) + case "turbo4v2": return (4, 2) + case "turbo3": return (3, 3) + case "turbo2": return (2, 2) + + default: return nil + } +} + +/// Namespace holding the dedup state for the "rotating layers kept fp16" +/// notice. Mirrors the `nonisolated(unsafe) static var` + `NSLock` idiom +/// used elsewhere in this file (e.g. `TurboQuantCodebook.precomputed`). +private enum RotatingSkipNotice { + /// Once per process: the notice is informational and repeating it per + /// generation adds noise without new signal. Tracking instances instead + /// would grow without bound, `newCache(parameters:)` allocates fresh + /// caches for every generation. + nonisolated(unsafe) static var logged = false + static let lock = NSLock() +} + +/// Log, at most once per set of `RotatingKVCache` instances, that a +/// requested TurboQuant kvScheme is leaving sliding-window layers at fp16. +/// +/// TurboQuant only compresses `KVCacheSimple` layers (see +/// `maybeTurboQuantizeKVCache` below); `RotatingKVCache` (Gemma-style +/// sliding-window) layers are already memory-bounded by their window size, +/// and their rotating/eviction storage layout is not compatible with the +/// sequential-append compression path, so they are intentionally left +/// untouched. The notice tells a user which layers kept fp16 rotating +/// caches so a partially engaged scheme is visible. +private func logRotatingKVCacheSkipOnce(cache: [KVCache]) { + let rotatingIndices = cache.indices.filter { cache[$0] is RotatingKVCache } + guard !rotatingIndices.isEmpty else { return } + + RotatingSkipNotice.lock.lock() + defer { RotatingSkipNotice.lock.unlock() } + + guard !RotatingSkipNotice.logged else { return } + RotatingSkipNotice.logged = true + + let indexList = rotatingIndices.map(String.init).joined(separator: ", ") + print( + "[TurboQuant] kvScheme requested KV compression, but layer(s) at index \(indexList) " + + "use RotatingKVCache (sliding-window) and will stay fp16. TurboQuant only " + + "compresses non-rotating (global) KV cache layers." + ) +} + +/// Convert eligible `KVCacheSimple` layers to `TurboQuantKVCache` once their +/// offset passes `quantizedKVStart`, transferring the accumulated KV state. +/// Mamba / wrapper caches are left untouched. `RotatingKVCache` layers +/// (Gemma-style sliding-window) are also left untouched -- see +/// `logRotatingKVCacheSkipOnce` -- since they are already memory-bounded by +/// their window size and their rotating storage layout does not fit the +/// sequential-append compression path. Called from `maybeQuantizeKVCache` +/// when `kvScheme` names a TurboQuant scheme. +public func maybeTurboQuantizeKVCache( + cache: inout [KVCache], + keyBits: Int, + valueBits: Int, + quantizedKVStart: Int +) { + logRotatingKVCacheSkipOnce(cache: cache) + + guard + cache.contains(where: { $0 is KVCacheSimple && $0.offset > quantizedKVStart }) + else { return } + + // Boundary layer protection: the first and last attention layers are + // disproportionately sensitive to KV quantization, keeping 2 on each + // end at FP16 recovers 37-91% of the quality gap at minimal + // compression cost. Non-attention layers (Mamba/wrapper caches) are + // excluded from the rank count so hybrids protect the right layers. + let kvLayerIndices = cache.indices.filter { cache[$0] is KVCacheSimple } + // Only engage when the scheme is actually fragile: turbo-quantized keys + // (keyBits 2-4) or 2-bit values. Raw/affine-8 keys with >=3-bit values + // are near-lossless everywhere and protection just costs compression. + let fragile = (keyBits > 0 && keyBits < 8) || valueBits <= 2 + let boundaryLayers = fragile ? min(2, kvLayerIndices.count / 2) : 0 + let protected = Set( + kvLayerIndices.prefix(boundaryLayers) + kvLayerIndices.suffix(boundaryLayers)) + + for i in 0 ..< cache.count { + guard let simple = cache[i] as? KVCacheSimple, cache[i].offset > quantizedKVStart, + !(cache[i] is TurboQuantKVCache) + else { continue } + + let state = cache[i].innerState() + let headDims: (key: Int, value: Int)? = + state.count >= 2 ? (state[0].dim(3), state[1].dim(3)) : nil + + if protected.contains(i) { + // Boundary layers use 8-bit affine instead of the turbo scheme: + // near-lossless protection for the quantization-sensitive first + // and last layers at a quarter of the fp16 cost. Skip the + // conversion (leave fp16) rather than crash when neither head + // dimension divides one of the supported affine group sizes. + guard + let headDims, + resolvedKVQuantizationGroupSize( + requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) != nil + else { continue } + cache[i] = simple.toQuantized(groupSize: 64, bits: 8) + continue + } + + // Affine-K mode (keyBits == 8) quantizes keys in groups, resolve the + // group size against this layer's head dimension up front so the + // TurboQuantKVCache instance never has to fall back at first encode. + // Skip layers whose head dimension divides none of the supported + // sizes rather than constructing a cache that would only crash on + // its first compress. Other modes (raw / symmetric turbo) don't + // group-quantize keys, so they carry the unused default unchanged. + var resolvedKeyGroupSize = 64 + if keyBits == 8 { + guard + let headDims, + let resolved = resolvedKVQuantizationGroupSize( + requested: 64, keyHeadDim: headDims.key, valueHeadDim: headDims.value) + else { continue } + resolvedKeyGroupSize = resolved + } + + let turbo = TurboQuantKVCache( + bits: max(keyBits, valueBits), keyBits: keyBits, valueBits: valueBits, + keyGroupSize: resolvedKeyGroupSize) + // Transfer existing KV data, trimmed to the live offset (the simple + // cache over-allocates in steps). + let offset = cache[i].offset + if state.count >= 2, offset > 0 { + let keys = state[0][.ellipsis, ..> shift); + + // Handle bits that spill across uint32 word boundary + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + value |= (packed_ptr[word_idx + 1] << ((uint)Bits - (uint)spill)); + } + value &= MASK; + + // Codebook lookup + accumulate dot product + acc += q_ptr[d] * cb[value]; + } + + // SIMD reduction across 32 lanes + acc = simd_sum(acc); + + // Lane 0 writes final score (scaled by stored norm) + if (thread_index_in_simdgroup == 0) { + scores[q_idx * token_count + k_idx] = acc * norm_val; + } + """ + + /// Fused encode kernel: norm + rotate + quantize + pack + norm correction in ONE dispatch. + /// + /// For each input vector [D]: + /// 1. Compute L2 norm (SIMD reduction) + /// 2. Normalize to unit vector + /// 3. Rotate: y = Π · x_unit (shared memory matmul) + /// 4. Quantize: find codebook index via boundary comparison + /// 5. Pack bits into uint32 words (atomic OR) + /// 6. Norm correction: compute reconstruction norm, store original_norm / recon_norm + /// + /// Norm correction compensates for quantization error so that + /// centroid[idx] * corrected_norm more accurately reconstructs the original vector. + /// This is why TurboQuant beats q8_0 on perplexity. + /// + /// Grid: (Dim, numRows, 1), one threadgroup per vector + /// Threadgroup: (Dim, 1, 1), all D threads cooperate + /// + /// Template params: Bits, Dim, PackedWidth, NumBoundaries (= 2^Bits - 1) + static let fusedEncodeSource = """ + constexpr uint LEVELS = 1u << Bits; + + uint d = thread_position_in_threadgroup.x; // dimension index (0..Dim-1) + uint row = thread_position_in_grid.y; // vector index (B*H*T) + + // --- Step 1: Load input value --- + float val = input[row * Dim + d]; + + // --- Step 2: Compute L2 norm (SIMD reduction) --- + float sq = val * val; + float norm_sq = simd_sum(sq); + // For Dim > 32, need threadgroup reduction + threadgroup float shared_norm[32]; // up to 32 SIMD groups (dim <= 1024) + uint sg_id = d / 32; + if (d % 32 == 0) { + shared_norm[sg_id] = norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_norm_sq = 0; + uint num_groups = (Dim + 31) / 32; + for (uint i = 0; i < num_groups; i++) { + total_norm_sq += shared_norm[i]; + } + float norm_val = sqrt(total_norm_sq); + float inv_norm = (norm_val > 1e-8f) ? (1.0f / norm_val) : 0.0f; + + // --- Step 3: Normalize --- + float unit_val = val * inv_norm; + + // --- Step 4: Rotate (y = Π · x_unit) via shared memory matmul --- + // Each thread d computes: y[d] = Σ_j rotation[d * Dim + j] * x_unit[j] + threadgroup float shared_unit[1024]; // max Dim = 1024 + shared_unit[d] = unit_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float rotated = 0.0f; + for (uint j = 0; j < Dim; j++) { + rotated += rotation[d * Dim + j] * shared_unit[j]; + } + + // --- Step 5: Quantize via branchless boundary comparison --- + // V2.1 optimization: use arithmetic sum of comparisons instead of branching. + // Metal compiles (rotated > boundaries[b]) to a predicated 0/1, summing these + // is branchless and avoids SIMD lane divergence. + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(rotated > boundaries[b]); + } + + // --- Step 6: Pack bits into uint32 word (atomic OR) --- + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + // Pack bits, use threadgroup shared memory to avoid atomic contention + // Each thread writes its index bits to shared, then thread 0 per word writes output + threadgroup uint shared_packed[64]; // max PackedWidth = 64 words + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Each dimension contributes its bits via atomic OR on threadgroup memory + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Write packed words to output (one thread per word) + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // --- Step 7: Norm correction --- + // Compute reconstruction norm: ||codebook[idx]||₂ for the quantized unit vector. + // Store corrected_norm = original_norm / recon_norm so that + // decode(centroid[idx] * corrected_norm) better approximates the original vector. + float centroid_val = codebook[idx]; + float recon_sq = centroid_val * centroid_val; + float recon_norm_sq = simd_sum(recon_sq); + // Threadgroup reduction for Dim > 32 + if (d % 32 == 0) { + shared_norm[sg_id] = recon_norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_recon_sq = 0; + for (uint i = 0; i < num_groups; i++) { + total_recon_sq += shared_norm[i]; + } + float recon_norm = sqrt(total_recon_sq); + float corrected_norm = (recon_norm > 1e-8f) ? (norm_val / recon_norm) : norm_val; + + if (d == 0) { + norms_out[row] = corrected_norm; + } + """ + + /// Fused WHT encode kernel: norm + WHT rotation + quantize + pack (NO norm correction). + /// + /// Same as fusedEncodeSource but replaces dense O(d²) matmul with Fast Walsh-Hadamard + /// Transform O(d log d) butterfly + random sign flip. 18× fewer ops for dim=128. + /// + /// Stores RAW norms (no reconstruction-norm correction), a measured choice, + /// not an omission: for values consumed by attention's weighted sums, raw + /// norms beat both matched-norm (attention cos 0.976 vs ≥0.980) and + /// projection-optimal (0.9798) rescaling at 4-bit; per-vector rescales bias + /// the decoded vectors toward zero while raw-norm errors stay unbiased and + /// average out across tokens. The dense fallback path keeps matched-norm + /// for canonical-formulation parity. + /// + /// WHT forward rotation: y = WHT(signs * x_unit) / sqrt(Dim) + /// The butterfly pattern: for each stage s in 0.. 1e-8f) ? (1.0f / norm_val) : 0.0f; + + // --- Step 3: Normalize + sign flip (fused) --- + // V2.1 optimization: pre-compute inv_norm * sign to eliminate one multiply per element. + // Instead of: unit_val = val * inv_norm; wht_val = sign * unit_val (2 muls) + // We do: wht_val = val * (inv_norm * sign) (1 mul + 1 FMA-friendly product) + float inv_norm_sign = inv_norm * wht_signs[d]; + float wht_val = val * inv_norm_sign; + + // --- Step 4: WHT rotation via cooperative SIMD shuffle --- + // V2.1 optimization: use simd_shuffle_xor for intra-SIMD butterfly stages + // (register-to-register, no shared memory or barriers needed for first 5 stages) + + // Phase 1: Intra-SIMD butterfly via simd_shuffle_xor (stages 0..min(LogDim,5)-1) + // Each stage s XORs lane indices at distance 2^s, effectively free on Apple GPU + uint log_dim_u = uint(LogDim); + uint simd_stages = min(log_dim_u, 5u); // 5 stages covers 32 lanes (2^5 = 32) + uint lane_in_simd = d % 32; + for (uint s = 0; s < simd_stages; s++) { + uint step = 1u << s; + float other = simd_shuffle_xor(wht_val, step); + wht_val = (lane_in_simd & step) ? (other - wht_val) : (other + wht_val); + } + + // Phase 2: Cross-SIMD-group butterfly via shared memory (stages 5..LogDim-1) + // Only needed when Dim > 32, these stages cross SIMD group boundaries + threadgroup float shared_buf[1024]; // max Dim = 1024 + if (log_dim_u > 5u) { + shared_buf[d] = wht_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint s = simd_stages; s < log_dim_u; s++) { + uint half_block = 1u << s; + uint block_size = half_block << 1; + uint block_id = d / block_size; + uint pos_in_block = d % block_size; + + float a, b; + if (pos_in_block < half_block) { + a = shared_buf[block_id * block_size + pos_in_block]; + b = shared_buf[block_id * block_size + pos_in_block + half_block]; + shared_buf[d] = a + b; + } else { + a = shared_buf[block_id * block_size + pos_in_block - half_block]; + b = shared_buf[block_id * block_size + pos_in_block]; + shared_buf[d] = a - b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + wht_val = shared_buf[d]; + } + + // Normalize: WHT has scale factor sqrt(Dim) + float inv_sqrt_dim = 1.0f / sqrt((float)Dim); + float rotated = wht_val * inv_sqrt_dim; + + // --- Step 5: Quantize via branchless boundary comparison --- + // V2.1 optimization: arithmetic sum avoids SIMD lane divergence + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(rotated > boundaries[b]); + } + + // --- Step 6: Pack bits into uint32 word (atomic OR) --- + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + threadgroup uint shared_packed[64]; + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // --- Step 7: Store raw norm (WHT is orthogonal, no norm correction needed) --- + // WHT preserves norms: ||WHT(x)||₂ = ||x||₂. Reconstruction norm ≈ original norm, + // so the correction ratio ≈ 1.0. Skipping saves codebook lookup + norm + division. + if (d == 0) { + norms_out[row] = norm_val; + } + """ + + /// Quantize, pack, and norm-correct pre-rotated calibrated vectors. + /// + /// The rotation stays in MLX ops (MSECodec.rotatedUnit) so kernel and + /// codec quantize bit-identical values; this kernel only does the parts + /// that are slow as MLX ops: the boundary compare (which broadcasts a + /// [rows, dim, levels] intermediate) and low-bit packing. The stored + /// norm is norms_in / ||codebook[idx] * scale|| per row, matching + /// MSECodec.encode(_:scale:). + /// + /// Template params: Bits, Dim, PackedWidth. + static let fusedQuantizePackScaledSource = """ + constexpr uint LEVELS = 1u << Bits; + + uint d = thread_position_in_threadgroup.x; // dimension index (0..Dim-1) + uint row = thread_position_in_grid.y; // vector index (B*H*T) + + float calibrated = rotated[row * Dim + d] / scale[d]; + uint idx = 0; + for (uint b = 0; b < LEVELS - 1; b++) { + idx += (uint)(calibrated > boundaries[b]); + } + + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint masked = idx & ((1u << Bits) - 1u); + + threadgroup uint shared_packed[64]; + if (d < PackedWidth) shared_packed[d] = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint primary_val = masked << shift; + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx], primary_val, memory_order_relaxed); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + uint spill_val = masked >> ((uint)Bits - (uint)spill); + atomic_fetch_or_explicit((threadgroup atomic_uint*)&shared_packed[word_idx + 1], spill_val, memory_order_relaxed); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (d < PackedWidth) { + packed_out[row * PackedWidth + d] = shared_packed[d]; + } + + // Norm correction against the calibrated reconstruction. + float centroid_val = codebook[idx] * scale[d]; + float recon_sq = centroid_val * centroid_val; + float recon_norm_sq = simd_sum(recon_sq); + threadgroup float shared_norm[32]; + uint sg_id = d / 32; + if (d % 32 == 0) { + shared_norm[sg_id] = recon_norm_sq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_recon_sq = 0; + uint num_groups = (Dim + 31) / 32; + for (uint i = 0; i < num_groups; i++) { + total_recon_sq += shared_norm[i]; + } + float recon_norm = sqrt(total_recon_sq); + + if (d == 0) { + float raw_norm = norms_in[row]; + norms_out[row] = (recon_norm > 1e-8f) ? (raw_norm / recon_norm) : raw_norm; + } + """ + + /// TurboFlashAttention Pass 1: Per-block partial attention with online softmax. + /// + /// Parallelizes across both query heads AND token blocks. Each SIMD group (32 lanes) + /// handles one (query, block) pair, producing partial online softmax state (m, l, o[D]). + /// Pass 2 merges partials across blocks. + /// + /// Grid: (32, totalQueries, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks + static let turboFlashPass1Source = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load key codebook into registers + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K dot product --- + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + float k_norm = key_norms[kv_idx * token_count + t]; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + + dot_partial += q_vals[i] * key_cb[k_value]; + } + + float score = simd_sum(dot_partial) * k_norm; + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// TurboFlashAttention Pass 1 (Causal): Per-block partial attention with causal masking. + /// + /// Same as turboFlashPass1Source but supports L>1 query chunks with causal masking. + /// Each query position q_within_L only attends to tokens where t <= q_offset + q_within_L. + /// Blocks that are entirely future-masked exit early. + /// + /// Grid: (32, totalQueries, numBlocks) where totalQueries = B * nQHeads * L + /// Threadgroup: (32, 1, 1) + /// + /// Additional template params: L (query chunk length), q_offset (absolute offset of first query) + /// Pass-1 variant: raw K scoring (KT template, matches the stored K + /// cache dtype), turbo-V accumulation. Same SIMD-parallel online-softmax + /// structure as the packed-K flash kernel, so the asymmetric family + /// gets the same decode throughput. + static let turboFlashPass1RawKSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K, K read raw f16 (no unpack, no rotation) --- + const device KT* k_raw_ptr = (const device KT*)k_raw + kv_idx * token_count * Dim + t * Dim; + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + dot_partial += q_vals[i] * float(k_raw_ptr[d]); + } + float score = simd_sum(dot_partial); + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// Pass-1 variant: 8-bit affine K scoring (inline dequant, scale/bias in + /// their native dtype via the KScaleT/KBiasT templates), turbo-V. + static let turboFlashPass1AffineKSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + uint kv_idx = q_idx / repeat_count; // map to KV head (GQA; L=1 only) + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K, K dequantized inline from 8-bit affine --- + const device uint* k_w_ptr = k_weights + (kv_idx * token_count + t) * (Dim / 4); + const device KScaleT* k_s_ptr = + (const device KScaleT*)k_scales + (kv_idx * token_count + t) * (Dim / KGroup); + const device KBiasT* k_b_ptr = + (const device KBiasT*)k_biases + (kv_idx * token_count + t) * (Dim / KGroup); + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + uint g = d / KGroup; + uint w8 = (k_w_ptr[d / 4] >> ((d % 4) * 8)) & 0xFFu; + dot_partial += q_vals[i] * (float(k_s_ptr[g]) * float(w8) + float(k_b_ptr[g])); + } + float score = simd_sum(dot_partial); + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + static let turboFlashPass1CausalSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + const uint L = params[3]; + const uint q_offset = params[4]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint q_idx = thread_position_in_grid.y; // query index (B*nQHeads*L) + uint block_idx = thread_position_in_grid.z; // token block index + + // For L>1, queries are laid out as [B * nQHeads * L, D] from reshape of [B, nQHeads, L, D]. + // q_idx = b * (nQHeads * L) + h * L + l + // We need: l (position within chunk) and kv_head (for GQA mapping). + uint q_within_L = q_idx % L; + uint q_head_idx = q_idx / L; // index into [B * nQHeads] + uint kv_idx = q_head_idx / repeat_count; // map to KV head (GQA) + + // Causal boundary: this query can attend to tokens 0..q_abs (inclusive) + uint q_abs = q_offset + q_within_L; + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Early exit: entire block is future-masked + if (t_start > q_abs) { + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = 0.0f; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = -INFINITY; + l_partials[ml_idx] = 0.0f; + } + return; + } + + // Clamp t_end to causal boundary + if (t_end > q_abs + 1) t_end = q_abs + 1; + + // Load key codebook into registers + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for this lane's dimensions + float q_vals[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + + // Online softmax state for this block + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + // Process tokens in this block (up to causal boundary) + for (uint t = t_start; t < t_end; t++) { + // --- Score: Q×K dot product --- + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + float k_norm = key_norms[kv_idx * token_count + t]; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + + dot_partial += q_vals[i] * key_cb[k_value]; + } + + float score = simd_sum(dot_partial) * k_norm; + + // --- Online softmax update + V accumulation --- + float new_m = max(m, score); + float exp_diff = exp(m - new_m); + float exp_score = exp(score - new_m); + + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) break; + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + + o[i] = o[i] * exp_diff + exp_score * (val_cb[v_value] * v_norm); + } + + l = l * exp_diff + exp_score; + m = new_m; + } + + // Write partial results: o[D], m, l + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o[i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m; + l_partials[ml_idx] = l; + } + """ + + /// TurboFlashAttention Pass 1 NR0=2: Multi-row amortized KV dequant. + /// + /// Ported from llama.cpp V2.1 fused decode kernel concept. Each SIMD group processes + /// NR0=2 queries against one KV block simultaneously. The key win: packed index unpacking + /// + codebook lookup for K and V is done ONCE and reused across both queries. + /// + /// Register budget per thread (NR0=2, DIMS_PER_LANE=4 for dim=128): + /// - 2 × DIMS_PER_LANE q_vals = 8 floats (query data) + /// - 2 × 1 m/l = 4 floats (online softmax state) + /// - 2 × DIMS_PER_LANE o = 8 floats (value accumulators) + /// - codebook regs shared = KEY_LEVELS + VAL_LEVELS floats + /// Total: ~24 extra floats vs NR0=1. Well within Apple GPU register file. + /// + /// Zero threadgroup memory: all score computation + softmax + V accumulation happen + /// in SIMD registers. No shared memory needed in pass 1 (same as NR0=1 baseline). + /// Note: pass 2 still needs threadgroup memory for dim>32 (cross-SIMD gather for rotation). + /// See turboFlashPass2FusedRotSource comments for details. + /// + /// Grid: (32, totalQueries/NR0, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks, NR0 + static let turboFlashPass1NR0Source = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; // SIMD lane (0-31) + uint query_group = thread_position_in_grid.y; // which group of NR0 queries + uint block_idx = thread_position_in_grid.z; // which KV block + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Load key codebook into registers (shared across all NR0 queries) + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) { + key_cb[i] = key_codebook[i]; + } + + // Load value codebook into registers (shared across all NR0 queries) + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) { + val_cb[i] = val_codebook[i]; + } + + // Load query values for ALL NR0 rows, each row's dims interleaved in registers + float q_vals[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[r * DIMS_PER_LANE + i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + } + + // Per-query KV head mapping (for GQA, each query may map to different KV head) + uint kv_indices[NR0]; + for (uint r = 0; r < NR0; r++) { + kv_indices[r] = (query_group * NR0 + r) / repeat_count; + } + + // Online softmax state, NR0 independent streams, all in registers + float m_state[NR0]; + float l_state[NR0]; + float o_state[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + m_state[r] = -INFINITY; + l_state[r] = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = 0.0f; + } + } + + // Process tokens in this block, KV dequant done ONCE, reused across NR0 queries + for (uint t = t_start; t < t_end; t++) { + // --- Dequant K for this token ONCE (amortized across NR0 queries) --- + // Each lane unpacks its dims' codebook values into registers + float k_decoded[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { k_decoded[i] = 0.0f; continue; } + + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + + // All NR0 rows share one KV head by construction: the Swift + // dispatcher only selects this kernel when repeat_count is a + // multiple of NR0, so an aligned group can never span a + // KV-head boundary and kv_indices[0] is exact. + const device uint32_t* k_packed_ptr = key_packed + kv_indices[0] * token_count * KeyPackedWidth + t * KeyPackedWidth; + + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + k_decoded[i] = key_cb[k_value]; + } + float k_norm = key_norms[kv_indices[0] * token_count + t]; + + // --- Dequant V for this token ONCE --- + float v_decoded[DIMS_PER_LANE]; + const device uint32_t* v_packed_ptr = val_packed + kv_indices[0] * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_indices[0] * token_count + t]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { v_decoded[i] = 0.0f; continue; } + + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + v_decoded[i] = val_cb[v_value] * v_norm; + } + + // --- Score + softmax + V accumulate for each of NR0 queries --- + // K/V dequant above is the expensive part, this loop is cheap ALU + for (uint r = 0; r < NR0; r++) { + // Dot product: q[r] · k (both already in registers) + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + dot_partial += q_vals[r * DIMS_PER_LANE + i] * k_decoded[i]; + } + float score = simd_sum(dot_partial) * k_norm; + + // Online softmax update + float new_m = max(m_state[r], score); + float exp_diff = exp(m_state[r] - new_m); + float exp_score = exp(score - new_m); + + // V accumulation (reusing pre-decoded values) + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = o_state[r * DIMS_PER_LANE + i] * exp_diff + exp_score * v_decoded[i]; + } + + l_state[r] = l_state[r] * exp_diff + exp_score; + m_state[r] = new_m; + } + } + + // Write partial results for all NR0 queries + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o_partials[partial_base + d] = o_state[r * DIMS_PER_LANE + i]; + } + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m_state[r]; + l_partials[ml_idx] = l_state[r]; + } + } + """ + + /// TurboFlashAttention Pass 1 NR0 Causal: Multi-row amortized KV dequant with causal masking. + /// + /// Same as turboFlashPass1NR0Source but each query within the NR0 group has its own + /// causal boundary. For L>1 prefill, q_within_L differs per row so each row may attend + /// to a different number of tokens. We compute the conservative (minimum) causal boundary + /// across the NR0 group for the shared K/V dequant, then mask per-row in the score loop. + /// + /// Grid: (32, totalQueries/NR0, numBlocks) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: KeyBits, ValueBits, Dim, KeyPackedWidth, ValuePackedWidth, + /// BlockSize, token_count, repeat_count, num_blocks, NR0, L, q_offset + static let turboFlashPass1NR0CausalSource = """ + // Per-call values arrive via the `params` input buffer, NOT as template + // args: template values are baked into the kernel name, so a varying + // token_count would JIT-compile and permanently cache one Metal + // library per generated token (unbounded host-memory growth + a + // shader compile every step). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint num_blocks = params[2]; + const uint L = params[3]; + const uint q_offset = params[4]; + constexpr uint KEY_MASK = (1u << KeyBits) - 1u; + constexpr uint KEY_LEVELS = 1u << KeyBits; + constexpr uint VAL_MASK = (1u << ValueBits) - 1u; + constexpr uint VAL_LEVELS = 1u << ValueBits; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint query_group = thread_position_in_grid.y; + uint block_idx = thread_position_in_grid.z; + + // Token range for this block + uint t_start = block_idx * BlockSize; + uint t_end = t_start + BlockSize; + if (t_end > (uint)token_count) t_end = (uint)token_count; + + // Compute per-row causal boundaries and find the maximum (most permissive) + // for the shared token loop. Per-row masking happens inside the score loop. + uint q_abs[NR0]; + uint max_q_abs = 0; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint q_within_L = q_idx % L; + q_abs[r] = q_offset + q_within_L; + if (q_abs[r] > max_q_abs) max_q_abs = q_abs[r]; + } + + // Early exit: entire block is future-masked for ALL NR0 queries + if (t_start > max_q_abs) { + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = 0.0f; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = -INFINITY; + l_partials[ml_idx] = 0.0f; + } + } + return; + } + + // Clamp t_end to the most permissive causal boundary + if (t_end > max_q_abs + 1) t_end = max_q_abs + 1; + + // Load codebooks (shared across all NR0 queries) + float key_cb[KEY_LEVELS]; + for (uint i = 0; i < KEY_LEVELS; i++) key_cb[i] = key_codebook[i]; + float val_cb[VAL_LEVELS]; + for (uint i = 0; i < VAL_LEVELS; i++) val_cb[i] = val_codebook[i]; + + // Load query values for all NR0 rows + float q_vals[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + q_vals[r * DIMS_PER_LANE + i] = (d < Dim) ? q_rot[q_idx * Dim + d] : 0.0f; + } + } + + // KV head mapping (use first query's head, same assumption as non-causal NR0) + uint q_head_idx_0 = (query_group * NR0) / L; + uint kv_idx = q_head_idx_0 / repeat_count; + + // Online softmax state, NR0 independent streams + float m_state[NR0]; + float l_state[NR0]; + float o_state[NR0 * DIMS_PER_LANE]; + for (uint r = 0; r < NR0; r++) { + m_state[r] = -INFINITY; + l_state[r] = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) o_state[r * DIMS_PER_LANE + i] = 0.0f; + } + + // Process tokens, KV dequant once, score per-row with causal mask + for (uint t = t_start; t < t_end; t++) { + // Dequant K once + float k_decoded[DIMS_PER_LANE]; + const device uint32_t* k_packed_ptr = key_packed + kv_idx * token_count * KeyPackedWidth + t * KeyPackedWidth; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { k_decoded[i] = 0.0f; continue; } + uint k_bit_offset = d * KeyBits; + uint k_word_idx = k_bit_offset / 32; + uint k_shift = k_bit_offset % 32; + uint k_value = (k_packed_ptr[k_word_idx] >> k_shift); + int k_spill = (int)k_shift + (int)KeyBits - 32; + if (k_spill > 0) { + k_value |= (k_packed_ptr[k_word_idx + 1] << ((uint)KeyBits - (uint)k_spill)); + } + k_value &= KEY_MASK; + k_decoded[i] = key_cb[k_value]; + } + float k_norm = key_norms[kv_idx * token_count + t]; + + // Dequant V once + float v_decoded[DIMS_PER_LANE]; + const device uint32_t* v_packed_ptr = val_packed + kv_idx * token_count * ValuePackedWidth + t * ValuePackedWidth; + float v_norm = val_norms[kv_idx * token_count + t]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d >= Dim) { v_decoded[i] = 0.0f; continue; } + uint v_bit_offset = d * ValueBits; + uint v_word_idx = v_bit_offset / 32; + uint v_shift = v_bit_offset % 32; + uint v_value = (v_packed_ptr[v_word_idx] >> v_shift); + int v_spill = (int)v_shift + (int)ValueBits - 32; + if (v_spill > 0) { + v_value |= (v_packed_ptr[v_word_idx + 1] << ((uint)ValueBits - (uint)v_spill)); + } + v_value &= VAL_MASK; + v_decoded[i] = val_cb[v_value] * v_norm; + } + + // Score + softmax + V for each query row (with per-row causal mask) + for (uint r = 0; r < NR0; r++) { + // Per-row causal: skip if this token is future for this specific query + if (t > q_abs[r]) continue; + + float dot_partial = 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + dot_partial += q_vals[r * DIMS_PER_LANE + i] * k_decoded[i]; + } + float score = simd_sum(dot_partial) * k_norm; + + float new_m = max(m_state[r], score); + float exp_diff = exp(m_state[r] - new_m); + float exp_score = exp(score - new_m); + + for (uint i = 0; i < DIMS_PER_LANE; i++) { + o_state[r * DIMS_PER_LANE + i] = o_state[r * DIMS_PER_LANE + i] * exp_diff + exp_score * v_decoded[i]; + } + l_state[r] = l_state[r] * exp_diff + exp_score; + m_state[r] = new_m; + } + } + + // Write partial results for all NR0 queries + for (uint r = 0; r < NR0; r++) { + uint q_idx = query_group * NR0 + r; + uint partial_base = (q_idx * num_blocks + block_idx) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) o_partials[partial_base + d] = o_state[r * DIMS_PER_LANE + i]; + } + if (lane == 0) { + uint ml_idx = q_idx * num_blocks + block_idx; + m_partials[ml_idx] = m_state[r]; + l_partials[ml_idx] = l_state[r]; + } + } + """ + + /// TurboFlashAttention Pass 2: Cross-block reduction. + /// + /// Merges partial online softmax states from pass 1 across token blocks. + /// Each SIMD group handles one query, iterating over all blocks to produce + /// the final normalized output. + /// + /// Grid: (32, totalQueries, 1) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: Dim, num_blocks + static let turboFlashPass2Source = """ + // num_blocks via params buffer (varies per call; see pass-1 note). + const uint num_blocks = params[0]; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint q_idx = thread_position_in_grid.y; + + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + for (uint b = 0; b < (uint)num_blocks; b++) { + uint ml_idx = q_idx * num_blocks + b; + + // All lanes read the same m/l (broadcast read from device memory) + float block_m = m_partials[ml_idx]; + float block_l = l_partials[ml_idx]; + + // Skip empty blocks + if (block_l == 0.0f) continue; + + float new_m = max(m, block_m); + float exp_old = exp(m - new_m); + float exp_block = exp(block_m - new_m); + + uint partial_base = (q_idx * num_blocks + b) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o[i] = o[i] * exp_old + o_partials[partial_base + d] * exp_block; + } + } + + l = l * exp_old + block_l * exp_block; + m = new_m; + } + + // Write normalized output + float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + output[q_idx * Dim + d] = o[i] * inv_l; + } + } + """ + + /// TurboFlashAttention Pass 2 with fused output rotation. + /// + /// Same as turboFlashPass2Source but applies inverse value rotation (Π_val) in-kernel + /// after merging partials, eliminating a separate MLX matmul dispatch. + /// Uses threadgroup shared memory to gather the full output vector across SIMD lanes, + /// then each lane computes rotated output as dot product with rotation matrix rows. + /// + /// Grid: (32, totalQueries, 1) + /// Threadgroup: (32, 1, 1) + /// + /// Template params: Dim, num_blocks + static let turboFlashPass2FusedRotSource = """ + // num_blocks via params buffer (varies per call; see pass-1 note). + const uint num_blocks = params[0]; + constexpr uint DIMS_PER_LANE = (Dim + 31) / 32; + + uint lane = thread_position_in_grid.x; + uint q_idx = thread_position_in_grid.y; + + float m = -INFINITY; + float l = 0.0f; + float o[DIMS_PER_LANE]; + for (uint i = 0; i < DIMS_PER_LANE; i++) o[i] = 0.0f; + + for (uint b = 0; b < (uint)num_blocks; b++) { + uint ml_idx = q_idx * num_blocks + b; + + float block_m = m_partials[ml_idx]; + float block_l = l_partials[ml_idx]; + + if (block_l == 0.0f) continue; + + float new_m = max(m, block_m); + float exp_old = exp(m - new_m); + float exp_block = exp(block_m - new_m); + + uint partial_base = (q_idx * num_blocks + b) * Dim; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + o[i] = o[i] * exp_old + o_partials[partial_base + d] * exp_block; + } + } + + l = l * exp_old + block_l * exp_block; + m = new_m; + } + + // Normalize + float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; + + // Gather normalized output into threadgroup shared memory for rotation + threadgroup float shared_out[Dim]; + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + shared_out[d] = o[i] * inv_l; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Apply inverse value rotation: output[d] = Σ_j shared_out[j] * Π_val[j][d] + // matmul(x, Π_val) reads column d of Π_val for output dimension d. + // Π_val is stored row-major [Dim, Dim], so column d = val_rotation[j * Dim + d] + for (uint i = 0; i < DIMS_PER_LANE; i++) { + uint d = lane + i * 32; + if (d < Dim) { + float acc = 0.0f; + for (uint j = 0; j < Dim; j++) { + acc += shared_out[j] * val_rotation[j * Dim + d]; + } + output[q_idx * Dim + d] = acc; + } + } + """ + + /// Value aggregation kernel: weighted sum of codebook-quantized values. + /// + /// output[d] = Σ_t weights[t] * norm[t] * codebook[val_idx[t,d]] + /// Result is in rotated space, caller applies inverse rotation. + /// + /// Grid: (32, totalHeads, ceil(Dim/32)) + /// Threadgroup: (32, 1, 1) + static let valueKernelSource = """ + // token_count/repeat_count/L via params buffer (vary per call; template + // args would bake them into the kernel name = one compiled Metal + // library cached per token). + const uint token_count = params[0]; + const uint repeat_count = params[1]; + const uint chunk_len = params[2]; + constexpr uint MASK = (1u << Bits) - 1u; + constexpr uint LEVELS = 1u << Bits; + + uint lane = thread_position_in_grid.x; + uint head_idx = thread_position_in_grid.y; + uint dim_block = thread_position_in_grid.z; + + uint d = dim_block * 32 + lane; + if (d >= Dim) return; + + // Rows are flattened [B, heads, L]; divide out the chunk length + // before the GQA repeat mapping. + uint kv_head = (head_idx / chunk_len) / repeat_count; + + // Load codebook + float cb[LEVELS]; + for (uint i = 0; i < LEVELS; i++) { + cb[i] = codebook[i]; + } + + float acc = 0.0f; + for (uint t = 0; t < (uint)token_count; t++) { + float w = weights[head_idx * token_count + t]; + if (w < 1e-6f) continue; // Sparse V: skip negligible attention weights + + float norm_val = norms[kv_head * token_count + t]; + const device uint32_t* packed_ptr = packed + kv_head * token_count * PackedWidth + t * PackedWidth; + + uint bit_offset = d * Bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint value = (packed_ptr[word_idx] >> shift); + + int spill = (int)shift + (int)Bits - 32; + if (spill > 0) { + value |= (packed_ptr[word_idx + 1] << ((uint)Bits - (uint)spill)); + } + value &= MASK; + + acc += w * norm_val * cb[value]; + } + + output[head_idx * Dim + d] = acc; + """ +} + +// MARK: - Kernel Dispatch Wrappers + +enum TurboQuantKernelOps { + + // Kernel caches + nonisolated(unsafe) private static var encodeKernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var scoreKernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var valueKernels: [String: MLXFast.MLXFastKernel] = [:] + private static let lock = NSLock() + + /// Fused encode: norm + rotate + quantize + pack + norm correction in single GPU dispatch. + /// + /// - Parameters: + /// - input: Raw vectors [numRows, D] float32 + /// - rotation: Rotation matrix Π [D, D] float32 + /// - boundaries: Codebook boundaries [2^bits - 1] float32 + /// - codebook: Centroids [2^bits] float32 (needed for norm correction) + /// - bits: Quantization bit-width + /// - dim: Vector dimension + /// - Returns: (packed: [numRows, PackedWidth] uint32, norms: [numRows] float32) + /// norms are norm-corrected: original_norm / reconstruction_norm + /// Kernel sources declare float buffers; normalize any half/bfloat input + /// so a f16/bf16 activation stream is not silently reinterpreted. + @inline(__always) + private static func f32(_ x: MLXArray) -> MLXArray { + x.dtype == .float32 ? x : x.asType(.float32) + } + + static func fusedEncode( + input: MLXArray, + rotation: MLXArray, + boundaries: MLXArray, + codebook: MLXArray, + bits: Int, + dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "encode_nc_\(bits)_\(dim)" // nc = norm-corrected + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = encodeKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_fused_encode_\(bits)_\(dim)", + inputNames: ["input", "rotation", "boundaries", "codebook"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedEncodeSource, + ensureRowContiguous: true + ) + lock.lock() + encodeKernels[key] = k + lock.unlock() + kernel = k + } + + let numRows = input.dim(0) + + let results = kernel( + [f32(input), f32(rotation), f32(boundaries), f32(codebook)], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + + return (packed: results[0], norms: results[1]) + } + + /// Fused WHT encode: norm + WHT rotation + quantize + pack (raw norms, no correction). + /// + /// Same as fusedEncode but uses O(d log d) Walsh-Hadamard butterfly instead of + /// O(d²) dense matmul. Only works for power-of-2 dimensions. + /// + /// WHT is orthogonal so norms are preserved, no norm correction needed. + /// Codebook is NOT passed to the kernel (saves one buffer bind + GPU transfer). + /// + /// - Parameters: + /// - input: Raw vectors [numRows, D] float32 + /// - whtSigns: Random ±1 signs [D] float32 + /// - boundaries: Codebook boundaries [2^bits - 1] float32 + /// - codebook: Centroids [2^bits] float32 (unused by kernel, kept in API for caller convenience) + /// - bits: Quantization bit-width + /// - dim: Vector dimension (must be power of 2) + /// - Returns: (packed: [numRows, PackedWidth] uint32, norms: [numRows] float32, raw norms) + static func fusedEncodeWHT( + input: MLXArray, + whtSigns: MLXArray, + boundaries: MLXArray, + codebook: MLXArray, + bits: Int, + dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let logDim = Int(log2(Double(dim))) + let key = "encode_wht_\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = encodeKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_fused_encode_wht_\(bits)_\(dim)", + inputNames: ["input", "wht_signs", "boundaries"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedEncodeWHTSource, + ensureRowContiguous: true + ) + lock.lock() + encodeKernels[key] = k + lock.unlock() + kernel = k + } + + let numRows = input.dim(0) + + // NOTE: codebook no longer passed, WHT kernel stores raw norms (no norm correction) + let results = kernel( + [f32(input), f32(whtSigns), f32(boundaries)], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), ("LogDim", logDim), + ], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + + return (packed: results[0], norms: results[1]) + } + + // Flash attention kernel caches + nonisolated(unsafe) private static var flashPass1Kernels: [String: MLXFast.MLXFastKernel] = [:] + nonisolated(unsafe) private static var flashPass1NR0Kernels: [String: MLXFast.MLXFastKernel] = + [:] + nonisolated(unsafe) private static var flashPass2Kernels: [String: MLXFast.MLXFastKernel] = [:] + + /// NR0: number of query rows processed per SIMD group in the multi-row amortized kernel. + /// + /// Ported from llama.cpp V2.1: each threadgroup loads K/V packed data once and reuses + /// it across NR0 queries. At NR0=2, the KV dequant cost is halved per query. + /// + /// NR0=2 is conservative, register pressure is ~24 extra floats per thread (for dim=128). + /// Apple M-series GPUs have 96 registers per thread (384 bytes), so this fits comfortably. + /// + /// Override via environment variable `TURBO_FLASH_NR0` (must be power of 2). + static let flashNR0: Int = { + if let envValue = ProcessInfo.processInfo.environment["TURBO_FLASH_NR0"], + let parsed = Int(envValue), parsed > 0, (parsed & (parsed - 1)) == 0 + { + return parsed + } + return 2 // default, conservative starting point + }() + + /// Default block size for TurboFlashAttention two-pass approach. + /// Each SIMD group processes this many tokens per block. + /// Tuned for M1 Max via sweep: B=64 wins or ties at all token counts (512-8192+). + /// Smaller blocks = more parallelism but more pass-2 merge work. + static let flashBlockSize = 64 + + /// Shared pass 1 dispatch, used by both causal and non-causal variants. + private static func dispatchFlashPass1( + source: String, cachePrefix: String, + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, queryChunkLength: Int = 1, queryOffset: Int = 0 + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + + let pass1Key = "\(cachePrefix)_\(keyBits)_\(valueBits)_\(dim)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_\(cachePrefix)_\(keyBits)_\(valueBits)_\(dim)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: source, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), + ] + // Varying values travel in a buffer, not templates, a template value + // is baked into the kernel name and would JIT + cache one Metal + // library per distinct token_count (i.e., per generated token). + let params = MLXArray( + [ + UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), + UInt32(queryChunkLength), UInt32(queryOffset), + ]) + + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// NR0 multi-row pass 1 dispatch, processes NR0 queries per SIMD group. + /// + /// Each SIMD group loads K/V packed data once and computes scores for NR0 queries. + /// The grid Y dimension is totalQueries/NR0 instead of totalQueries. + /// Output shapes are the same as NR0=1 (partials indexed by original q_idx). + /// + /// Precondition: totalQueries must be divisible by NR0. + private static func dispatchFlashPass1NR0( + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, nr0: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let queryGroups = totalQ / nr0 + + let pass1Key = "flash_p1_nr0_\(keyBits)_\(valueBits)_\(dim)_\(nr0)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1NR0Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_nr0_\(keyBits)_\(valueBits)_\(dim)_\(nr0)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1NR0Source, + ensureRowContiguous: true + ) + lock.lock() + flashPass1NR0Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("NR0", nr0), + ] + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + + // Grid Y = queryGroups (totalQ / NR0), not totalQ + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, queryGroups, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// NR0 multi-row causal pass 1 dispatch, processes NR0 queries with per-row causal masking. + /// + /// Same as dispatchFlashPass1NR0 but supports causal masking for L>1 prefill. + /// Each query in the NR0 group has its own causal boundary. + /// + /// Precondition: totalQueries must be divisible by NR0. + private static func dispatchFlashPass1NR0Causal( + rotatedQueries: MLXArray, + keyPacked: MLXArray, keyNorms: MLXArray, keyCodebook: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, + keyBits: Int, valueBits: Int, dim: Int, + blockSize: Int, nr0: Int, + queryChunkLength: Int, queryOffset: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let kpw = TurboQuantPacking.packedWidth(count: dim, bits: keyBits) + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let queryGroups = totalQ / nr0 + + let pass1Key = "flash_p1_nr0_causal_\(keyBits)_\(valueBits)_\(dim)_\(nr0)" + let pass1Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1NR0Kernels[pass1Key] { + pass1Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_nr0_causal_\(keyBits)_\(valueBits)_\(dim)_\(nr0)", + inputNames: [ + "q_rot", "key_packed", "key_norms", "key_codebook", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1NR0CausalSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1NR0Kernels[pass1Key] = k + lock.unlock() + pass1Kernel = k + } + + let template: [(String, Int)] = [ + ("KeyBits", keyBits), ("ValueBits", valueBits), + ("Dim", dim), ("KeyPackedWidth", kpw), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("NR0", nr0), + ] + let params = MLXArray( + [ + UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), + UInt32(queryChunkLength), UInt32(queryOffset), + ]) + + let partials = pass1Kernel( + [ + f32(rotatedQueries), keyPacked, f32(keyNorms), f32(keyCodebook), + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: template, + grid: (32, queryGroups, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + + return (oPartials: partials[0], mPartials: partials[1], lPartials: partials[2]) + } + + /// Shared pass 2 dispatch, with optional fused output rotation. + /// + /// When `valRotation` is provided, the inverse value rotation (Π_val) is applied + /// in-kernel using threadgroup shared memory, eliminating a separate MLX matmul dispatch. + /// Output is in original (non-rotated) space. + /// + /// When `valRotation` is nil, output is in rotated V space (caller must apply inverse rotation). + private static func dispatchFlashPass2( + oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray, + dim: Int, numBlocks: Int, totalQ: Int, + valRotation: MLXArray? = nil + ) -> MLXArray { + let fused = valRotation != nil + let pass2Key = fused ? "flash_p2_fused_\(dim)" : "flash_p2_\(dim)" + let pass2Kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass2Kernels[pass2Key] { + pass2Kernel = cached + lock.unlock() + } else { + lock.unlock() + let k: MLXFast.MLXFastKernel + if fused { + k = MLXFast.metalKernel( + name: "turbo_flash_p2_fused_\(dim)", + inputNames: [ + "o_partials", "m_partials", "l_partials", "val_rotation", "params", + ], + outputNames: ["output"], + source: TurboQuantMetalKernels.turboFlashPass2FusedRotSource, + ensureRowContiguous: true + ) + } else { + k = MLXFast.metalKernel( + name: "turbo_flash_p2_\(dim)", + inputNames: ["o_partials", "m_partials", "l_partials", "params"], + outputNames: ["output"], + source: TurboQuantMetalKernels.turboFlashPass2Source, + ensureRowContiguous: true + ) + } + lock.lock() + flashPass2Kernels[pass2Key] = k + lock.unlock() + pass2Kernel = k + } + + let params = MLXArray([UInt32(numBlocks)]) + let inputs: [MLXArray] = + fused + ? [oPartials, mPartials, lPartials, valRotation!, params] + : [oPartials, mPartials, lPartials, params] + + return pass2Kernel( + inputs, + template: [ + ("Dim", dim) + ], + grid: (32, totalQ, 1), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ, dim]], + outputDTypes: [.float32] + )[0] + } + + nonisolated(unsafe) private static var quantizePackKernels: [String: MLXFast.MLXFastKernel] = + [:] + + /// Quantize + pack + norm-correct pre-rotated calibrated vectors. + /// `rotated` is [rows, dim] from MSECodec.rotatedUnit, `rawNorms` [rows]. + static func fusedQuantizePackScaled( + rotated: MLXArray, rawNorms: MLXArray, scale: MLXArray, + boundaries: MLXArray, codebook: MLXArray, bits: Int, dim: Int + ) -> (packed: MLXArray, norms: MLXArray) { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "qpack_scaled_\(bits)_\(dim)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = quantizePackKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_qpack_scaled_\(bits)_\(dim)", + inputNames: ["rotated", "norms_in", "scale", "boundaries", "codebook"], + outputNames: ["packed_out", "norms_out"], + source: TurboQuantMetalKernels.fusedQuantizePackScaledSource, + ensureRowContiguous: true + ) + lock.lock() + quantizePackKernels[key] = k + lock.unlock() + kernel = k + } + let numRows = rotated.dim(0) + let results = kernel( + [f32(rotated), f32(rawNorms), f32(scale), f32(boundaries), f32(codebook)], + template: [("Bits", bits), ("Dim", dim), ("PackedWidth", pw)], + grid: (dim, numRows, 1), + threadGroup: (dim, 1, 1), + outputShapes: [[numRows, pw], [numRows]], + outputDTypes: [.uint32, .float32] + ) + return (packed: results[0], norms: results[1]) + } + + /// TurboFlashAttention: two-pass fused Score + Online Softmax + Value. + /// + /// Pass 1: Parallelizes across (query × token_block) pairs. Each SIMD group processes + /// BlockSize tokens, producing partial online softmax state (m, l, o[D]). + /// Pass 2: Merges partial states across blocks to produce final normalized output. + /// + /// Eliminates intermediate score and attention weight arrays entirely. + /// + /// - Parameter valRotation: Optional [D, D] inverse value rotation matrix. When provided, + /// rotation is fused into pass 2, eliminating a separate MLX matmul dispatch. + /// Output is in original space. When nil, output is in rotated V space. + /// - Parameter blockSize: Tokens per block (default: flashBlockSize). Smaller = more parallelism + /// but more pass-2 merge work. Must be > 0. + /// - Returns: Output [totalQ, D] float32 + /// Pass-1 dispatch, raw-fp16 K + turbo-V. Reuses the packed-K flash + /// structure and pass 2. + private static func dispatchFlashPass1RawK( + rotatedQueries: MLXArray, rawKeys: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, blockSize: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let key = "flash_p1_rawk_\(valueBits)_\(dim)_\(rawKeys.dtype)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_flash_p1_rawk_\(valueBits)_\(dim)_\(rawKeys.dtype)", + inputNames: [ + "q_rot", "k_raw", "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1RawKSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[key] = k + lock.unlock() + kernel = k + } + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + // rawKeys stays in its native dtype (KT template): a per-step cast + // would copy the whole growing cache every token and push bfloat16 + // values above the float16 range to infinity. + let partials = kernel( + [ + f32(rotatedQueries), rawKeys, + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: [ + ("ValueBits", valueBits), ("Dim", dim), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("KT", rawKeys.dtype), + ], + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + return (partials[0], partials[1], partials[2]) + } + + /// Pass-1 dispatch, 8-bit affine K + turbo-V. + private static func dispatchFlashPass1AffineK( + rotatedQueries: MLXArray, + kWeights: MLXArray, kScales: MLXArray, kBiases: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, kGroup: Int, + blockSize: Int + ) -> (oPartials: MLXArray, mPartials: MLXArray, lPartials: MLXArray) { + let vpw = TurboQuantPacking.packedWidth(count: dim, bits: valueBits) + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let key = "flash_p1_affk_\(valueBits)_\(dim)_\(kGroup)_\(kScales.dtype)_\(kBiases.dtype)" + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = flashPass1Kernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: + "turbo_flash_p1_affk_\(valueBits)_\(dim)_\(kGroup)_\(kScales.dtype)_\(kBiases.dtype)", + inputNames: [ + "q_rot", "k_weights", "k_scales", "k_biases", + "val_packed", "val_norms", "val_codebook", "params", + ], + outputNames: ["o_partials", "m_partials", "l_partials"], + source: TurboQuantMetalKernels.turboFlashPass1AffineKSource, + ensureRowContiguous: true + ) + lock.lock() + flashPass1Kernels[key] = k + lock.unlock() + kernel = k + } + let params = MLXArray( + [UInt32(tokenCount), UInt32(repeatCount), UInt32(numBlocks), 1, 0] as [UInt32]) + // kScales/kBiases stay in their native dtype (KScaleT/KBiasT + // templates): they are cache-resident and grow with token_count, so + // a per-step cast would copy them every token. + let partials = kernel( + [ + f32(rotatedQueries), kWeights, kScales, kBiases, + valPacked, f32(valNorms), f32(valCodebook), params, + ], + template: [ + ("ValueBits", valueBits), ("Dim", dim), ("ValuePackedWidth", vpw), + ("BlockSize", blockSize), ("KGroup", kGroup), + ("KScaleT", kScales.dtype), ("KBiasT", kBiases.dtype), + ], + grid: (32, totalQ, numBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ * numBlocks, dim], [totalQ, numBlocks], [totalQ, numBlocks]], + outputDTypes: [.float32, .float32, .float32] + ) + return (partials[0], partials[1], partials[2]) + } + + /// Flash decode, raw-fp16 K + turbo-V (single decode step, L=1). + static func turboFlashRawK( + rotatedQueries: MLXArray, rawKeys: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, + valRotation: MLXArray? = nil, blockSize: Int? = nil + ) -> MLXArray { + let bs = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + bs - 1) / bs + let totalQ = rotatedQueries.dim(0) + let (o, m, l) = dispatchFlashPass1RawK( + rotatedQueries: rotatedQueries, rawKeys: rawKeys, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, valueBits: valueBits, + dim: dim, blockSize: bs) + return dispatchFlashPass2( + oPartials: o, mPartials: m, lPartials: l, dim: dim, numBlocks: numBlocks, + totalQ: totalQ, valRotation: valRotation) + } + + /// Flash decode, 8-bit affine K + turbo-V (single decode step, L=1). + static func turboFlashAffineK( + rotatedQueries: MLXArray, + kWeights: MLXArray, kScales: MLXArray, kBiases: MLXArray, + valPacked: MLXArray, valNorms: MLXArray, valCodebook: MLXArray, + tokenCount: Int, repeatCount: Int, valueBits: Int, dim: Int, kGroup: Int, + valRotation: MLXArray? = nil, blockSize: Int? = nil + ) -> MLXArray { + let bs = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + bs - 1) / bs + let totalQ = rotatedQueries.dim(0) + let (o, m, l) = dispatchFlashPass1AffineK( + rotatedQueries: rotatedQueries, + kWeights: kWeights, kScales: kScales, kBiases: kBiases, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, valueBits: valueBits, + dim: dim, kGroup: kGroup, blockSize: bs) + return dispatchFlashPass2( + oPartials: o, mPartials: m, lPartials: l, dim: dim, numBlocks: numBlocks, + totalQ: totalQ, valRotation: valRotation) + } + + static func turboFlashAttention( + rotatedQueries: MLXArray, + keyPacked: MLXArray, + keyNorms: MLXArray, + keyCodebook: MLXArray, + valPacked: MLXArray, + valNorms: MLXArray, + valCodebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + keyBits: Int, + valueBits: Int, + dim: Int, + valRotation: MLXArray? = nil, + blockSize: Int? = nil + ) -> MLXArray { + let blockSize = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let nr0 = flashNR0 + + // Use NR0 multi-row kernel when totalQ is evenly divisible by NR0 and NR0 > 1. + // Falls back to NR0=1 (original kernel) for remainder queries or when NR0=1. + // NR0 groups read one kv head (kv_indices[0]) for all rows; only valid + // when the GQA repeat factor is a multiple of nr0 so aligned groups can + // never span a KV-head boundary. MHA and odd repeat factors take the + // per-row kernel. + let useNR0 = + nr0 > 1 && totalQ % nr0 == 0 && totalQ >= nr0 && repeatCount % nr0 == 0 + + let oPartials: MLXArray + let mPartials: MLXArray + let lPartials: MLXArray + + if useNR0 { + (oPartials, mPartials, lPartials) = dispatchFlashPass1NR0( + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, nr0: nr0 + ) + } else { + (oPartials, mPartials, lPartials) = dispatchFlashPass1( + source: TurboQuantMetalKernels.turboFlashPass1Source, + cachePrefix: "flash_p1", + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize + ) + } + + return dispatchFlashPass2( + oPartials: oPartials, mPartials: mPartials, lPartials: lPartials, + dim: dim, numBlocks: numBlocks, totalQ: totalQ, + valRotation: valRotation + ) + } + + /// TurboFlashAttention with causal masking for L>1 prefill chunks. + /// + /// Same as turboFlashAttention but each query position only attends to tokens + /// where t <= queryOffset + q_within_L. Eliminates the need to materialize + /// the full [nQHeads, L, T] score matrix for causal masking. + /// + /// - Parameter queryChunkLength: Number of query positions in the chunk (L) + /// - Parameter queryOffset: Absolute position of the first query in the chunk + /// - Returns: Output [totalQ, D] float32 + static func turboFlashAttentionCausal( + rotatedQueries: MLXArray, + keyPacked: MLXArray, + keyNorms: MLXArray, + keyCodebook: MLXArray, + valPacked: MLXArray, + valNorms: MLXArray, + valCodebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + keyBits: Int, + valueBits: Int, + dim: Int, + queryChunkLength: Int, + queryOffset: Int, + valRotation: MLXArray? = nil, + blockSize: Int? = nil + ) -> MLXArray { + let blockSize = blockSize ?? flashBlockSize + let numBlocks = (tokenCount + blockSize - 1) / blockSize + let totalQ = rotatedQueries.dim(0) + let nr0 = flashNR0 + + // NR0 groups read one kv head (kv_indices[0]) for all rows; only valid + // when the GQA repeat factor is a multiple of nr0 so aligned groups can + // never span a KV-head boundary. MHA and odd repeat factors take the + // per-row kernel. + let useNR0 = + nr0 > 1 && totalQ % nr0 == 0 && totalQ >= nr0 && repeatCount % nr0 == 0 + + let oPartials: MLXArray + let mPartials: MLXArray + let lPartials: MLXArray + + if useNR0 { + (oPartials, mPartials, lPartials) = dispatchFlashPass1NR0Causal( + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, nr0: nr0, + queryChunkLength: queryChunkLength, queryOffset: queryOffset + ) + } else { + (oPartials, mPartials, lPartials) = dispatchFlashPass1( + source: TurboQuantMetalKernels.turboFlashPass1CausalSource, + cachePrefix: "flash_p1_causal", + rotatedQueries: rotatedQueries, + keyPacked: keyPacked, keyNorms: keyNorms, keyCodebook: keyCodebook, + valPacked: valPacked, valNorms: valNorms, valCodebook: valCodebook, + tokenCount: tokenCount, repeatCount: repeatCount, + keyBits: keyBits, valueBits: valueBits, dim: dim, + blockSize: blockSize, + queryChunkLength: queryChunkLength, queryOffset: queryOffset + ) + } + + return dispatchFlashPass2( + oPartials: oPartials, mPartials: mPartials, lPartials: lPartials, + dim: dim, numBlocks: numBlocks, totalQ: totalQ, + valRotation: valRotation + ) + } + + /// Compute Q×K attention scores from packed codebook indices. + /// + /// - Parameters: + /// - rotatedQueries: Pre-rotated queries [totalQ, D] (already scaled) + /// - packed: Packed key indices [totalKVHeads, T, PackedWidth] uint32 + /// - norms: Key norms [totalKVHeads, T] float32 + /// - codebook: Centroids [2^bits] float32 + /// - tokenCount: Number of cached tokens + /// - repeatCount: GQA repeat factor (nQHeads / nKVHeads) + /// - bits: MSE bit-width + /// - dim: Vector dimension + /// - Returns: Scores [totalQ, T] float32 + static func mseScore( + rotatedQueries: MLXArray, + packed: MLXArray, + norms: MLXArray, + codebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + bits: Int, + dim: Int, queryChunkLength: Int = 1 + ) -> MLXArray { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = scoreKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_score_\(bits)_\(dim)", + inputNames: ["q_rot", "packed", "norms", "codebook", "params"], + outputNames: ["scores"], + source: TurboQuantMetalKernels.scoreKernelSource, + ensureRowContiguous: true + ) + lock.lock() + scoreKernels[key] = k + lock.unlock() + kernel = k + } + + let totalQ = rotatedQueries.dim(0) + + let params = MLXArray([UInt32(tokenCount), UInt32(repeatCount), UInt32(queryChunkLength)]) + return kernel( + [f32(rotatedQueries), packed, f32(norms), f32(codebook), params], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (32, totalQ, tokenCount), + threadGroup: (32, 1, 1), + outputShapes: [[totalQ, tokenCount]], + outputDTypes: [.float32] + )[0] + } + + /// Compute weighted sum of values from packed codebook indices. + /// + /// Result is in ROTATED space, caller must apply inverse rotation. + /// + /// - Returns: [totalHeads, D] float32 (rotated space) + static func mseWeightedSum( + weights: MLXArray, + packed: MLXArray, + norms: MLXArray, + codebook: MLXArray, + tokenCount: Int, + repeatCount: Int, + bits: Int, + dim: Int, queryChunkLength: Int = 1 + ) -> MLXArray { + let pw = TurboQuantPacking.packedWidth(count: dim, bits: bits) + let key = "\(bits)_\(dim)" + + let kernel: MLXFast.MLXFastKernel + lock.lock() + if let cached = valueKernels[key] { + kernel = cached + lock.unlock() + } else { + lock.unlock() + let k = MLXFast.metalKernel( + name: "turbo_value_\(bits)_\(dim)", + inputNames: ["weights", "packed", "norms", "codebook", "params"], + outputNames: ["output"], + source: TurboQuantMetalKernels.valueKernelSource, + ensureRowContiguous: true + ) + lock.lock() + valueKernels[key] = k + lock.unlock() + kernel = k + } + + let totalHeads = weights.dim(0) + let dimBlocks = (dim + 31) / 32 + + let params = MLXArray([UInt32(tokenCount), UInt32(repeatCount), UInt32(queryChunkLength)]) + return kernel( + [f32(weights), packed, f32(norms), f32(codebook), params], + template: [ + ("Bits", bits), ("Dim", dim), ("PackedWidth", pw), + ], + grid: (32, totalHeads, dimBlocks), + threadGroup: (32, 1, 1), + outputShapes: [[totalHeads, dim]], + outputDTypes: [.float32] + )[0] + } +} diff --git a/Libraries/MLXVLM/Gemma4AssistantRegistration.swift b/Libraries/MLXVLM/Gemma4AssistantRegistration.swift index a8c1ca4d5..d1fe9c2f1 100644 --- a/Libraries/MLXVLM/Gemma4AssistantRegistration.swift +++ b/Libraries/MLXVLM/Gemma4AssistantRegistration.swift @@ -25,13 +25,21 @@ import MLXLMCommon /// ``` public enum Gemma4AssistantRegistration { public static func register() async { - await MTPDrafterTypeRegistry.shared.registerModelType( - "gemma4_assistant", - creator: { data in - let config = try JSONDecoder().decode( - Gemma4AssistantConfiguration.self, from: data) - return Gemma4AssistantDraftModel(config) - } - ) + // `gemma4_assistant` (E-series, 26B-A4B, 31B drafters) and + // `gemma4_unified_assistant` (the 12B drafter, whose target is the + // `gemma4_unified` VLM) share the drafter implementation: the unified + // variant differs only in its `text_config` (`gemma4_unified_text`, + // `attention_k_eq_v`, `num_global_key_value_heads`), which + // `Gemma4TextConfiguration` already decodes. + for modelType in ["gemma4_assistant", "gemma4_unified_assistant"] { + await MTPDrafterTypeRegistry.shared.registerModelType( + modelType, + creator: { data in + let config = try JSONDecoder().decode( + Gemma4AssistantConfiguration.self, from: data) + return Gemma4AssistantDraftModel(config) + } + ) + } } } diff --git a/Libraries/MLXVLM/Models/Gemma4.swift b/Libraries/MLXVLM/Models/Gemma4.swift index 5e32b878f..2fee2d7a5 100644 --- a/Libraries/MLXVLM/Models/Gemma4.swift +++ b/Libraries/MLXVLM/Models/Gemma4.swift @@ -7,15 +7,11 @@ import MLXNN // Based on https://github.com/Blaizzy/mlx-vlm/tree/main/mlx_vlm/models/gemma4 private enum Gemma4Error: LocalizedError { - case imageTokenCountMismatch(expectedVisionTokens: Int, actualPromptTokens: Int) case multimodalTokenCountMismatch(kind: String, featureTokens: Int, promptTokens: Int) case imagePlaceholderMismatch(images: Int, placeholders: Int) var errorDescription: String? { switch self { - case .imageTokenCountMismatch(let expectedVisionTokens, let actualPromptTokens): - return - "Gemma4 image token count mismatch: vision encoder produced \(expectedVisionTokens) soft tokens, but the prompt contains \(actualPromptTokens) image tokens." case .multimodalTokenCountMismatch(let kind, let featureTokens, let promptTokens): return "Gemma4 \(kind) token count mismatch: encoder produced \(featureTokens) soft tokens, but the prompt contains \(promptTokens) \(kind) tokens." @@ -586,6 +582,7 @@ public struct Gemma4Configuration: Codable, Sendable { public let visionSoftTokensPerImage: Int public let audioSoftTokensPerImage: Int public let audioMsPerToken: Int + public let visionSoftTokensPerVideoFrame: Int public let tieWordEmbeddings: Bool private let _vocabularySize: Int? @@ -612,6 +609,7 @@ public struct Gemma4Configuration: Codable, Sendable { case visionSoftTokensPerImage = "vision_soft_tokens_per_image" case audioSoftTokensPerImage = "audio_soft_tokens_per_image" case audioMsPerToken = "audio_ms_per_token" + case visionSoftTokensPerVideoFrame = "vision_soft_tokens_per_video_frame" case tieWordEmbeddings = "tie_word_embeddings" case _vocabularySize = "vocab_size" case _hiddenSize = "hidden_size" @@ -643,6 +641,8 @@ public struct Gemma4Configuration: Codable, Sendable { try c.decodeIfPresent(Int.self, forKey: CodingKeys.audioSoftTokensPerImage) ?? 750 audioMsPerToken = try c.decodeIfPresent(Int.self, forKey: CodingKeys.audioMsPerToken) ?? 40 + visionSoftTokensPerVideoFrame = + try c.decodeIfPresent(Int.self, forKey: CodingKeys.visionSoftTokensPerVideoFrame) ?? 70 tieWordEmbeddings = try c.decodeIfPresent(Bool.self, forKey: CodingKeys.tieWordEmbeddings) ?? textConfiguration.tieWordEmbeddings @@ -2043,7 +2043,8 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { inputIds: MLXArray, image: LMInput.ProcessedImage? = nil, audioFeatures: MLXArray? = nil, - audioMask: MLXArray? = nil + audioMask: MLXArray? = nil, + video: LMInput.ProcessedVideo? = nil ) throws -> (MLXArray, MLXArray?) { var inputsEmbeds = languageModel.model.embedTokens(inputIds) inputsEmbeds = @@ -2053,15 +2054,17 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { var perLayerInputs: MLXArray? = nil if config.textConfiguration.hiddenSizePerLayerInput > 0 { - let imageMask = inputIds .== config.imageTokenId - let audioTokenMask = - if let audioTokenId = config.audioTokenId { - inputIds .== audioTokenId - } else { - MLXArray.zeros(like: imageMask) - } - let textMask = logicalNot(logicalOr(imageMask, audioTokenMask)) - let perLayerTokens = MLX.where(textMask, inputIds, MLXArray.zeros(like: inputIds)) + // Per-layer inputs are text-only: zero out every multimodal soft token + // (image / audio / video) so their ids don't index the PLE embedding. + var multimodalMask = inputIds .== config.imageTokenId + if let audioTokenId = config.audioTokenId { + multimodalMask = multimodalMask | (inputIds .== audioTokenId) + } + if let videoTokenId = config.videoTokenId { + multimodalMask = multimodalMask | (inputIds .== videoTokenId) + } + let perLayerTokens = MLX.where( + logicalNot(multimodalMask), inputIds, MLXArray.zeros(like: inputIds)) perLayerInputs = languageModel.model.getPerLayerInputs(perLayerTokens) } @@ -2095,9 +2098,9 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { let expectedImageTokens = imageMask.asType(.int32).sum().item(Int.self) if expectedImageTokens != imageFeatures.dim(1) { - throw Gemma4Error.imageTokenCountMismatch( - expectedVisionTokens: imageFeatures.dim(1), - actualPromptTokens: expectedImageTokens) + throw Gemma4Error.multimodalTokenCountMismatch( + kind: "image", featureTokens: imageFeatures.dim(1), + promptTokens: expectedImageTokens) } var imageMaskExpanded = expandedDimensions(imageMask, axis: -1) @@ -2130,21 +2133,66 @@ public final class Gemma4: Module, VLMModel, KVCacheDimensionProvider { ) } + // Gemma 4 has no separate video encoder — each video frame runs through the + // same vision tower as images (producing `visionSoftTokensPerImage` pooled + // tokens per frame) and is then truncated to the smaller per-frame video + // budget before scattering onto the `