diff --git a/apps/omlx-mac/Resources/Localizable.xcstrings b/apps/omlx-mac/Resources/Localizable.xcstrings index 5d23f35c2..b009a7f39 100644 --- a/apps/omlx-mac/Resources/Localizable.xcstrings +++ b/apps/omlx-mac/Resources/Localizable.xcstrings @@ -11512,6 +11512,24 @@ } } }, + "profile.detail.acceleration.turboquant.mid_prefill" : { + "comment" : "TurboQuant profile-chip suffix indicating opt-in conversion when a full prefill chunk cannot fit", + "extractionState" : "manual", + "localizations" : { + "en" : { + "stringUnit" : { + "state" : "translated", + "value" : "prefill pressure" + } + }, + "ru" : { + "stringUnit" : { + "state" : "translated", + "value" : "prefill pressure" + } + } + } + }, "profile.detail.acceleration.turboquant.skip" : { "comment" : "TurboQuant skip-last-N suffix; placeholder is the layer count", "extractionState" : "extracted_with_value", @@ -22007,20 +22025,56 @@ } } }, + "settings.experimental.turboquant.mid_prefill.label" : { + "comment" : "Subordinate TurboQuant row label for converting the KV cache when a full prefill chunk cannot fit", + "extractionState" : "manual", + "localizations" : { + "en" : { + "stringUnit" : { + "state" : "translated", + "value" : "Convert under prefill pressure" + } + }, + "ru" : { + "stringUnit" : { + "state" : "translated", + "value" : "Convert under prefill pressure" + } + } + } + }, + "settings.experimental.turboquant.mid_prefill.sub" : { + "comment" : "Help text for converting the TurboQuant KV cache once under prefill memory pressure", + "extractionState" : "manual", + "localizations" : { + "en" : { + "stringUnit" : { + "state" : "translated", + "value" : "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill." + } + }, + "ru" : { + "stringUnit" : { + "state" : "translated", + "value" : "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill." + } + } + } + }, "settings.experimental.turboquant.sub" : { - "comment" : "Sublabel describing TurboQuant KV cache", + "comment" : "Sublabel describing the normal TurboQuant KV cache conversion path", "extractionState" : "manual", "localizations" : { "en" : { "stringUnit" : { "state" : "translated", - "value" : "Quantize the KV cache during prefill. Saves memory at a small quality cost." + "value" : "Convert the KV cache after prefill for generation. Saves memory at a small quality cost." } }, "ru" : { "stringUnit" : { "state" : "translated", - "value" : "Квантовать KV-кэш во время префилла. Экономит память с небольшой потерей качества." + "value" : "Квантовать KV-кэш после префилла для генерации. Экономит память с небольшой потерей качества." } }, "zh-Hans" : { diff --git a/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift b/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift index bd9f05e57..dfd9dd1a9 100644 --- a/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift +++ b/apps/omlx-mac/Sources/AppView/Screens/ModelSettingsScreen.swift @@ -1111,6 +1111,15 @@ private struct ExperimentalSection: View { .help(vm.vlmMtpEnabled ? vlmMtpOwnsSpeculativePathReason : "") } } + if vm.showsTurboquantMidPrefill { + Row(label: turboquantMidPrefillLabel, + sublabel: turboquantMidPrefillHelp) { + Toggle(turboquantMidPrefillLabel, + isOn: vm.bindProfile($vm.turboquantMidPrefill)) + .labelsHidden().toggleStyle(.switch) + .help(turboquantMidPrefillHelp) + } + } // IndexCache (DSA-only — surface to the user that the row // only applies to models whose config matches the DSA set). @@ -1370,8 +1379,20 @@ private struct ExperimentalSection: View { private var turboquantSublabel: String { if vm.vlmMtpEnabled { return vlmMtpOwnsSpeculativePathReason } return String(localized: "settings.experimental.turboquant.sub", - defaultValue: "Quantize the KV cache during prefill. Saves memory at a small quality cost.", - comment: "Sublabel describing TurboQuant KV cache") + defaultValue: "Convert the KV cache after prefill for generation. Saves memory at a small quality cost.", + comment: "Sublabel describing the normal TurboQuant KV cache conversion path") + } + + private var turboquantMidPrefillLabel: String { + String(localized: "settings.experimental.turboquant.mid_prefill.label", + defaultValue: "Convert under prefill pressure", + comment: "Subordinate TurboQuant row label for converting the KV cache when a full prefill chunk cannot fit") + } + + private var turboquantMidPrefillHelp: String { + String(localized: "settings.experimental.turboquant.mid_prefill.sub", + defaultValue: "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", + comment: "Help text for converting the TurboQuant KV cache once under prefill memory pressure") } private var specprefillSublabel: String { diff --git a/apps/omlx-mac/Sources/AppView/Screens/ProfileViews.swift b/apps/omlx-mac/Sources/AppView/Screens/ProfileViews.swift index 077d8cf34..fa5ea22da 100644 --- a/apps/omlx-mac/Sources/AppView/Screens/ProfileViews.swift +++ b/apps/omlx-mac/Sources/AppView/Screens/ProfileViews.swift @@ -10,6 +10,13 @@ import SwiftUI +/// Profile-chip detail state for TurboQuant's opt-in mid-prefill conversion. +/// The child flag is meaningful only while its parent feature is enabled. +func turboquantMidPrefillDetailIsActive(_ settings: [String: AnyCodable]) -> Bool { + boolOf(settings[ProfileSettingsKey.turboquantKvEnabled]) == true + && boolOf(settings[ProfileSettingsKey.turboquantMidPrefill]) == true +} + // MARK: - Scope colors / labels /// Per-scope visual treatment. Lifted from omlx-screens.jsx:878-884 @@ -980,7 +987,12 @@ struct ProfileDetailCard: View { defaultValue: "skip \(skip)", comment: "TurboQuant skip-last-N suffix; placeholder is the layer count") : nil - let parts = [bitsText, skipText].compactMap { $0 } + let midPrefillText = turboquantMidPrefillDetailIsActive(s) + ? String(localized: "profile.detail.acceleration.turboquant.mid_prefill", + defaultValue: "prefill pressure", + comment: "TurboQuant profile-chip suffix indicating opt-in conversion when a full prefill chunk cannot fit") + : nil + let parts = [bitsText, skipText, midPrefillText].compactMap { $0 } return parts.isEmpty ? baseName : String(localized: "profile.detail.acceleration.turboquant.with_parts", diff --git a/apps/omlx-mac/Sources/AppView/Utils/ProfileWorkingState.swift b/apps/omlx-mac/Sources/AppView/Utils/ProfileWorkingState.swift index fe262bcf6..e1b9c5f32 100644 --- a/apps/omlx-mac/Sources/AppView/Utils/ProfileWorkingState.swift +++ b/apps/omlx-mac/Sources/AppView/Utils/ProfileWorkingState.swift @@ -129,6 +129,7 @@ enum ProfileSettingsKey { static let trustRemoteCode = "trust_remote_code" static let turboquantKvEnabled = "turboquant_kv_enabled" static let turboquantKvBits = "turboquant_kv_bits" + static let turboquantMidPrefill = "turboquant_mid_prefill" static let indexCacheFreq = "index_cache_freq" static let specprefillEnabled = "specprefill_enabled" static let specprefillDraftModel = "specprefill_draft_model" diff --git a/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift b/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift index 060002c11..29fd26ab5 100644 --- a/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift +++ b/apps/omlx-mac/Sources/AppView/ViewModels/ModelSettingsScreenVM.swift @@ -34,7 +34,7 @@ final class ModelSettingsScreenVM { case trustRemoteCode case reasoningParser case chatTemplateKwargs - case turboquantKvEnabled, turboquantKvBits + case turboquantKvEnabled, turboquantKvBits, turboquantMidPrefill case indexCacheEnabled, indexCacheFreq case specprefillEnabled, specprefillDraftModel, specprefillKeepPct, specprefillThreshold case dflashEnabled, dflashDraftModel, dflashMaxCtx @@ -253,6 +253,7 @@ final class ModelSettingsScreenVM { // Experimental: TurboQuant KV var turboquantKvEnabled: Bool = false var turboquantKvBits: String = "4" + var turboquantMidPrefill: Bool = false // Experimental: IndexCache (DSA-only) var indexCacheEnabled: Bool = false @@ -329,6 +330,10 @@ final class ModelSettingsScreenVM { return Self.diffusionConfigModelTypes.contains(type) } + var showsTurboquantMidPrefill: Bool { + !isDiffusionModel && turboquantKvEnabled + } + private func isDiffusionUnsupportedField(_ field: Field) -> Bool { switch field { case .topP, .topK, .minP, .repetitionPenalty, .presencePenalty: @@ -339,7 +344,7 @@ final class ModelSettingsScreenVM { return true case .forceSampling, .reasoningParser: return true - case .turboquantKvEnabled, .turboquantKvBits: + case .turboquantKvEnabled, .turboquantKvBits, .turboquantMidPrefill: return true case .indexCacheEnabled, .indexCacheFreq: return true @@ -462,6 +467,8 @@ final class ModelSettingsScreenVM { ) self.turboquantKvEnabled = s?.turboquantKvEnabled ?? false self.turboquantKvBits = s?.turboquantKvBits.map { Self.formatBits($0) } ?? "4" + self.turboquantMidPrefill = !self.isDiffusionModel + && (s?.turboquantMidPrefill ?? false) self.indexCacheEnabled = s?.indexCacheFreq != nil self.indexCacheFreq = s?.indexCacheFreq.map(String.init) ?? "4" self.specprefillEnabled = s?.specprefillEnabled ?? false @@ -594,6 +601,8 @@ final class ModelSettingsScreenVM { patch.forcedCtKwargs = pair.forced ?? [] case .turboquantKvEnabled: patch.turboquantKvEnabled = turboquantKvEnabled case .turboquantKvBits: patch.turboquantKvBits = Double(turboquantKvBits) + case .turboquantMidPrefill: + patch.turboquantMidPrefill = turboquantMidPrefill case .indexCacheEnabled: patch.indexCacheFreq = indexCacheEnabled ? (Int(indexCacheFreq) ?? 4) : 0 case .indexCacheFreq: @@ -886,6 +895,10 @@ final class ModelSettingsScreenVM { if turboquantKvEnabled, let bits = Double(turboquantKvBits) { out[ProfileSettingsKey.turboquantKvBits] = AnyCodable(bits) } + putBool( + ProfileSettingsKey.turboquantMidPrefill, + turboquantMidPrefill + ) if indexCacheEnabled, let n = Int(indexCacheFreq), n >= 2 { out[ProfileSettingsKey.indexCacheFreq] = AnyCodable(n) } diff --git a/apps/omlx-mac/Sources/Net/DTO/ModelsDTO.swift b/apps/omlx-mac/Sources/Net/DTO/ModelsDTO.swift index 876d6d73a..30aff6610 100644 --- a/apps/omlx-mac/Sources/Net/DTO/ModelsDTO.swift +++ b/apps/omlx-mac/Sources/Net/DTO/ModelsDTO.swift @@ -102,6 +102,7 @@ struct ModelSettingsDTO: Codable, Equatable, Sendable { // Experimental: TurboQuant KV cache let turboquantKvEnabled: Bool? let turboquantKvBits: Double? + let turboquantMidPrefill: Bool? // Experimental: IndexCache (DSA models only) let indexCacheFreq: Int? // Experimental: SpecPrefill @@ -166,6 +167,7 @@ struct ModelSettingsPatch: Encodable, Equatable, Sendable { // Experimental: TurboQuant KV var turboquantKvEnabled: Bool? = nil var turboquantKvBits: Double? = nil + var turboquantMidPrefill: Bool? = nil // Experimental: IndexCache var indexCacheFreq: Int? = nil // Experimental: SpecPrefill diff --git a/apps/omlx-mac/Tests/oMLXTests/DTOFixtureTests.swift b/apps/omlx-mac/Tests/oMLXTests/DTOFixtureTests.swift index 78c84752f..04e67dcae 100644 --- a/apps/omlx-mac/Tests/oMLXTests/DTOFixtureTests.swift +++ b/apps/omlx-mac/Tests/oMLXTests/DTOFixtureTests.swift @@ -125,9 +125,41 @@ final class DTOFixtureTests: XCTestCase { if let first = list.models.first { XCTAssertFalse(first.id.isEmpty, "ModelDTO.id must be non-empty.") XCTAssertEqual(first.displayName, "deepsweet/Qwen3.6-27B-UD-MLX-4bit") + XCTAssertEqual(first.settings?.turboquantMidPrefill, false) } } + func testTurboquantMidPrefillResponseAllowsMissingField() throws { + let data = Data(#"{}"#.utf8) + + let settings = try Self.makeDecoder().decode(ModelSettingsDTO.self, from: data) + + XCTAssertNil(settings.turboquantMidPrefill) + } + + func testTurboquantMidPrefillResponseDecodesPresentTrue() throws { + let data = Data(#"{"turboquant_mid_prefill":true}"#.utf8) + + let settings = try Self.makeDecoder().decode(ModelSettingsDTO.self, from: data) + + XCTAssertEqual(settings.turboquantMidPrefill, true) + } + + func testTurboquantMidPrefillPatchUsesExactWireKey() throws { + var patch = ModelSettingsPatch() + patch.turboquantMidPrefill = true + let encoder = JSONEncoder() + encoder.keyEncodingStrategy = .convertToSnakeCase + + let body = try XCTUnwrap( + JSONSerialization.jsonObject(with: encoder.encode(patch)) as? [String: Any] + ) + + XCTAssertEqual(body.count, 1) + XCTAssertEqual(body["turboquant_mid_prefill"] as? Bool, true) + XCTAssertNil(body["turboquantMidPrefill"]) + } + // MARK: - Profile list (per-model) func testModelProfilesFixtureDecodes() throws { @@ -135,6 +167,10 @@ final class DTOFixtureTests: XCTestCase { let resp = try Self.makeDecoder().decode(ProfileListResponse.self, from: data) XCTAssertNotNil(resp.profiles, "Profiles array must be present even when empty.") + XCTAssertEqual( + resp.profiles.first?.settings?[ProfileSettingsKey.turboquantMidPrefill]?.value as? Bool, + false + ) } // MARK: - Profile templates diff --git a/apps/omlx-mac/Tests/oMLXTests/Fixtures/model-profiles.json b/apps/omlx-mac/Tests/oMLXTests/Fixtures/model-profiles.json index 59e1c1522..826ed9b48 100644 --- a/apps/omlx-mac/Tests/oMLXTests/Fixtures/model-profiles.json +++ b/apps/omlx-mac/Tests/oMLXTests/Fixtures/model-profiles.json @@ -11,6 +11,7 @@ "max_context_window": 131072, "force_sampling": false, "turboquant_kv_enabled": false, + "turboquant_mid_prefill": false, "dflash_enabled": false, "max_tokens": 262144, "thinking_budget_tokens": 8192, diff --git a/apps/omlx-mac/Tests/oMLXTests/Fixtures/models.json b/apps/omlx-mac/Tests/oMLXTests/Fixtures/models.json index a1cfdbe87..5fee67423 100644 --- a/apps/omlx-mac/Tests/oMLXTests/Fixtures/models.json +++ b/apps/omlx-mac/Tests/oMLXTests/Fixtures/models.json @@ -49,6 +49,7 @@ "reasoning_parser": null, "turboquant_kv_enabled": false, "turboquant_kv_bits": 4.0, + "turboquant_mid_prefill": false, "turboquant_skip_last": true, "specprefill_enabled": false, "specprefill_draft_model": null, diff --git a/apps/omlx-mac/Tests/oMLXTests/LocalizationSmokeTests.swift b/apps/omlx-mac/Tests/oMLXTests/LocalizationSmokeTests.swift index c2eddcfb0..e8993f54f 100644 --- a/apps/omlx-mac/Tests/oMLXTests/LocalizationSmokeTests.swift +++ b/apps/omlx-mac/Tests/oMLXTests/LocalizationSmokeTests.swift @@ -73,6 +73,9 @@ final class LocalizationSmokeTests: XCTestCase { "bench.context.header.title", "bench.context.section.configuration", // Settings + helpers "settings.section.basic", "settings.advanced.experimental.section", + "settings.experimental.turboquant.mid_prefill.label", + "settings.experimental.turboquant.mid_prefill.sub", + "profile.detail.acceleration.turboquant.mid_prefill", "appearance.row.menubar_icon", "appearance.row.menubar_icon.restore", // Menubar + updates "menubar.item.quit", "menubar.stats.session_section", diff --git a/apps/omlx-mac/Tests/oMLXTests/ModelSettingsScreenVMTests.swift b/apps/omlx-mac/Tests/oMLXTests/ModelSettingsScreenVMTests.swift index dc5321dd2..03f7c1b8a 100644 --- a/apps/omlx-mac/Tests/oMLXTests/ModelSettingsScreenVMTests.swift +++ b/apps/omlx-mac/Tests/oMLXTests/ModelSettingsScreenVMTests.swift @@ -60,6 +60,73 @@ final class ModelSettingsScreenVMTests: XCTestCase { XCTAssertFalse(values.contains("model-MTPLX-runtime")) } + func testTurboquantMidPrefillDefaultsFalseAndHidden() { + let vm = ModelSettingsScreenVM() + + XCTAssertFalse(vm.turboquantMidPrefill) + XCTAssertFalse(vm.showsTurboquantMidPrefill) + XCTAssertEqual(ProfileSettingsKey.turboquantMidPrefill, "turboquant_mid_prefill") + } + + func testTurboquantMidPrefillProfileStateRoundTripsWhenParentEnabled() { + let vm = ModelSettingsScreenVM() + vm.model = makeModel(id: "text-model", configModelType: "qwen3_5") + vm.turboquantKvEnabled = true + vm.turboquantMidPrefill = true + + let settings = vm.currentSettingsDict() + + XCTAssertTrue(vm.showsTurboquantMidPrefill) + XCTAssertEqual( + settings["turboquant_mid_prefill"]?.value as? Bool, + true + ) + XCTAssertTrue(turboquantMidPrefillDetailIsActive(settings)) + } + + func testTurboquantMidPrefillStaysSavedButInertWhenParentDisabled() { + let vm = ModelSettingsScreenVM() + vm.model = makeModel(id: "text-model", configModelType: "qwen3_5") + vm.turboquantKvEnabled = false + vm.turboquantMidPrefill = true + + let settings = vm.currentSettingsDict() + + XCTAssertFalse(vm.showsTurboquantMidPrefill) + XCTAssertEqual( + settings["turboquant_mid_prefill"]?.value as? Bool, + true + ) + XCTAssertFalse(turboquantMidPrefillDetailIsActive(settings)) + } + + func testTurboquantMidPrefillIsExcludedForDiffusion() { + let vm = ModelSettingsScreenVM() + vm.model = makeModel(id: "diffusion-model", configModelType: "diffusion_gemma") + vm.turboquantKvEnabled = true + vm.turboquantMidPrefill = true + + let settings = vm.currentSettingsDict() + + XCTAssertFalse(vm.showsTurboquantMidPrefill) + XCTAssertNil(settings["turboquant_mid_prefill"]) + XCTAssertFalse(turboquantMidPrefillDetailIsActive(settings)) + } + + func testTurboquantProfileDetailRequiresBothFlags() { + let parentOnly: [String: AnyCodable] = [ + ProfileSettingsKey.turboquantKvEnabled: AnyCodable(true), + ProfileSettingsKey.turboquantMidPrefill: AnyCodable(false), + ] + let childOnly: [String: AnyCodable] = [ + ProfileSettingsKey.turboquantKvEnabled: AnyCodable(false), + ProfileSettingsKey.turboquantMidPrefill: AnyCodable(true), + ] + + XCTAssertFalse(turboquantMidPrefillDetailIsActive(parentOnly)) + XCTAssertFalse(turboquantMidPrefillDetailIsActive(childOnly)) + } + private func makeModel(id: String, configModelType: String?) -> ModelDTO { ModelDTO( id: id, diff --git a/omlx/admin/benchmark.py b/omlx/admin/benchmark.py index 3e754d7a3..8d3156ef2 100644 --- a/omlx/admin/benchmark.py +++ b/omlx/admin/benchmark.py @@ -301,6 +301,7 @@ def _derive_feature_flags(model_settings: Any) -> list[dict]: "model_type_override", "index_cache_freq", "turboquant_kv_enabled", + "turboquant_mid_prefill", "turboquant_kv_bits", "turboquant_skip_last", "specprefill_enabled", @@ -361,13 +362,15 @@ def _filter_uploaded_settings(model_settings: Any) -> Optional[dict]: if len(json.dumps(filtered, separators=(",", ":"))) > _MAX_UPLOADED_SETTINGS_BYTES: logger.warning( "Benchmark: model settings snapshot exceeded " - f"{_MAX_UPLOADED_SETTINGS_BYTES} bytes, uploading accelerator flags only" + f"{_MAX_UPLOADED_SETTINGS_BYTES} bytes, uploading accelerator provenance only" ) filtered = { spec.attr: filtered[spec.attr] for spec in _FEATURE_FLAG_SPECS if spec.attr in filtered } + if "turboquant_mid_prefill" in raw: + filtered["turboquant_mid_prefill"] = bool(raw["turboquant_mid_prefill"]) return filtered diff --git a/omlx/admin/i18n/en.json b/omlx/admin/i18n/en.json index 036d6140d..bf1db36cf 100644 --- a/omlx/admin/i18n/en.json +++ b/omlx/admin/i18n/en.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Truncate large tool results (e.g. file reads) to a token limit", "modal.model_settings.limit_tool_placeholder": "e.g. 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV Cache", - "modal.model_settings.turboquant_kv_hint": "Compress KV cache using vector quantization. Lower bits = more compression, higher bits = better quality. Supports 2 to 8 bits.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits per channel", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/es.json b/omlx/admin/i18n/es.json index 7c6ff757d..0b095b73d 100644 --- a/omlx/admin/i18n/es.json +++ b/omlx/admin/i18n/es.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Recortar resultados grandes de herramientas (ej. lecturas de archivos) a un límite de tokens", "modal.model_settings.limit_tool_placeholder": "ej. 2000", "modal.model_settings.turboquant_kv": "Caché KV TurboQuant", - "modal.model_settings.turboquant_kv_hint": "Comprimir caché KV usando cuantización vectorial. Bits más bajos = más compresión, bits más altos = mejor calidad. Soporta de 2 a 8 bits.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits por canal", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/fr.json b/omlx/admin/i18n/fr.json index 9124fef9a..3e65961d2 100644 --- a/omlx/admin/i18n/fr.json +++ b/omlx/admin/i18n/fr.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Tronquez les résultats d'outil volumineux (par ex. lectures de fichiers) à une limite de tokens", "modal.model_settings.limit_tool_placeholder": "par ex. 2000", "modal.model_settings.turboquant_kv": "Cache KV TurboQuant", - "modal.model_settings.turboquant_kv_hint": "Compressez le cache KV en utilisant la quantification vectorielle. Bits inférieurs = plus de compression, bits supérieurs = meilleure qualité. Supporte 2 à 8 bits.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits par canal", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/ja.json b/omlx/admin/i18n/ja.json index bc4eb963e..1fc1aa84b 100644 --- a/omlx/admin/i18n/ja.json +++ b/omlx/admin/i18n/ja.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "大きなツール結果（ファイル読み込みなど）をトークン制限に合わせて切り詰めます", "modal.model_settings.limit_tool_placeholder": "例: 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV Cache", - "modal.model_settings.turboquant_kv_hint": "Compress KV cache using vector quantization. Reduces memory ~60-75% with near-lossless quality for long context.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits per channel", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/ko.json b/omlx/admin/i18n/ko.json index d1ab9c9d6..c5b258fbf 100644 --- a/omlx/admin/i18n/ko.json +++ b/omlx/admin/i18n/ko.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Tool calling(예: 파일 읽기) 결과를 정해진 크기로 자릅니다", "modal.model_settings.limit_tool_placeholder": "예: 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV Cache", - "modal.model_settings.turboquant_kv_hint": "Compress KV cache using vector quantization. Reduces memory ~60-75% with near-lossless quality for long context.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits per channel", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/pt-BR.json b/omlx/admin/i18n/pt-BR.json index 185a54ead..e991c0af8 100644 --- a/omlx/admin/i18n/pt-BR.json +++ b/omlx/admin/i18n/pt-BR.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Truncar resultados grandes de ferramentas (ex: leituras de arquivo) para um limite de tokens", "modal.model_settings.limit_tool_placeholder": "ex: 2000", "modal.model_settings.turboquant_kv": "Cache KV TurboQuant", - "modal.model_settings.turboquant_kv_hint": "Comprimir o cache KV usando quantização vetorial. Menos bits = mais compressão, mais bits = melhor qualidade. Suporta de 2 a 8 bits.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Bits por canal", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Drafts several tokens per step with the model's built-in MTP head. Up to ~1.5x faster decoding. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/ru.json b/omlx/admin/i18n/ru.json index b5f01f3b0..82b9ab942 100644 --- a/omlx/admin/i18n/ru.json +++ b/omlx/admin/i18n/ru.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "Обрезать большие результаты инструментов, например чтение файлов, до лимита токенов", "modal.model_settings.limit_tool_placeholder": "например, 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV-кэш", - "modal.model_settings.turboquant_kv_hint": "Сжимайте KV-кэш с помощью векторного квантования. Меньшее число бит = сильнее сжатие, большее = лучше качество. Поддерживает от 2 до 8 бит.", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "Битов на канал", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "Использует встроенную MTP-голову модели для черновой генерации нескольких токенов за шаг. Декодирование до ~1,5× быстрее. (Qwen 3.5/3.6, DeepSeek-V4, GLM-5.2)", diff --git a/omlx/admin/i18n/zh-TW.json b/omlx/admin/i18n/zh-TW.json index b270e0a8a..4eed7d94d 100644 --- a/omlx/admin/i18n/zh-TW.json +++ b/omlx/admin/i18n/zh-TW.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "將過長的回傳結果（如檔案讀取）截斷至指定 Token 數量上限", "modal.model_settings.limit_tool_placeholder": "例如 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV Cache", - "modal.model_settings.turboquant_kv_hint": "使用向量量化壓縮 KV 快取。較低位元 = 更多壓縮，較高位元 = 更佳品質。支援 2 到 8 位元。", + "modal.model_settings.turboquant_kv_hint": "Compress the KV cache with vector quantization after prefill for generation. Ordinary TurboQuant does not convert during cold prefill. Lower bits use less memory; higher bits preserve more quality.", + "modal.model_settings.turboquant_mid_prefill": "Convert under prefill pressure", + "modal.model_settings.turboquant_mid_prefill_hint": "When a full prefill chunk cannot fit, convert the growing KV cache once and continue with TurboQuant. Requires this to be the only loaded model. Adds a one-time pause and may slow the rest of prefill.", "modal.model_settings.turboquant_kv_bits_label": "每通道位元數", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "使用模型內建的 MTP 頭，每步產生多個 token。解碼速度最快可達約 1.5 倍。（Qwen 3.5/3.6、DeepSeek-V4、GLM-5.2）", diff --git a/omlx/admin/i18n/zh.json b/omlx/admin/i18n/zh.json index 234ae23f9..3faf888ae 100644 --- a/omlx/admin/i18n/zh.json +++ b/omlx/admin/i18n/zh.json @@ -559,7 +559,9 @@ "modal.model_settings.limit_tool_result_hint": "截断过长的工具结果（如文件读取）到 Token 限制", "modal.model_settings.limit_tool_placeholder": "例如 2000", "modal.model_settings.turboquant_kv": "TurboQuant KV 缓存", - "modal.model_settings.turboquant_kv_hint": "使用向量量化压缩 KV Cache。位数越低压缩越多，位数越高质量越好。支持 2 到 8 位。", + "modal.model_settings.turboquant_kv_hint": "在预填充完成后，使用向量量化压缩生成阶段的 KV 缓存。普通 TurboQuant 不会在冷预填充期间转换缓存。位数越低，内存占用越少；位数越高，质量保留越好。", + "modal.model_settings.turboquant_mid_prefill": "预填充内存压力下转换", + "modal.model_settings.turboquant_mid_prefill_hint": "当完整的预填充块无法放入内存时，将增长中的 KV 缓存转换一次，然后使用 TurboQuant 继续预填充。此功能要求它是唯一已加载的模型，会造成一次暂停，并可能减慢剩余的预填充。", "modal.model_settings.turboquant_kv_bits_label": "每通道位数", "modal.model_settings.lightning_mtp": "Lightning MTP", "modal.model_settings.lightning_mtp_hint": "使用模型内置 MTP 头每步预测多个 Token。解码速度提升约 1.5 倍。（Qwen 3.5/3.6、DeepSeek-V4、GLM-5.2）", diff --git a/omlx/admin/oq_manager.py b/omlx/admin/oq_manager.py index 835a7395b..713779b40 100644 --- a/omlx/admin/oq_manager.py +++ b/omlx/admin/oq_manager.py @@ -12,21 +12,38 @@ import logging import time import uuid +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path -from typing import Callable, Optional +from typing import Optional try: import mlx.core as mx - HAS_MLX = True + HAS_MLX = mx is not None except ImportError: HAS_MLX = False from ..model_discovery import _has_vision_subconfig +from ..utils.metal_sync import _conversion_coordinator, _sync_and_clear_cache logger = logging.getLogger(__name__) +def _run_process_guarded_metal( + fn: Callable[..., object], + /, + *args: object, + **kwargs: object, +) -> object: + """Run independent oQ Metal work only when no mid-prefill owner exists.""" + with _conversion_coordinator.background_metal_operation(): + if HAS_MLX: + _sync_and_clear_cache() + try: + return fn(*args, **kwargs) + finally: + if HAS_MLX: + _sync_and_clear_cache() class _QuantCancelled(Exception): """Raised by progress callback when task is cancelled.""" @@ -479,15 +496,6 @@ async def cancel_quantization(self, task_id: str) -> bool: except (asyncio.CancelledError, Exception): pass - # GPU cleanup after thread is done - if HAS_MLX: - for _attempt in range(3): - try: - mx.synchronize() - mx.clear_cache() - break - except Exception: - await asyncio.sleep(1.0) logger.info(f"oQ quantization cancelled: {task.model_name} (task_id={task_id})") return True @@ -525,16 +533,6 @@ async def _run_quantization(self, task_id: str) -> None: if task_id in self._cancelled: return - # Ensure GPU is clean before starting (previous task may have been cancelled) - # Metal command buffers need full sync + cache clear after cancellation - if HAS_MLX: - for _ in range(3): - try: - mx.synchronize() - mx.clear_cache() - break - except Exception: - await asyncio.sleep(1.0) # Phase 1: Loading task.status = QuantStatus.LOADING @@ -571,6 +569,7 @@ def _progress_cb( from ..oq import quantize_oq_streaming await asyncio.to_thread( + _run_process_guarded_metal, quantize_oq_streaming, task.model_path, task.output_path, @@ -600,6 +599,7 @@ def _progress_cb( _progress_cb("saving", 97.0, "Merging MTP head...") await asyncio.to_thread( + _run_process_guarded_metal, combine_mtp_into_output, task.output_path, task.mtp_assistant_model_path, diff --git a/omlx/admin/routes.py b/omlx/admin/routes.py index e41f168dc..f73df7619 100644 --- a/omlx/admin/routes.py +++ b/omlx/admin/routes.py @@ -126,6 +126,7 @@ class ModelSettingsRequest(BaseModel): thinking_budget_tokens: int | None = None # TurboQuant KV cache (mlx-vlm backend) turboquant_kv_enabled: bool | None = None + turboquant_mid_prefill: bool | None = None turboquant_kv_bits: float | None = None # SpecPrefill (experimental) specprefill_enabled: bool | None = None @@ -516,6 +517,7 @@ def _sanitize_diffusion_settings_dict(settings: dict) -> None: settings["thinking_budget_enabled"] = False settings["guided_grammar_enabled"] = False settings["turboquant_kv_enabled"] = False + settings["turboquant_mid_prefill"] = False settings["turboquant_kv_bits"] = 4 settings["turboquant_skip_last"] = True settings["specprefill_enabled"] = False @@ -591,6 +593,7 @@ def _sanitize_diffusion_model_settings(settings) -> None: settings.index_cache_freq = None settings.turboquant_kv_enabled = False + settings.turboquant_mid_prefill = False settings.turboquant_kv_bits = 4 settings.turboquant_skip_last = True settings.specprefill_enabled = False @@ -2238,6 +2241,10 @@ async def update_model_settings( # TurboQuant KV cache settings if "turboquant_kv_enabled" in sent: current_settings.turboquant_kv_enabled = request.turboquant_kv_enabled or False + if "turboquant_mid_prefill" in sent: + current_settings.turboquant_mid_prefill = ( + request.turboquant_mid_prefill or False + ) if "turboquant_kv_bits" in sent: current_settings.turboquant_kv_bits = request.turboquant_kv_bits or 4 # SpecPrefill settings diff --git a/omlx/admin/static/js/dashboard.js b/omlx/admin/static/js/dashboard.js index 2afb9e610..a477d61e6 100644 --- a/omlx/admin/static/js/dashboard.js +++ b/omlx/admin/static/js/dashboard.js @@ -25,6 +25,7 @@ 'max_tool_result_tokens', 'index_cache_freq', 'turboquant_kv_enabled', + 'turboquant_mid_prefill', 'turboquant_kv_bits', 'turboquant_skip_last', 'specprefill_enabled', @@ -1156,6 +1157,10 @@ } continue; } + if (k === 'turboquant_mid_prefill') { + out.turboquant_mid_prefill = !!ms.turboquant_mid_prefill; + continue; + } if (k === 'guided_grammar_enabled') { out.guided_grammar_enabled = !!ms.guided_grammar_enabled; continue; @@ -1516,6 +1521,7 @@ enableIndexCache: !!(s.index_cache_freq), index_cache_freq: s.index_cache_freq || null, turboquant_kv_enabled: s.turboquant_kv_enabled || false, + turboquant_mid_prefill: s.turboquant_mid_prefill || false, turboquant_kv_bits: s.turboquant_kv_bits || 4, specprefill_enabled: s.specprefill_enabled || false, specprefill_draft_model: s.specprefill_draft_model || '', @@ -2033,6 +2039,7 @@ forced_ct_kwargs: forcedCtKwargs.length > 0 ? forcedCtKwargs : null, turboquant_kv_enabled: this.modelSettings.turboquant_kv_enabled, + turboquant_mid_prefill: !!this.modelSettings.turboquant_mid_prefill, turboquant_kv_bits: this.modelSettings.turboquant_kv_enabled ? (parseFloat(this.modelSettings.turboquant_kv_bits) || 4) : 4, @@ -2117,6 +2124,7 @@ guided_grammar: null, max_tool_result_tokens: 0, turboquant_kv_enabled: false, + turboquant_mid_prefill: false, turboquant_kv_bits: 4, specprefill_enabled: false, specprefill_draft_model: null, @@ -2205,6 +2213,7 @@ this.modelSettings.max_tool_result_tokens = 0; this.modelSettings.ctKwargEntries = []; this.modelSettings.turboquant_kv_enabled = false; + this.modelSettings.turboquant_mid_prefill = false; this.modelSettings.turboquant_kv_bits = 4; this.modelSettings.specprefill_enabled = false; this.modelSettings.specprefill_draft_model = null; diff --git a/omlx/admin/templates/dashboard/_modal_model_settings.html b/omlx/admin/templates/dashboard/_modal_model_settings.html index 88c1cd1ed..01b17bb73 100644 --- a/omlx/admin/templates/dashboard/_modal_model_settings.html +++ b/omlx/admin/templates/dashboard/_modal_model_settings.html @@ -695,10 +695,14 @@

{{
- {{ t('modal.model_settings.turboquant_kv') }} -

{{ t('modal.model_settings.turboquant_kv_hint') }}

+ {{ t('modal.model_settings.turboquant_kv') }} +

{{ t('modal.model_settings.turboquant_kv_hint') }}

-
- - @@ -723,6 +729,24 @@

{{ +
+
+ {{ t('modal.model_settings.turboquant_mid_prefill') }} +

{{ t('modal.model_settings.turboquant_mid_prefill_hint') }}

+
+ +

diff --git a/omlx/cache/paged_ssd_cache.py b/omlx/cache/paged_ssd_cache.py index 57a2fb7ac..9c06d9e8d 100644 --- a/omlx/cache/paged_ssd_cache.py +++ b/omlx/cache/paged_ssd_cache.py @@ -1491,6 +1491,22 @@ def __init__( f"existing_files={self._index.count}{disk_info}" ) + def set_expected_block_payload_bytes(self, payload_bytes: int) -> None: + """Resize the writer queue from one conservative serialized block size.""" + if payload_bytes <= 0: + raise ValueError("expected block payload must be positive") + block_tokens = max(1, int(self._expected_block_size_tokens)) + per_token = max(1, (int(payload_bytes) + block_tokens - 1) // block_tokens) + cap = _compute_max_pending_writes( + block_size_tokens=block_tokens, + kv_bytes_per_token=per_token, + ) + with self._write_queue.mutex: + self._expected_kv_bytes_per_token = per_token + self._max_pending_writes = cap + self._write_queue.maxsize = cap + self._write_queue.not_full.notify_all() + # --- Hot cache helpers --- @staticmethod @@ -2289,7 +2305,6 @@ def save_block( file_path = self._get_file_path(block_hash) try: - # Prepare arrays for safetensors. Three layer_data shapes are # accepted: # - ``('__nstate__', class_name, [elem0, elem1, ...])`` — V3 @@ -2305,9 +2320,7 @@ def save_block( # etc.) has been migrated to emit ``__nstate__`` markers yet. arrays = {} has_pooling_cache_delta = False - cache_list_meta = ( - {} - ) # Per-layer sidecar metadata (sub_count, state_count, etc.) + cache_list_meta = {} # Per-layer sidecar metadata (sub_count, state_count, etc.) # Shim; module-level to avoid a recursive-closure refcount # cycle pinning `arrays` — see _store_nstate_elements_flat. @@ -2987,8 +3000,7 @@ def load_block_with_metadata( self._stats["hits"] += 1 self._stats["hot_cache_hits"] += 1 logger.debug( - f"Loaded block with metadata from hot cache: " - f"{block_hash.hex()[:16]}..." + f"Loaded block with metadata from hot cache: {block_hash.hex()[:16]}..." ) return cache_data, metadata_dict @@ -3326,17 +3338,13 @@ def set_expected_layer_signature( new_signature = list(layer_cache_types) new_canonical = _canonicalize_layer_cache_types(new_signature) - new_bits = ( - float(turboquant_kv_bits) if turboquant_kv_bits is not None else None - ) + new_bits = float(turboquant_kv_bits) if turboquant_kv_bits is not None else None with self._lock: old_signature = self._expected_layer_cache_types old_canonical = _canonicalize_layer_cache_types(old_signature) bits_changed = new_bits != self._expected_turboquant_kv_bits - subtypes_changed = ( - cachelist_subtypes != self._expected_cachelist_subtypes - ) + subtypes_changed = cachelist_subtypes != self._expected_cachelist_subtypes if ( old_canonical == new_canonical and not bits_changed diff --git a/omlx/engine/base.py b/omlx/engine/base.py index 5419b9abd..8a1255025 100644 --- a/omlx/engine/base.py +++ b/omlx/engine/base.py @@ -71,7 +71,7 @@ async def _run_scheduler_preflight_with_cleanup_retry( request_id: str | None, eviction_callback: Any | None, executor: Any | None = None, -) -> None: +) -> bool: """Run route preflight after transient post-request cleanup settles. A finished request can remain resident while its asynchronous cache store @@ -86,6 +86,7 @@ async def _run_scheduler_preflight_with_cleanup_retry( """ deadline = time.monotonic() + _PREFLIGHT_CLEANUP_WAIT_TIMEOUT_S waited_for_cleanup = False + callback_attempted = False while True: eviction_request = scheduler.preflight_eviction_request( @@ -97,7 +98,7 @@ async def _run_scheduler_preflight_with_cleanup_retry( num_prompt_tokens=num_prompt_tokens, request_id=request_id, ) - return + return callback_attempted cleanup_pending_fn = getattr( scheduler, "has_pending_route_preflight_cleanup", None @@ -146,12 +147,13 @@ async def _run_scheduler_preflight_with_cleanup_retry( "Running preflight LRU eviction for request %s", eviction_request.request_id, ) + callback_attempted = True await eviction_callback(eviction_request) scheduler.preflight_or_raise( num_prompt_tokens=num_prompt_tokens, request_id=request_id, ) - return + return callback_attempted @dataclass @@ -410,7 +412,7 @@ async def preflight_chat( tools: Optional[list] = None, request_id: Optional[str] = None, **kwargs, - ) -> None: + ) -> bool | None: """Optional prefill-memory preflight check for chat requests. Default no-op; engines that implement the prefill memory guard @@ -418,6 +420,10 @@ async def preflight_chat( actual estimation logic. The base no-op lets simpler engines (SimpleEngine, embedding/reranker engines, test stubs) be invoked from the server endpoints without additional wrapping. + + An implementation may return ``True`` when route preflight actually + invoked an eviction callback and admission must consume that request's + retry budget. ``None`` retains the legacy no-op contract. """ return None @@ -426,10 +432,12 @@ async def preflight_completion( prompt: str, request_id: Optional[str] = None, **kwargs, - ) -> None: + ) -> bool | None: """Optional prefill-memory preflight check for completion requests. - See :meth:`preflight_chat` for the rationale. + Implementations may report a route-level eviction callback attempt + using the same request-local contract as :meth:`preflight_chat`. + ``None`` retains the legacy no-op behavior. """ return None diff --git a/omlx/engine/batched.py b/omlx/engine/batched.py index 8eaaf2e54..bb334d827 100644 --- a/omlx/engine/batched.py +++ b/omlx/engine/batched.py @@ -25,6 +25,24 @@ logger = logging.getLogger(__name__) +def _configure_turboquant_scheduler(scheduler: Any, model_settings: Any | None) -> None: + """Propagate effective TurboQuant settings to a newly-created scheduler.""" + turboquant_active = bool(getattr(model_settings, "turboquant_kv_enabled", False)) + scheduler._turboquant_mid_prefill = turboquant_active and bool( + getattr(model_settings, "turboquant_mid_prefill", False) + ) + if not turboquant_active: + return + + scheduler._turboquant_kv_bits = float( + getattr(model_settings, "turboquant_kv_bits", 4) + ) + scheduler._turboquant_skip_last = getattr( + model_settings, "turboquant_skip_last", True + ) + scheduler._set_model_info_for_monitor() + + # Optional Harmony adapter import try: from ..adapter.harmony import preprocess_harmony_messages @@ -85,8 +103,8 @@ async def _preflight_or_raise_with_eviction( *, num_prompt_tokens: int, request_id: str | None, - ) -> None: - await _run_scheduler_preflight_with_cleanup_retry( + ) -> bool: + callback_attempted = await _run_scheduler_preflight_with_cleanup_retry( scheduler, num_prompt_tokens=num_prompt_tokens, request_id=request_id, @@ -97,6 +115,7 @@ async def _preflight_or_raise_with_eviction( None, ), ) + return callback_attempted @property def model_name(self) -> str: @@ -433,17 +452,12 @@ def _load_model_sync(): await self._engine.engine.start() - # TurboQuant KV cache: propagate bits to scheduler + # TurboQuant KV cache: propagate effective settings to scheduler scheduler = self._engine.engine.scheduler - if self._model_settings is not None: - tq_enabled = getattr(self._model_settings, "turboquant_kv_enabled", False) - if tq_enabled: - tq_bits = float(getattr(self._model_settings, "turboquant_kv_bits", 4)) - scheduler._turboquant_kv_bits = tq_bits - scheduler._turboquant_skip_last = getattr( - self._model_settings, "turboquant_skip_last", True - ) - scheduler._set_model_info_for_monitor() + scheduler._prefill_eviction_callback_configured = ( + self._prefill_eviction_callback is not None + ) + _configure_turboquant_scheduler(scheduler, self._model_settings) scheduler.refresh_ssd_layer_signature() # SpecPrefill: load draft model and pass to scheduler @@ -737,11 +751,15 @@ async def generate( # stream_generate so the non-streaming path is not silently ignored. specprefill_kwargs = self._pop_specprefill_kwargs(kwargs) tools = kwargs.pop("tools", None) + prefill_eviction_callback_attempted = bool( + kwargs.pop("prefill_eviction_callback_attempted", False) + ) output = await self._engine.generate( prompt=prompt, sampling_params=sampling_params, tools=tools, + prefill_eviction_callback_attempted=prefill_eviction_callback_attempted, **specprefill_kwargs, ) @@ -813,6 +831,9 @@ async def stream_generate( # SpecPrefill: pass per-request overrides to engine specprefill_kwargs = self._pop_specprefill_kwargs(kwargs) tools = kwargs.pop("tools", None) + prefill_eviction_callback_attempted = bool( + kwargs.pop("prefill_eviction_callback_attempted", False) + ) engine = self._engine request_id = await engine.add_request( @@ -820,6 +841,7 @@ async def stream_generate( sampling_params=sampling_params, tools=tools, skip_cache_store=bool(kwargs.get("skip_cache_store", False)), + prefill_eviction_callback_attempted=prefill_eviction_callback_attempted, **specprefill_kwargs, ) @@ -938,7 +960,7 @@ async def preflight_chat( tools: list[dict] | None = None, request_id: str | None = None, **kwargs, - ) -> None: + ) -> bool | None: """Early prefill memory check for chat completions. Tokenizes the templated prompt and asks the scheduler whether the @@ -987,7 +1009,7 @@ async def preflight_chat( if scheduler is None: _warn_scheduler_unreachable_once(self, "preflight_chat") return - await self._preflight_or_raise_with_eviction( + return await self._preflight_or_raise_with_eviction( scheduler, num_prompt_tokens=num_tokens, request_id=request_id ) @@ -996,7 +1018,7 @@ async def preflight_completion( prompt: str, request_id: str | None = None, **kwargs, - ) -> None: + ) -> bool | None: """Early prefill memory check for plain /v1/completions calls. See ``preflight_chat`` for the rationale. @@ -1017,7 +1039,7 @@ async def preflight_completion( if scheduler is None: _warn_scheduler_unreachable_once(self, "preflight_completion") return - await self._preflight_or_raise_with_eviction( + return await self._preflight_or_raise_with_eviction( scheduler, num_prompt_tokens=num_tokens, request_id=request_id ) diff --git a/omlx/engine/dflash.py b/omlx/engine/dflash.py index bedaf30c0..34cec5063 100644 --- a/omlx/engine/dflash.py +++ b/omlx/engine/dflash.py @@ -119,9 +119,7 @@ def _ms(key: str) -> str: # Precision suffixes Poolside uses for drafts trained against a quantized # target ("Laguna-S-2.1-DFlash-NVFP4" etc.). Anything else after "-DFlash-" # (e.g. z-lab's "-b16" block-size suffix) makes no precision claim. -_DRAFT_PRECISION_TAGS = frozenset( - {"NVFP4", "INT4", "INT8", "FP8", "FP4", "MXFP4"} -) +_DRAFT_PRECISION_TAGS = frozenset({"NVFP4", "INT4", "INT8", "FP8", "FP4", "MXFP4"}) def _canonical_precision_tag(value: object) -> str | None: @@ -298,6 +296,8 @@ def __init__( fallback_engine_type: str = "batched", scheduler_config: Any | None = None, omlx_ssd_cache_dir: str | Path | None = None, + prefill_eviction_callback: Any | None = None, + mid_prefill_process_claim_callback: Any | None = None, ): super().__init__() self._model_name = model_name @@ -317,6 +317,8 @@ def __init__( self._omlx_ssd_cache_dir = ( Path(omlx_ssd_cache_dir) if omlx_ssd_cache_dir else None ) + self._prefill_eviction_callback = prefill_eviction_callback + self._mid_prefill_process_claim_callback = mid_prefill_process_claim_callback self._target_model = None self._target_ops = None @@ -725,6 +727,21 @@ def _end_runtime_cache_request(manager: Any | None) -> None: except Exception as exc: logger.debug(f"dflash cache end_request failed: {exc}") + async def _claim_fallback_mid_prefill_process(self) -> None: + scheduler = self.scheduler + if ( + scheduler is None + or getattr(scheduler, "_turboquant_mid_prefill", False) is not True + ): + return + claim_callback = self._mid_prefill_process_claim_callback + if not callable(claim_callback): + raise RuntimeError( + "TurboQuant mid-prefill DFlash fallback requires an " + "EnginePool-owned process-exclusive Metal lane" + ) + await claim_callback(scheduler) + async def _evict_dflash_and_start_fallback(self) -> None: """Evict dflash models from memory, verify release, then start fallback engine.""" from dflash_mlx.cache.manager import shutdown_runtime_cache_manager @@ -788,24 +805,53 @@ async def _evict_dflash_and_start_fallback(self) -> None: else: logger.warning("DFlash model eviction: memory settle timed out") - # Start fallback engine + # Start fallback engine. if self._fallback_engine_type == "vlm": from .vlm import VLMBatchedEngine - self._fallback_engine = VLMBatchedEngine( + fallback_engine: BaseEngine = VLMBatchedEngine( model_name=self._model_name, scheduler_config=self._scheduler_config, model_settings=self._model_settings, + prefill_eviction_callback=self._prefill_eviction_callback, ) else: from .batched import BatchedEngine - self._fallback_engine = BatchedEngine( + fallback_engine = BatchedEngine( model_name=self._model_name, scheduler_config=self._scheduler_config, model_settings=self._model_settings, + prefill_eviction_callback=self._prefill_eviction_callback, ) - await self._fallback_engine.start() + self._fallback_engine = fallback_engine + try: + await fallback_engine.start() + await self._claim_fallback_mid_prefill_process() + except Exception: + scheduler = self.scheduler + process_owner = getattr(scheduler, "_metal_process_owner", None) + try: + await fallback_engine.stop() + except Exception: + logger.warning( + "Failed to stop rejected DFlash fallback engine", + exc_info=True, + ) + finally: + close_process_owner = getattr(process_owner, "close", None) + if callable(close_process_owner): + try: + close_process_owner() + except Exception: + logger.warning( + "Failed to close rejected DFlash fallback owner", + exc_info=True, + ) + self._fallback_engine = None + self._in_fallback_mode = False + self._loaded = False + raise self._in_fallback_mode = True logger.info(f"DFlash fallback engine started: {self._fallback_engine_type}") @@ -930,7 +976,7 @@ async def preflight_chat( tools: list | None = None, request_id: str | None = None, **kwargs, - ) -> None: + ) -> bool | None: """Prefill-memory preflight for chat requests. DFlash bypasses the scheduler, so it implements the front-door guard @@ -943,10 +989,9 @@ async def preflight_chat( if not self._loaded: await self.start() if self._in_fallback_mode and self._fallback_engine is not None: - await self._fallback_engine.preflight_chat( + return await self._fallback_engine.preflight_chat( messages, tools=tools, request_id=request_id, **kwargs ) - return if self._prefill_guard is None: _warn_scheduler_unreachable_once( self, "preflight_chat", "primary-mode prefill guard unavailable" @@ -981,15 +1026,14 @@ async def preflight_completion( prompt: str, request_id: str | None = None, **kwargs, - ) -> None: + ) -> bool | None: """Prefill-memory preflight for plain completions. See ``preflight_chat``.""" if not self._loaded: await self.start() if self._in_fallback_mode and self._fallback_engine is not None: - await self._fallback_engine.preflight_completion( + return await self._fallback_engine.preflight_completion( prompt, request_id=request_id, **kwargs ) - return if self._prefill_guard is None: _warn_scheduler_unreachable_once( self, diff --git a/omlx/engine/vlm.py b/omlx/engine/vlm.py index 75d3c7d49..af0b0f8bc 100644 --- a/omlx/engine/vlm.py +++ b/omlx/engine/vlm.py @@ -62,6 +62,25 @@ logger = logging.getLogger(__name__) + +def _configure_turboquant_scheduler(scheduler: Any, model_settings: Any | None) -> None: + """Propagate effective TurboQuant settings to a newly-created scheduler.""" + turboquant_active = bool(getattr(model_settings, "turboquant_kv_enabled", False)) + scheduler._turboquant_mid_prefill = turboquant_active and bool( + getattr(model_settings, "turboquant_mid_prefill", False) + ) + if not turboquant_active: + return + + scheduler._turboquant_kv_bits = float( + getattr(model_settings, "turboquant_kv_bits", 4) + ) + scheduler._turboquant_skip_last = getattr( + model_settings, "turboquant_skip_last", True + ) + scheduler._set_model_info_for_monitor() + + # OCR model types that require special handling. # unlimited-ocr keeps its dashed config model_type (mlx-vlm resolves it to the # unlimited_ocr package via MODEL_REMAPPING), so key it in the dashed form to @@ -1314,8 +1333,8 @@ async def _preflight_or_raise_with_eviction( *, num_prompt_tokens: int, request_id: str | None, - ) -> None: - await _run_scheduler_preflight_with_cleanup_retry( + ) -> bool: + callback_attempted = await _run_scheduler_preflight_with_cleanup_retry( scheduler, num_prompt_tokens=num_prompt_tokens, request_id=request_id, @@ -1326,6 +1345,7 @@ async def _preflight_or_raise_with_eviction( None, ), ) + return callback_attempted @property def model_name(self) -> str: @@ -1700,21 +1720,19 @@ def _load_vlm_sync(): # TurboQuant KV cache scheduler = self._engine.engine.scheduler - if self._model_settings is not None: - tq_enabled = getattr(self._model_settings, "turboquant_kv_enabled", False) - if tq_enabled: - from ..patches.turboquant_attention import ( - apply_turboquant_attention_patch, - ) + scheduler._prefill_eviction_callback_configured = ( + self._prefill_eviction_callback is not None + ) + tq_enabled = bool(getattr(self._model_settings, "turboquant_kv_enabled", False)) + if tq_enabled: + from ..patches.turboquant_attention import ( + apply_turboquant_attention_patch, + ) - apply_turboquant_attention_patch() - tq_bits = float(getattr(self._model_settings, "turboquant_kv_bits", 4)) - scheduler._turboquant_kv_bits = tq_bits - scheduler._turboquant_skip_last = getattr( - self._model_settings, "turboquant_skip_last", True - ) - scheduler._set_model_info_for_monitor() - logger.info(f"TurboQuant KV cache enabled for VLM: {tq_bits} bits") + apply_turboquant_attention_patch() + tq_bits = float(getattr(self._model_settings, "turboquant_kv_bits", 4)) + logger.info(f"TurboQuant KV cache enabled for VLM: {tq_bits} bits") + _configure_turboquant_scheduler(scheduler, self._model_settings) # head_dim=256 long-context prefill -> O(L) tiled SDPA kernel. See # batched.py for rationale. Passthrough-safe; strictly gated route. @@ -3113,6 +3131,9 @@ async def generate( # SpecPrefill: forward per-request overrides to the engine, mirroring # stream_generate so the non-streaming path is not silently ignored. specprefill_kwargs = self._pop_specprefill_kwargs(kwargs) + prefill_eviction_callback_attempted = bool( + kwargs.pop("prefill_eviction_callback_attempted", False) + ) output = await self._engine.generate( prompt=prompt, @@ -3122,6 +3143,7 @@ async def generate( vlm_image_hash=vlm_image_hash, vlm_cache_key_start=vlm_cache_key_start, vlm_cache_key_ranges=vlm_cache_key_ranges, + prefill_eviction_callback_attempted=prefill_eviction_callback_attempted, **specprefill_kwargs, ) @@ -3222,6 +3244,9 @@ async def stream_generate( # SpecPrefill: pass per-request overrides specprefill_kwargs = self._pop_specprefill_kwargs(kwargs) + prefill_eviction_callback_attempted = bool( + kwargs.pop("prefill_eviction_callback_attempted", False) + ) engine = self._engine request_id = await engine.add_request( @@ -3233,6 +3258,7 @@ async def stream_generate( vlm_cache_key_start=vlm_cache_key_start, vlm_cache_key_ranges=vlm_cache_key_ranges, skip_cache_store=bool(kwargs.get("skip_cache_store", False)), + prefill_eviction_callback_attempted=prefill_eviction_callback_attempted, **specprefill_kwargs, ) @@ -3349,7 +3375,7 @@ async def preflight_chat( tools: list[dict] | None = None, request_id: str | None = None, **kwargs, - ) -> None: + ) -> bool: """Early prefill memory check for chat completions (VLM path). The actual VLM prompt is built by ``_process_chat_messages`` → @@ -3381,6 +3407,10 @@ async def preflight_chat( ``BatchedEngine.preflight_chat`` for the upstream rationale (avoiding the ``StreamingResponse`` 200 commit so HTTP 400 actually reaches the client). + + Returns ``True`` only when this preflight actually invoked the async + eviction callback, allowing the route to consume the request's one + scheduler retry without storing cross-request state. """ if not self._loaded: await self.start() @@ -3392,7 +3422,7 @@ async def preflight_chat( stop=kwargs.get("stop"), kwargs=kwargs, ) - return + return False template_tools = convert_tools_for_template(tools) if tools else None ct_kwargs = kwargs.get("chat_template_kwargs") partial = kwargs.get("is_partial") @@ -3429,7 +3459,7 @@ async def preflight_chat( "surface the error", type(e).__name__, ) - return + return False # Count images from the ORIGINAL messages (the stripped # ``text_messages`` no longer has the image content-parts). num_tokens += _count_image_tokens_real( @@ -3442,8 +3472,8 @@ async def preflight_chat( scheduler = getattr(getattr(self._engine, "engine", None), "scheduler", None) if scheduler is None: _warn_scheduler_unreachable_once(self, "preflight_chat") - return - await self._preflight_or_raise_with_eviction( + return False + return await self._preflight_or_raise_with_eviction( scheduler, num_prompt_tokens=num_tokens, request_id=request_id ) @@ -3452,8 +3482,13 @@ async def preflight_completion( prompt: str, request_id: str | None = None, **kwargs, - ) -> None: - """Early prefill memory check for plain /v1/completions calls (VLM).""" + ) -> bool: + """Early prefill memory check for plain /v1/completions calls (VLM). + + Returns ``True`` only when this preflight completed after invoking the + async eviction callback. The route carries that request-local result to + admission so the scheduler cannot invoke the same callback twice. + """ if not self._loaded: await self.start() if self.is_diffusion_model: @@ -3461,7 +3496,7 @@ async def preflight_completion( stop=kwargs.get("stop"), kwargs=kwargs, ) - return + return False try: num_tokens = len(self._tokenizer.encode(prompt)) except Exception as e: @@ -3471,12 +3506,12 @@ async def preflight_completion( "path will surface the error", type(e).__name__, ) - return + return False scheduler = getattr(getattr(self._engine, "engine", None), "scheduler", None) if scheduler is None: _warn_scheduler_unreachable_once(self, "preflight_completion") - return - await self._preflight_or_raise_with_eviction( + return False + return await self._preflight_or_raise_with_eviction( scheduler, num_prompt_tokens=num_tokens, request_id=request_id ) diff --git a/omlx/engine_core.py b/omlx/engine_core.py index e025b316a..eac5ea1a2 100644 --- a/omlx/engine_core.py +++ b/omlx/engine_core.py @@ -50,6 +50,7 @@ ) from .utils.fatal import FATAL_TEARDOWN_TIMEOUT_S, fatal_exit from .utils.hardware import format_bytes +from .utils.metal_sync import _conversion_coordinator logger = logging.getLogger(__name__) @@ -128,6 +129,25 @@ def _init_mlx_thread() -> None: logger.info(f"MLX executor thread initialized: generation_stream = {stream}") +class _ProcessGuardedMLXExecutor(concurrent.futures.ThreadPoolExecutor): + """Reject global Metal work while a mid-prefill engine owns the process.""" + + @staticmethod + def _run_guarded(fn: Callable[..., Any], args: tuple, kwargs: dict) -> Any: + _conversion_coordinator.assert_background_metal_allowed() + return fn(*args, **kwargs) + + def submit( + self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any + ) -> concurrent.futures.Future[Any]: + return super().submit(self._run_guarded, fn, args, kwargs) + + +def _claim_turboquant_mid_prefill_process(owner: Any) -> None: + """Drain the global executor stream, then make ``owner`` process-exclusive.""" + mx.synchronize() + _conversion_coordinator.claim_process_exclusive(owner) + def get_mlx_executor() -> concurrent.futures.ThreadPoolExecutor: """Get or create the global MLX executor (lazy singleton). @@ -139,7 +159,7 @@ def get_mlx_executor() -> concurrent.futures.ThreadPoolExecutor: """ global _global_mlx_executor if _global_mlx_executor is None: - _global_mlx_executor = concurrent.futures.ThreadPoolExecutor( + _global_mlx_executor = _ProcessGuardedMLXExecutor( max_workers=1, thread_name_prefix="mlx-global", initializer=_init_mlx_thread, @@ -222,34 +242,60 @@ def __init__( self._engine_id = engine_id or str(uuid.uuid4()) self._owns_model = False self._closed = False - - # Acquire model ownership + self._metal_registration_active = False + self._mlx_executor: concurrent.futures.ThreadPoolExecutor | None = None + _conversion_coordinator.register_engine(self) + self._metal_registration_active = True + + # Acquire model ownership and finish construction while the process + # registration blocks a concurrent mid-prefill claim. Any constructor + # failure must release that registration immediately; a traceback can + # otherwise retain this half-built engine and block all later loads. registry = get_registry() - registry.acquire( - model=model, - engine=self, - engine_id=self._engine_id, - force=force_model_ownership, - ) - self._owns_model = True - - # Per-engine executor with dedicated mx.Stream (#1248). - # Each EngineCore gets its own thread + GPU stream so different - # models can run scheduler.step() concurrently. - self._mlx_stream = mx.new_thread_local_stream(mx.default_device()) - self._mlx_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=1, - thread_name_prefix=f"mlx-engine-{self._engine_id[:8]}", - ) + try: + registry.acquire( + model=model, + engine=self, + engine_id=self._engine_id, + force=force_model_ownership, + ) + self._owns_model = True + + # Per-engine executor with dedicated mx.Stream (#1248). + # Each EngineCore gets its own thread + GPU stream so different + # models can run scheduler.step() concurrently. + self._mlx_stream = mx.new_thread_local_stream(mx.default_device()) + self._mlx_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix=f"mlx-engine-{self._engine_id[:8]}", + ) - # Create scheduler with per-engine stream - scheduler_config = self.config.scheduler_config or SchedulerConfig() - self.scheduler = Scheduler( - model=model, - tokenizer=tokenizer, - config=scheduler_config, - stream=self._mlx_stream, - ) + # Create scheduler with per-engine stream + scheduler_config = self.config.scheduler_config or SchedulerConfig() + self.scheduler = Scheduler( + model=model, + tokenizer=tokenizer, + config=scheduler_config, + stream=self._mlx_stream, + ) + self.scheduler._metal_process_owner = self + except BaseException: + if self._mlx_executor is not None: + self._mlx_executor.shutdown(wait=False) + self._mlx_executor = None + if self._owns_model: + try: + registry.release(self.model, self._engine_id) + except Exception: + logger.warning( + "Engine %s: model release failed during initialization", + self._engine_id, + exc_info=True, + ) + self._owns_model = False + _conversion_coordinator.unregister_engine(self) + self._metal_registration_active = False + raise # Output collectors for low-latency streaming (vLLM pattern) self._output_collectors: Dict[str, RequestOutputCollector] = {} @@ -304,6 +350,15 @@ async def stop(self) -> None: self._loop = None logger.info("Engine stopped") + async def claim_turboquant_mid_prefill_process(self) -> None: + """Claim process-exclusive Metal access after all load work completes.""" + loop = asyncio.get_running_loop() + await loop.run_in_executor( + get_mlx_executor(), + _claim_turboquant_mid_prefill_process, + self, + ) + def is_running(self) -> bool: """Check if engine is running.""" return self._running @@ -546,6 +601,7 @@ async def add_request( specprefill_system_end: Optional[int] = None, skip_cache_store: bool = False, tools: list[dict[str, Any]] | None = None, + prefill_eviction_callback_attempted: bool = False, ) -> str: """ Add a request for processing. @@ -562,6 +618,8 @@ async def add_request( specprefill: Per-request SpecPrefill override (True/False/None) specprefill_keep_pct: Per-request keep rate override specprefill_threshold: Per-request threshold override (min tokens) + prefill_eviction_callback_attempted: Whether route preflight already + invoked the async LRU callback for this request. Returns: The request ID @@ -585,6 +643,7 @@ async def add_request( vlm_cache_key_start=vlm_cache_key_start, vlm_cache_key_ranges=vlm_cache_key_ranges, skip_cache_store=skip_cache_store, + prefill_eviction_retries=int(prefill_eviction_callback_attempted), ) # SpecPrefill: resolve per-request settings. @@ -1224,6 +1283,10 @@ def close(self) -> None: _immortal_mlx_streams.append(self._mlx_stream) self._mlx_executor = None + if getattr(self, "_metal_registration_active", False): + _conversion_coordinator.unregister_engine(self) + self._metal_registration_active = False + logger.debug(f"Engine {self._engine_id} closed") def __del__(self): diff --git a/omlx/engine_pool.py b/omlx/engine_pool.py index ee9880617..330bbf0c2 100644 --- a/omlx/engine_pool.py +++ b/omlx/engine_pool.py @@ -44,6 +44,7 @@ ModelNotFoundError, ModelTooLargeError, ModelUnavailableError, + TurboQuantProcessExclusiveError, describe_ceiling_binding, ) from .model_discovery import discover_models, format_size @@ -334,6 +335,10 @@ def add(key: str, value: object) -> None: turboquant_active = bool(data.get("turboquant_kv_enabled", False)) add("turboquant_kv_enabled", turboquant_active) if turboquant_active: + add( + "turboquant_mid_prefill", + bool(data.get("turboquant_mid_prefill", False)), + ) add("turboquant_kv_bits", data.get("turboquant_kv_bits", 4)) add("turboquant_skip_last", data.get("turboquant_skip_last", True)) @@ -818,6 +823,14 @@ async def get_engine( entry = self._entries.get(model_id) if not entry: raise ModelNotFoundError(model_id, list(self._entries.keys())) + try: + self._raise_if_other_mid_prefill_model_owns_process(model_id) + except TurboQuantProcessExclusiveError as exc: + raise ModelLoadingError( + model_id, + f"Model '{model_id}' cannot load while process-exclusive " + f"Metal access is unavailable: {exc}", + ) from exc expected_signature = self._engine_runtime_signature( model_id, runtime_settings, @@ -895,9 +908,7 @@ async def get_engine( # only its language weights, so admit it by the text-only estimate # instead of the vision-inclusive file size (#2385). admission_size = entry.estimated_size - if entry.text_only_size and ( - force_lm or entry.engine_type == "batched" - ): + if entry.text_only_size and (force_lm or entry.engine_type == "batched"): admission_size = entry.text_only_size ceiling = self._current_ceiling() @@ -907,9 +918,7 @@ async def get_engine( best_effort = ceiling > 0 if ceiling > 0: soft_target = self._admission_soft_target() - evict_target = ( - min(soft_target, ceiling) if soft_target > 0 else ceiling - ) + evict_target = min(soft_target, ceiling) if soft_target > 0 else ceiling evicted_any = unloaded_for_admission while True: # Consult the tracked accumulator alongside live memory: @@ -1200,6 +1209,53 @@ def _resolve_scheduler_from_engine(engine: object) -> object | None: except AttributeError: return None + def _raise_if_other_mid_prefill_model_owns_process( + self, + model_id: str, + ) -> None: + for other_id, entry in self._entries.items(): + if other_id == model_id or entry.engine is None: + continue + scheduler = self._resolve_scheduler_from_engine(entry.engine) + if ( + scheduler is not None + and getattr(scheduler, "_turboquant_mid_prefill", False) is True + ): + raise TurboQuantProcessExclusiveError( + "TurboQuant mid-prefill requires process-exclusive Metal " + f"access; unload '{other_id}' before loading '{model_id}'" + ) + + async def _claim_turboquant_mid_prefill_process( + self, + model_id: str, + scheduler: object, + ) -> None: + other_ids = [ + other_id + for other_id, entry in self._entries.items() + if other_id != model_id and entry.engine is not None + ] + if other_ids: + raise TurboQuantProcessExclusiveError( + "TurboQuant mid-prefill requires process-exclusive Metal " + "access; unload all other engines before enabling it " + f"(loaded: {', '.join(sorted(other_ids))})" + ) + + process_owner = getattr(scheduler, "_metal_process_owner", None) + claim_process = getattr( + process_owner, + "claim_turboquant_mid_prefill_process", + None, + ) + if not callable(claim_process): + raise RuntimeError( + "TurboQuant mid-prefill requires an EngineCore-owned " + "process-exclusive Metal lane" + ) + await claim_process() + def _is_idle_for_prefill_eviction(self, entry: EngineEntry) -> bool: engine = entry.engine if engine is None or entry.is_pinned or entry.is_loading or entry.in_use > 0: @@ -1429,11 +1485,28 @@ async def _unload_engine(self, model_id: str) -> None: logger.info(f"Unloading model: {model_id} (immediate abort)") pre_unload_active = mx.get_active_memory() + scheduler = self._resolve_scheduler_from_engine(entry.engine) + process_owner = ( + getattr(scheduler, "_metal_process_owner", None) + if scheduler is not None + else None + ) try: await entry.engine.stop() except Exception as e: logger.warning(f"Error stopping engine for {model_id}: {e}") + finally: + close_process_owner = getattr(process_owner, "close", None) + if callable(close_process_owner): + try: + close_process_owner() + except Exception: + logger.warning( + "Error closing engine owner for %s", + model_id, + exc_info=True, + ) # #1595: the immediate-abort stop() above tears the engine down without the normal # per-request completion callbacks, so a non-streaming engine's active_requests @@ -1607,13 +1680,18 @@ async def _reclaim() -> None: # barrier's small-model tolerance floor. target = pre_load_memory + 2 * 1024**3 current = 0 + blocked_rounds = 0 for _round in range(6): await asyncio.sleep(0.5 if _round == 0 else 1.0) gc.collect() - await loop.run_in_executor( - get_mlx_executor(), - lambda: (mx.synchronize(), mx.clear_cache()), - ) + try: + await loop.run_in_executor( + get_mlx_executor(), + lambda: (mx.synchronize(), mx.clear_cache()), + ) + except TurboQuantProcessExclusiveError: + blocked_rounds += 1 + continue current = max(mx.get_active_memory(), get_phys_footprint()) if current <= target: logger.info( @@ -1623,6 +1701,15 @@ async def _reclaim() -> None: ) self._wake_process_memory_enforcer() return + if blocked_rounds == 6: + logger.warning( + "Post-failed-load reclaim for '%s' remained blocked for " + "all 6 rounds by process-exclusive TurboQuant Metal " + "ownership; deferred buffers remain until exclusivity ends.", + model_id, + ) + self._wake_process_memory_enforcer() + return logger.warning( f"Post-failed-load reclaim for '{model_id}' did not settle: " f"current={format_size(current)} " @@ -1695,6 +1782,35 @@ async def _load_engine( # model families; let VLMBatchedEngine handle MTP-enabled VLMs. pass + async def prefill_eviction_callback( + eviction_request: object, + *, + _model_id: str = model_id, + ) -> bool: + return await self._evict_idle_lru_for_prefill( + exclude_model_id=_model_id, + eviction_request=eviction_request, + ) + + async def mid_prefill_process_claim_callback( + scheduler: object, + *, + _model_id: str = model_id, + ) -> None: + async with self._lock: + try: + await self._claim_turboquant_mid_prefill_process( + _model_id, + scheduler, + ) + except TurboQuantProcessExclusiveError as exc: + raise ModelLoadingError( + _model_id, + f"Model '{_model_id}' cannot switch to its " + "mid-prefill fallback while process-exclusive " + f"Metal access is unavailable: {exc}", + ) from exc + # Check if DFlash is enabled -- takes priority over engine type # since DFlash has its own model loading pipeline engine = None @@ -1736,10 +1852,16 @@ async def _load_engine( omlx_ssd_cache_dir=getattr( self._scheduler_config, "paged_ssd_cache_dir", None ), + prefill_eviction_callback=prefill_eviction_callback, + mid_prefill_process_claim_callback=( + mid_prefill_process_claim_callback + ), ) logger.info( f"DFlash enabled for {model_id}, draft={dflash_draft}" ) + except TurboQuantProcessExclusiveError: + raise except ImportError: logger.warning( f"DFlash enabled for {model_id} but dflash-mlx is not installed. " @@ -1761,16 +1883,6 @@ async def _load_engine( else False ) - async def prefill_eviction_callback( - eviction_request: object, - *, - _model_id: str = model_id, - ) -> bool: - return await self._evict_idle_lru_for_prefill( - exclude_model_id=_model_id, - eviction_request=eviction_request, - ) - # Create engine based on engine type (if DFlash not active) if engine is None: if effective_type == "embedding": @@ -1818,6 +1930,8 @@ async def prefill_eviction_callback( try: await engine.start() + except TurboQuantProcessExclusiveError: + raise except Exception as start_error: if _is_dflash_engine: # DFlash engine failed to start -- fall back to the @@ -1855,6 +1969,8 @@ async def prefill_eviction_callback( ) try: await engine.start() + except TurboQuantProcessExclusiveError: + raise except Exception as fallback_error: raise RuntimeError( f"DFlash load failed: {start_error}; " @@ -1893,6 +2009,8 @@ async def prefill_eviction_callback( ) try: await engine.start() + except TurboQuantProcessExclusiveError: + raise except Exception as fallback_error: raise RuntimeError( f"LM load failed (force_lm=True): {start_error}; " @@ -1929,6 +2047,8 @@ async def prefill_eviction_callback( ) try: await engine.start() + except TurboQuantProcessExclusiveError: + raise except Exception as fallback_error: raise RuntimeError( f"VLM load failed: {start_error}; " @@ -1938,7 +2058,7 @@ async def prefill_eviction_callback( entry.model_type = "llm" entry.engine_type = "batched" logger.info( - f"Successfully loaded {model_id} as LLM " f"(fallback from VLM)" + f"Successfully loaded {model_id} as LLM (fallback from VLM)" ) else: raise @@ -2069,6 +2189,49 @@ def _load_drafter_sync(path: str = drafter_path): "the request.", ) + scheduler = self._resolve_scheduler_from_engine(engine) + if ( + scheduler is not None + and getattr(scheduler, "_turboquant_mid_prefill", False) is True + ): + process_owner = getattr( + scheduler, + "_metal_process_owner", + None, + ) + try: + await self._claim_turboquant_mid_prefill_process( + model_id, + scheduler, + ) + except Exception: + entry.engine = None + self._current_model_memory = max( + 0, + self._current_model_memory - entry.estimated_size, + ) + load_completed = False + try: + await engine.stop() + except Exception: + logger.warning( + "Failed to stop non-exclusive mid-prefill engine %s", + model_id, + exc_info=True, + ) + finally: + close_process_owner = getattr(process_owner, "close", None) + if callable(close_process_owner): + try: + close_process_owner() + except Exception: + logger.warning( + "Failed to close non-exclusive mid-prefill " + "owner for %s", + model_id, + exc_info=True, + ) + raise logger.info( f"Loaded model: {model_id} " f"(actual: {format_size(entry.actual_size)}, " @@ -2086,6 +2249,12 @@ def _load_drafter_sync(path: str = drafter_path): # inflated and the memory-ceiling admission check rejects all # subsequent loads until a server restart. self._schedule_failed_load_reclaim(model_id, pre_load_memory) + if isinstance(exc, TurboQuantProcessExclusiveError): + raise ModelLoadingError( + model_id, + f"Model '{model_id}' cannot load while process-exclusive " + f"Metal access is unavailable: {exc}", + ) from exc if not entry.abort_loading and not entry_detached: self._mark_load_failure(entry, exc) logger.exception( diff --git a/omlx/exceptions.py b/omlx/exceptions.py index c19dce323..cc45dac4b 100644 --- a/omlx/exceptions.py +++ b/omlx/exceptions.py @@ -542,6 +542,10 @@ def __init__(self, required: int, current: int, message: str): super().__init__(message) +class TurboQuantProcessExclusiveError(RuntimeError): + """Raised when temporary process-wide Metal activity blocks exclusivity.""" + + class ModelLoadingError(EnginePoolError): """Raised when a model load is unavailable, blocked, or invalid.""" diff --git a/omlx/memory_monitor.py b/omlx/memory_monitor.py index 976aa466e..37cc5eb04 100644 --- a/omlx/memory_monitor.py +++ b/omlx/memory_monitor.py @@ -90,9 +90,7 @@ def register_attention_bias_transient(dtype_size: float | None) -> None: model load/swap: the setting is process-wide, like the tiled head_dim registry above.""" global _ATTENTION_BIAS_TRANSIENT_DTYPE_SIZE - _ATTENTION_BIAS_TRANSIENT_DTYPE_SIZE = ( - float(dtype_size) if dtype_size else None - ) + _ATTENTION_BIAS_TRANSIENT_DTYPE_SIZE = float(dtype_size) if dtype_size else None def estimate_unfused_sdpa_call_bytes( @@ -542,19 +540,34 @@ def estimate_block_memory( Estimate memory usage for a KV cache block. Args: - block_size: Number of tokens in the block - num_layers: Override stored num_layers - num_kv_heads: Override stored num_kv_heads - head_dim: Override stored head_dim - dtype_size: Override stored dtype_size + block_size: Number of tokens in the block. + num_layers: Override the number of full-attention KV cache layers. + num_kv_heads: Override stored num_kv_heads. + head_dim: Override stored head_dim. + dtype_size: Override stored dtype_size. Returns: Estimated memory in bytes for one block. """ - layers = num_layers or self._num_layers or 32 # Default for ~7B model - kv_heads = num_kv_heads or self._num_kv_heads or 8 - dim = head_dim or self._head_dim or 128 - dtype = dtype_size or self._dtype_size + if num_layers is not None: + layers = num_layers + elif self._num_kv_cache_layers is not None: + layers = self._num_kv_cache_layers + elif self._num_layers is not None: + layers = self._num_layers + else: + layers = 32 # Default for ~7B model + kv_heads = ( + num_kv_heads + if num_kv_heads is not None + else (self._num_kv_heads if self._num_kv_heads is not None else 8) + ) + dim = ( + head_dim + if head_dim is not None + else (self._head_dim if self._head_dim is not None else 128) + ) + dtype = self._dtype_size if dtype_size is None else dtype_size if ( self._kv_bytes_per_token_override is not None @@ -572,7 +585,32 @@ def estimate_block_memory( return total - def estimate_prompt_kv_bytes(self, num_tokens: int) -> float: + def estimate_paged_writer_block_memory(self, block_size: int) -> float: + """Estimate one queued SSD block, including boundary snapshot state.""" + if block_size <= 0: + return 0 + + total = self.estimate_block_memory(block_size) + if self._prefill_memory_profile is not None: + profile_bytes = self._prefill_memory_profile.estimate_resident_kv_bytes( + block_size, + chunk_tokens=block_size, + ) + total = max(total, profile_bytes) + + if self._rotating_layer_specs: + kv_heads = self._num_kv_heads or 0 + dim = self._head_dim or 0 + if kv_heads and dim: + per_token = kv_heads * dim * self._score_dtype_size * 2 + for count, window in self._rotating_layer_specs: + total += count * (window + block_size - 1) * per_token + + return total + self._fixed_state_bytes + + def estimate_prompt_kv_bytes( + self, num_tokens: int, *, dtype_size: float | None = None + ) -> float: """ Estimate KV cache memory for a prompt of given length. @@ -581,6 +619,8 @@ def estimate_prompt_kv_bytes(self, num_tokens: int) -> float: Args: num_tokens: Number of prompt tokens. + dtype_size: Optional stored-KV bytes per element for this request + phase. Defaults to the model-level cache width. Returns: Estimated KV cache memory in bytes. @@ -599,7 +639,7 @@ def estimate_prompt_kv_bytes(self, num_tokens: int) -> float: layers = self._num_layers or 0 kv_heads = self._num_kv_heads or 0 dim = self._head_dim or 0 - dtype = self._dtype_size + dtype = self._dtype_size if dtype_size is None else float(dtype_size) if not (layers and kv_heads and dim): return 0 @@ -612,7 +652,11 @@ def estimate_prompt_kv_bytes(self, num_tokens: int) -> float: return num_tokens * per_token def estimate_resident_kv_bytes( - self, num_tokens: int, *, chunk_tokens: int = 1 + self, + num_tokens: int, + *, + chunk_tokens: int = 1, + dtype_size: float | None = None, ) -> float: """Exact-shape resident KV bytes a prefill of ``num_tokens`` adds. @@ -639,7 +683,7 @@ def estimate_resident_kv_bytes( return self._prefill_memory_profile.estimate_resident_kv_bytes( num_tokens, chunk_tokens=chunk_tokens ) - total = self.estimate_prompt_kv_bytes(num_tokens) + total = self.estimate_prompt_kv_bytes(num_tokens, dtype_size=dtype_size) if self._rotating_layer_specs: kv_heads = self._num_kv_heads or 0 @@ -704,7 +748,9 @@ def _estimate_sdpa_activation_bytes(self, query_tokens: int, kv_len: int) -> int and query_tokens > 1 and kv_len >= _SDPA_TILED_MIN_KV_LEN ): - tile_scores = n_q * query_tokens * min(kv_tile, kv_len) * self._score_dtype_size + tile_scores = ( + n_q * query_tokens * min(kv_tile, kv_len) * self._score_dtype_size + ) return output + tile_scores + bias return ( @@ -715,7 +761,12 @@ def _estimate_sdpa_activation_bytes(self, query_tokens: int, kv_len: int) -> int ) def estimate_prefill_peak_bytes( - self, new_tokens: int, chunk_size: int, *, cached_tokens: int = 0 + self, + new_tokens: int, + chunk_size: int, + *, + cached_tokens: int = 0, + dtype_size: float | None = None, ) -> float: """ Estimate per-request prefill peak memory contribution (KV + SDPA). @@ -749,6 +800,8 @@ def estimate_prefill_peak_bytes( still typecheck — but they get the under-counting behavior this method was designed to fix, so always pass it when the value is available. + dtype_size: Optional stored-KV bytes per element for this request + phase. SDPA activation width remains the compute dtype. Returns: Per-request peak contribution in bytes (KV + SDPA). Returns 0 if @@ -778,7 +831,11 @@ def estimate_prefill_peak_bytes( # The cached portion is already counted in the caller's current-usage # baseline. Resident math includes window-capped sliding-window # layers and measured fixed state, not just full-attention KVCache. - kv = self.estimate_resident_kv_bytes(new_tokens, chunk_tokens=eff_chunk) + kv = self.estimate_resident_kv_bytes( + new_tokens, + chunk_tokens=eff_chunk, + dtype_size=dtype_size, + ) return attn + kv def estimate_chunk_transient_bytes(self, n_tokens: int, kv_len: int) -> int: @@ -803,6 +860,29 @@ def estimate_chunk_transient_bytes(self, n_tokens: int, kv_len: int) -> int: ) return self._estimate_sdpa_activation_bytes(n_tokens, kv_len) + def estimate_turboquant_prefill_attention_bytes( + self, + query_tokens: int, + kv_len: int, + *, + bits: float, + ) -> int: + """Return the source-structural long-prefill TurboQuant workspace.""" + from .turboquant_kv import ( + estimate_turboquant_prefill_attention_workspace_bytes, + ) + + return estimate_turboquant_prefill_attention_workspace_bytes( + query_tokens=query_tokens, + kv_len=kv_len, + num_query_heads=self._num_attention_heads or 0, + num_kv_heads=self._num_kv_heads or 0, + head_dim=self._head_dim or 0, + bits=bits, + compute_dtype_size=self._score_dtype_size, + causal=True, + ) + def estimate_blocks_to_free(self, bytes_to_free: int, block_size: int) -> int: """ Estimate number of blocks to evict to free the given bytes. @@ -1156,9 +1236,7 @@ def make_prefill_memory_profile( ) # Fixed-state recurrent caches (GDN/Mamba). Matches # Scheduler._cache_tree_has_arrays_cache plus mlx-lm's MambaCache. -_ARRAYS_CACHE_CLASS_NAMES = frozenset( - {"ArraysCache", "SizedArraysCache", "MambaCache"} -) +_ARRAYS_CACHE_CLASS_NAMES = frozenset({"ArraysCache", "SizedArraysCache", "MambaCache"}) def collect_kv_layer_specs( @@ -1487,7 +1565,10 @@ def raise_if_prefill_exceeds( request_id = f"preflight-{_uuid.uuid4().hex[:8]}" logger.warning( "Preflight rejected (%d tokens, cached=%d, request_id=%s): %s", - num_prompt_tokens, cached_tokens, request_id, message, + num_prompt_tokens, + cached_tokens, + request_id, + message, ) raise PrefillMemoryExceededError( message=message, diff --git a/omlx/model_profiles.py b/omlx/model_profiles.py index 4afffb853..aaaf2c18d 100644 --- a/omlx/model_profiles.py +++ b/omlx/model_profiles.py @@ -43,6 +43,7 @@ # Model-specific fields — eligible for per-model profiles only (never templates). MODEL_SPECIFIC_PROFILE_FIELDS = ( "turboquant_kv_enabled", + "turboquant_mid_prefill", "turboquant_kv_bits", "turboquant_skip_last", "dflash_enabled", diff --git a/omlx/model_settings.py b/omlx/model_settings.py index 87e935bed..36ee95044 100644 --- a/omlx/model_settings.py +++ b/omlx/model_settings.py @@ -107,6 +107,7 @@ class ModelSettings: guided_grammar_enabled: Whether a default guided grammar is active. guided_grammar: Default EBNF grammar for constrained decoding. turboquant_kv_enabled: Enable TurboQuant KV cache compression. + turboquant_mid_prefill: Convert the growing KV cache when prefill cannot fit. turboquant_kv_bits: TurboQuant bit depth (2/2.5/3/3.5/4/6/8). turboquant_skip_last: Skip last KVCache layer to prevent corruption. specprefill_enabled: Enable SpecPrefill (experimental sparse prefill for MoE). @@ -194,6 +195,7 @@ class ModelSettings: # TurboQuant KV cache (mlx-vlm backend) turboquant_kv_enabled: bool = False + turboquant_mid_prefill: bool = False turboquant_kv_bits: float = 4 # 2, 2.5, 3, 3.5, 4, 6, 8 turboquant_skip_last: bool = ( True # Skip last KVCache layer (prevents corruption on sensitive models) diff --git a/omlx/patches/turboquant_attention.py b/omlx/patches/turboquant_attention.py index 40a92df91..280cff45e 100644 --- a/omlx/patches/turboquant_attention.py +++ b/omlx/patches/turboquant_attention.py @@ -18,13 +18,18 @@ from typing import Optional import mlx.core as mx +from ..turboquant_kv import ( + TURBOQUANT_PREFILL_KEY_CHUNK_TOKENS, + TURBOQUANT_PREFILL_QUERY_BLOCK_TOKENS, +) + logger = logging.getLogger(__name__) _PATCHED = False _LONG_PREFILL_QUANTIZED_THRESHOLD = 8192 -_LONG_PREFILL_QUERY_BLOCK_SIZE = 256 -_LONG_PREFILL_KEY_CHUNK_SIZE = 16384 +_LONG_PREFILL_QUERY_BLOCK_SIZE = TURBOQUANT_PREFILL_QUERY_BLOCK_TOKENS +_LONG_PREFILL_KEY_CHUNK_SIZE = TURBOQUANT_PREFILL_KEY_CHUNK_TOKENS # MTP verify is a decode-shaped multi-row call (q_len = 1 + draft depth <= 9). # Above this floor a multi-row call is genuine (chunked) prefill. _DECODE_MULTIROW_MAX_Q_LEN = 15 @@ -621,10 +626,11 @@ def patched_sdpa( total_tokens = _state_length(keys_state) except Exception: total_tokens = 0 - if ( - total_tokens > _LONG_PREFILL_QUANTIZED_THRESHOLD - and hasattr(real_cache, "quantized_attention") - ): + if total_tokens > _LONG_PREFILL_QUANTIZED_THRESHOLD: + if not hasattr(real_cache, "quantized_attention"): + raise RuntimeError( + "Long TurboQuant prefill requires quantized attention" + ) old_query_block_size = getattr( real_cache, "prefill_query_block_size", None ) @@ -643,12 +649,6 @@ def patched_sdpa( scale=scale, mask=mask, ) - except Exception: - logger.debug( - "TurboQuant quantized prefill attention failed; " - "falling back to dequantize+SDPA", - exc_info=True, - ) finally: if old_query_block_size is not None: real_cache.prefill_query_block_size = old_query_block_size diff --git a/omlx/request.py b/omlx/request.py index a30b85043..5ee309cdc 100644 --- a/omlx/request.py +++ b/omlx/request.py @@ -210,6 +210,12 @@ def vlm_extra_key_ranges_for_cache( prefill_eviction_retries: int = ( 0 # Per-request prefill-headroom eviction phase counter ) + prefill_started_at: float | None = ( + None # First external/chunked prefill attempt; survives LRU retries + ) + turboquant_mid_prefill_attempted: bool = ( + False # Durable across prefill-OOM requeues; one attempt per request + ) # Request-scoped tool schemas used by protocol output parsers. tools: list[dict[str, Any]] | None = None diff --git a/omlx/scheduler.py b/omlx/scheduler.py index 6dd211b1c..2d16ca14f 100644 --- a/omlx/scheduler.py +++ b/omlx/scheduler.py @@ -16,9 +16,11 @@ import gc import importlib import logging +import math import os import threading import time +import traceback from array import array from collections import OrderedDict, defaultdict, deque from collections.abc import Callable @@ -68,10 +70,18 @@ run_vlm_mtp_decode, vlm_mtp_positioned_sampling_available, ) +from .turboquant_kv import ( + TURBOQUANT_CONVERSION_SLICE_TOKENS, + TurboQuantConversionStats, + convert_kv_cache_sliced, + estimate_turboquant_conversion_peak_bytes, + turboquant_mse_bytes_per_element, +) from .utils.fatal import FATAL_TEARDOWN_TIMEOUT_S, fatal_exit from .utils.generation_config import load_generation_config_token_ids from .utils.hardware import format_bytes from .utils.metal_sync import ( + _conversion_coordinator, _default_generation_stream, _mx_buffer_access_lock, _sync_and_clear_cache, @@ -405,10 +415,74 @@ def __init__(self, aborted_uids: list[int], processed_tokens: int): self.aborted_uids = aborted_uids self.processed_tokens = processed_tokens super().__init__( - f"Prefill aborted for UIDs {aborted_uids} " f"at {processed_tokens} tokens" + f"Prefill aborted for UIDs {aborted_uids} at {processed_tokens} tokens" ) +def _format_and_detach_exception(exc: BaseException) -> str: + """Format an exception, then remove every traceback/frame reference.""" + formatted = "".join( + traceback.format_exception(type(exc), exc, exc.__traceback__) + ).rstrip() + current: BaseException | None = exc + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + next_exception = current.__cause__ or current.__context__ + current_traceback = current.__traceback__ + current.__traceback__ = None + current.__cause__ = None + current.__context__ = None + if current_traceback is not None: + traceback.clear_frames(current_traceback) + current = next_exception + return formatted + + +class _PrefillKVPhase(Enum): + """Resident full-attention KV representation during prefill.""" + + DENSE = "dense" + TURBOQUANT = "turboquant" + UNSUPPORTED = "unsupported" + INVALID_PARTIAL = "invalid_partial" + + +@dataclass +class _PrefillContext: + """Request-scoped phase and telemetry for one prefill.""" + + request_id: str + loop_label: str + phase: _PrefillKVPhase + conversion_eligible: bool + started_at: float | None = None + conversion_attempted: bool = False + mid_triggered: bool = False + trigger_tokens: int | None = None + conversion_seconds: float = 0.0 + converted_layers: int = 0 + conversion_slices: int = 0 + skipped_dense_layers: int = 0 + memory_before_bytes: int = 0 + memory_after_bytes: int = 0 + post_trigger_tokens: int = 0 + conversion_completed_at: float | None = None + summary_logged: bool = False + + +@dataclass(frozen=True, slots=True) +class _GuardedTurboQuantConversion: + """Successful conversion samples and telemetry.""" + + stats: TurboQuantConversionStats + before_bytes: int + after_bytes: int + estimated_peak_bytes: int + conversion_seconds: float + completed_at: float + + @dataclass class PrefillEvictionRequest: """Internal request for async LRU model eviction before prefill.""" @@ -420,6 +494,7 @@ class PrefillEvictionRequest: predicted_transient_bytes: int requested_tokens: int reason: str + processed_tokens: int = 0 class _PrefillEvictionNeeded(Exception): @@ -483,6 +558,7 @@ class _PrefillState: boundary_enabled: bool # Whether boundary snapshots are active block_size: int # Copied from config.paged_cache_block_size total_length: int # len(original tokens) for completeness + prefill_context: _PrefillContext | None = None # Pre-built insert-time params (set by _schedule_waiting before enqueuing) sampler: Any = None sm: Any = None @@ -1096,7 +1172,7 @@ def _ckvcache_extend_passthrough(self, other): # Surface which ones so a regression in Llama-4 batching is visible # to operators without diffing the patch against installed mlx_lm. logger.info( - "ChunkedKVCache patch: methods already present upstream, " "skipped: %s", + "ChunkedKVCache patch: methods already present upstream, skipped: %s", ", ".join(_ckvcache_methods_skipped), ) except ImportError: @@ -1595,6 +1671,12 @@ def __init__( # TurboQuant KV cache (set by engine if model_settings has it enabled) self._turboquant_kv_bits: float | None = None self._turboquant_skip_last: bool = True + self._turboquant_mid_prefill: bool = False + self._prefill_dense_kv_dtype_size: float | None = None + self._prefill_tq_kv_dtype_size: float | None = None + # Route-level preflight can only defer a dense-memory rejection when + # model cache construction has confirmed a convertible dense layout. + self._turboquant_preflight_conversion_eligible: bool | None = None # Memoized MLA-architecture detection (see _model_uses_mla / #1613). self._mla_model: bool | None = None self._glm_dsa_adaptive_prefill = None @@ -1742,6 +1824,10 @@ def __init__( # default level even while it ran an order of magnitude slower; one # line per request keeps that visible without a line per chunk. self._throttle_notified_requests: set[str] = set() + # Engine wrappers set this after construction. ``None`` keeps direct + # Scheduler users on the legacy one-shot path; ``False`` avoids a + # pointless pause when no async eviction callback exists. + self._prefill_eviction_callback_configured: bool | None = None self._pending_prefill_eviction_request: PrefillEvictionRequest | None = None self._memory_admission_blocked_request_id: str | None = None self._memory_admission_blocked_since: float = 0.0 @@ -1755,6 +1841,9 @@ def __init__( self._prefill_transient_tracker = PrefillTransientTracker( model_id=_tracker_model_id ) + self._prefill_tq_transient_tracker = PrefillTransientTracker( + model_id=_tracker_model_id + ) # One-shot probe of the GDN/Mamba fixed recurrent-state footprint, # armed by _set_model_info_for_monitor when ArraysCache layers exist # and taken after the first prefill chunk's eval. @@ -1947,9 +2036,9 @@ def __init__( # Streaming detokenizers for proper UTF-8 handling (one per active request) # NOTE: No pooling - each request gets a fresh instance to prevent state contamination - self._request_detokenizers: dict[str, Any] = ( - {} - ) # request_id → active detokenizer + self._request_detokenizers: dict[ + str, Any + ] = {} # request_id → active detokenizer # Protocol-specific output parser support (e.g. Harmony, Gemma 4) self._output_parser_factory: OutputParserFactory | None = None @@ -2947,14 +3036,10 @@ def _ok(c: Any) -> bool: return False if isinstance(c, CacheList): # A KVCache member inside a CacheList converts fine at - # runtime, but the prefix/SSD store paths dispatch on the - # layer class ("CacheList") and have no TurboQuant - # sub-state serialization: the converted member's - # NamedTuple state is flattened to an anonymous tuple on - # store and rebuilt as a corrupt dense cache on restore. - # Until CacheList-level TQ serialization exists, exclude - # composite layers that contain a convertible KVCache - # (e.g. inkling's CacheList(KVCache, ArraysCache)). + # runtime, but prefix/SSD storage dispatches on the outer + # CacheList class and cannot serialize TurboQuant sub-state. + # Flattening that state would restore a corrupt dense cache, + # so reject this layout until CacheList has a TQ-aware format. if any(type(inner) is KVCache for inner in c.caches): if not getattr(self, "_tq_cachelist_guard_logged", False): self._tq_cachelist_guard_logged = True @@ -2969,6 +3054,240 @@ def _ok(c: Any) -> bool: return bool(prompt_cache) and all(_ok(c) for c in prompt_cache) + def _confirm_turboquant_preflight_conversion_eligibility( + self, + prompt_cache: list[Any], + ) -> bool: + """Return True only when prompt_cache is an all-dense convertible layout.""" + if not isinstance(prompt_cache, list) or not prompt_cache: + return False + if not self._turboquant_eligible(prompt_cache): + return False + family_targets = [ + cache_obj + for cache_obj in prompt_cache + if _is_turboquant_kv_family_cache(cache_obj) + ] + if not family_targets: + return False + return all(isinstance(target, _MLXKVCache) for target in family_targets) + + def _classify_prefill_cache( + self, prompt_cache: list[Any] + ) -> tuple[_PrefillKVPhase, bool]: + """Classify only complete dense or complete configured-TQ cache states.""" + raw_bits = getattr(self, "_turboquant_kv_bits", None) + if not isinstance(raw_bits, (int, float)) or isinstance(raw_bits, bool): + return _PrefillKVPhase.UNSUPPORTED, False + if not prompt_cache or not self._turboquant_eligible(prompt_cache): + return _PrefillKVPhase.UNSUPPORTED, False + + family_indices = [ + index + for index, cache_obj in enumerate(prompt_cache) + if _is_turboquant_kv_family_cache(cache_obj) + ] + if not family_indices: + return _PrefillKVPhase.UNSUPPORTED, False + + skip_last = ( + bool(getattr(self, "_turboquant_skip_last", True)) + and len(family_indices) > 1 + ) + skipped_index = family_indices[-1] if skip_last else None + expected_tokens = _cache_layer_token_count(prompt_cache[family_indices[0]]) + dense_targets = 0 + turboquant_targets = 0 + target_count = len(family_indices) - (1 if skipped_index is not None else 0) + expected_bits = float(raw_bits) + + for index in family_indices: + cache_obj = prompt_cache[index] + if _cache_layer_token_count(cache_obj) != expected_tokens: + return _PrefillKVPhase.INVALID_PARTIAL, False + is_turboquant = _is_turboquant_kv_cache(cache_obj) + if index == skipped_index: + if is_turboquant or not isinstance(cache_obj, _MLXKVCache): + return _PrefillKVPhase.INVALID_PARTIAL, False + continue + if is_turboquant: + cache_bits = getattr(cache_obj, "bits", None) + if not isinstance(cache_bits, (int, float)) or isinstance( + cache_bits, bool + ): + return _PrefillKVPhase.INVALID_PARTIAL, False + if not math.isclose( + float(cache_bits), expected_bits, rel_tol=0.0, abs_tol=1e-6 + ): + return _PrefillKVPhase.INVALID_PARTIAL, False + turboquant_targets += 1 + elif isinstance(cache_obj, _MLXKVCache): + dense_targets += 1 + else: + return _PrefillKVPhase.INVALID_PARTIAL, False + + if dense_targets == target_count and turboquant_targets == 0: + return _PrefillKVPhase.DENSE, True + if turboquant_targets == target_count and dense_targets == 0: + return _PrefillKVPhase.TURBOQUANT, False + return _PrefillKVPhase.INVALID_PARTIAL, False + + def _has_populated_turboquant_target( + self, + prompt_cache: list[Any], + ) -> bool: + """Return whether a mid-prefill conversion has dense history to keep.""" + family_indices = [ + index + for index, cache_obj in enumerate(prompt_cache) + if _is_turboquant_kv_family_cache(cache_obj) + ] + skipped_index = ( + family_indices[-1] + if bool(getattr(self, "_turboquant_skip_last", True)) + and len(family_indices) > 1 + else None + ) + return any( + index != skipped_index + and isinstance(prompt_cache[index], _MLXKVCache) + and _cache_layer_token_count(prompt_cache[index]) > 0 + for index in family_indices + ) + + def _discard_failed_prefill_cache( + self, request: "Request", prompt_cache: list[Any] + ) -> None: + """Drop every reference to a cache that may be partly converted.""" + if getattr(request, "prompt_cache", None) is prompt_cache: + request.prompt_cache = None + prompt_cache.clear() + _sync_and_clear_cache(self._stream) + + def _reclaim_failed_turboquant_conversion( + self, + request: "Request", + prompt_cache: list[Any], + ) -> None: + """Drop failed conversion state, collect it, then clear Metal buffers.""" + if getattr(request, "prompt_cache", None) is prompt_cache: + request.prompt_cache = None + prompt_cache.clear() + try: + gc.collect() + except Exception as exc: + reclaim_traceback = _format_and_detach_exception(exc) + del exc + logger.error( + "TurboQuant failure GC failed for %s:\n%s", + request.request_id, + reclaim_traceback, + ) + try: + _sync_and_clear_cache(self._stream) + except Exception as exc: + reclaim_traceback = _format_and_detach_exception(exc) + del exc + logger.error( + "TurboQuant failure Metal clear failed for %s:\n%s", + request.request_id, + reclaim_traceback, + ) + + def _new_prefill_context( + self, + request: "Request", + prompt_cache: list[Any], + *, + loop_label: str, + ) -> _PrefillContext: + """Create request-local phase state and reject partial caches.""" + mid_prefill_enabled = bool(getattr(self, "_turboquant_mid_prefill", False)) + if not mid_prefill_enabled and not any( + _is_turboquant_kv_cache(cache_obj) for cache_obj in prompt_cache + ): + return _PrefillContext( + request_id=request.request_id, + loop_label=loop_label, + phase=_PrefillKVPhase.DENSE, + conversion_eligible=False, + started_at=None, + conversion_attempted=False, + ) + phase, conversion_eligible = self._classify_prefill_cache(prompt_cache) + if not mid_prefill_enabled: + conversion_eligible = False + if phase is _PrefillKVPhase.UNSUPPORTED: + phase = _PrefillKVPhase.DENSE + if phase is _PrefillKVPhase.INVALID_PARTIAL: + current = self._current_usage_bytes() + limit = self._prefill_abort_cap() or None + self._discard_failed_prefill_cache(request, prompt_cache) + raise PrefillMemoryExceededError( + message=( + "TurboQuant prefill cache was partially converted or used " + "an incompatible bit width; the request cache was discarded" + ), + request_id=request.request_id, + estimated_bytes=current, + limit_bytes=limit, + ) + if not mid_prefill_enabled: + return _PrefillContext( + request_id=request.request_id, + loop_label=loop_label, + phase=phase, + conversion_eligible=False, + started_at=None, + conversion_attempted=False, + ) + started_at = request.prefill_started_at + telemetry_eligible = bool( + getattr(self, "_turboquant_mid_prefill", False) + and getattr(self, "_turboquant_kv_bits", None) is not None + and phase is _PrefillKVPhase.DENSE + and conversion_eligible + ) + if telemetry_eligible and started_at is None: + started_at = time.perf_counter() + request.prefill_started_at = started_at + return _PrefillContext( + request_id=request.request_id, + loop_label=loop_label, + phase=phase, + conversion_eligible=conversion_eligible, + started_at=started_at, + conversion_attempted=request.turboquant_mid_prefill_attempted, + ) + + def _prefill_phase_dtype_size(self, phase: _PrefillKVPhase) -> float | None: + """Return the full-attention KV width for one complete cache phase.""" + if phase is _PrefillKVPhase.TURBOQUANT: + width = getattr(self, "_prefill_tq_kv_dtype_size", None) + else: + width = getattr(self, "_prefill_dense_kv_dtype_size", None) + if isinstance(width, (int, float)) and not isinstance(width, bool): + return float(width) + return None + + def _prefill_transient_tracker_for_phase( + self, phase: _PrefillKVPhase + ) -> PrefillTransientTracker | None: + """Return the history trained under the same resident KV width.""" + if phase is _PrefillKVPhase.TURBOQUANT: + return getattr(self, "_prefill_tq_transient_tracker", None) + return getattr(self, "_prefill_transient_tracker", None) + + def _raise_if_prefill_cancelled( + self, request_id: str, processed_tokens: int + ) -> None: + """Raise the existing prefill-abort signal at a conversion boundary.""" + if request_id not in self._pending_abort_ids: + return + uid = self.request_id_to_uid.get(request_id) + aborted_uids = [uid] if uid is not None else [] + raise _PrefillAbortedError(aborted_uids, processed_tokens) + def _apply_turboquant_kv_empty(self, prompt_cache: list[Any]) -> None: """Replace empty KVCache layers with empty TurboQuantKVCache. @@ -3012,47 +3331,429 @@ def _apply_turboquant_kv_empty(self, prompt_cache: list[Any]) -> None: ) def _apply_turboquant_kv_convert(self, prompt_cache: list[Any]) -> None: - """Convert populated KVCache data to TurboQuantKVCache via from_cache(). - - Called AFTER fp16 prefill completes (or on an SSD-restored fp16 - cache): the completed full-precision KV is quantized once, so prefill - hidden states stay exact and quantization error only enters at - decode-time reads. This is the key difference from #717/#771, which - quantized on the fly during prefill and corrupted hidden states. - """ + """Convert populated KVCache data with the ordinary final-only path.""" from mlx_lm.models.cache import CacheList, KVCache from mlx_vlm.turboquant import TurboQuantKVCache kv_indices = [ - i for i, c in enumerate(prompt_cache) if _is_turboquant_kv_family_cache(c) + index + for index, cache_obj in enumerate(prompt_cache) + if _is_turboquant_kv_family_cache(cache_obj) ] skip_last = self._turboquant_skip_last and len(kv_indices) > 1 - last_kv_idx = kv_indices[-1] if skip_last else -1 + last_kv_index = kv_indices[-1] if skip_last else -1 converted = 0 bits = float(self._turboquant_kv_bits) - for i, cache_obj in enumerate(prompt_cache): + for index, cache_obj in enumerate(prompt_cache): if isinstance(cache_obj, KVCache): - if i == last_kv_idx: + if index == last_kv_index: continue - prompt_cache[i] = TurboQuantKVCache.from_cache(cache_obj, bits=bits) + prompt_cache[index] = TurboQuantKVCache.from_cache( + cache_obj, + bits=bits, + ) converted += 1 elif isinstance(cache_obj, CacheList): new_caches = [] - for c in cache_obj.caches: - if isinstance(c, KVCache): - new_caches.append(TurboQuantKVCache.from_cache(c, bits=bits)) + for inner_cache in cache_obj.caches: + if isinstance(inner_cache, KVCache): + new_caches.append( + TurboQuantKVCache.from_cache(inner_cache, bits=bits) + ) converted += 1 else: - new_caches.append(c) + new_caches.append(inner_cache) cache_obj.caches = tuple(new_caches) if converted > 0: - skip_msg = ", skipped last KVCache layer" if skip_last else "" + skip_message = ", skipped last KVCache layer" if skip_last else "" logger.info( - f"TurboQuant: converted {converted}/{len(prompt_cache)} " - f"cache layers to {bits}-bit{skip_msg}" + "TurboQuant: converted %d/%d cache layers to %s-bit%s", + converted, + len(prompt_cache), + bits, + skip_message, ) + def _apply_turboquant_kv_convert_sliced( + self, + prompt_cache: list[Any], + *, + check_cancelled: Callable[[], None] | None = None, + log_result: bool = True, + ) -> TurboQuantConversionStats: + """Convert populated dense KV in bounded slices, one layer at a time.""" + raw_bits = getattr(self, "_turboquant_kv_bits", None) + if not isinstance(raw_bits, (int, float)) or isinstance(raw_bits, bool): + raise ValueError("TurboQuant KV bit width is not configured") + stats = convert_kv_cache_sliced( + prompt_cache, + bits=float(raw_bits), + skip_last=bool(getattr(self, "_turboquant_skip_last", True)), + slice_tokens=TURBOQUANT_CONVERSION_SLICE_TOKENS, + stream=getattr(self, "_stream", None), + check_cancelled=check_cancelled, + ) + if log_result and stats.converted_layers > 0: + skip_message = ( + ", skipped last KVCache layer" if stats.skipped_dense_layers > 0 else "" + ) + logger.info( + "TurboQuant: converted %d/%d cache layers to %s-bit%s", + stats.converted_layers, + len(prompt_cache), + float(raw_bits), + skip_message, + ) + return stats + + def _run_guarded_turboquant_conversion( + self, + *, + request: "Request", + prompt_cache: list[Any], + processed_tokens: int, + safety_cap: int, + conversion_label: str, + log_result: bool, + ) -> _GuardedTurboQuantConversion: + """Reserve and run one conversion under the process-wide exclusive gate.""" + result: _GuardedTurboQuantConversion | None = None + with ( + _conversion_coordinator.conversion( + process_owner=getattr(self, "_metal_process_owner", None) + ) as conversion_owner, + _mx_buffer_access_lock, + ): + before = self._current_usage_bytes() + conversion_peak = 0 + estimate_traceback: str | None = None + try: + raw_bits = getattr(self, "_turboquant_kv_bits", None) + if not isinstance(raw_bits, (int, float)) or isinstance(raw_bits, bool): + raise ValueError("TurboQuant KV bit width is not configured") + conversion_peak = estimate_turboquant_conversion_peak_bytes( + prompt_cache, + bits=float(raw_bits), + skip_last=bool(getattr(self, "_turboquant_skip_last", True)), + slice_tokens=TURBOQUANT_CONVERSION_SLICE_TOKENS, + ) + except Exception as exc: + estimate_traceback = _format_and_detach_exception(exc) + del exc + + if estimate_traceback is not None: + logger.error( + "TurboQuant %s conversion estimate failed for %s:\n%s", + conversion_label, + request.request_id, + estimate_traceback, + ) + self._reclaim_failed_turboquant_conversion( + request, + prompt_cache, + ) + raise PrefillMemoryExceededError( + message=( + f"TurboQuant {conversion_label} conversion could not " + "be bounded; the request cache was discarded" + ), + request_id=request.request_id, + estimated_bytes=before, + limit_bytes=safety_cap or None, + ) from None + + accepted, estimated_peak = _conversion_coordinator.try_reserve( + conversion_owner, + current_bytes=before, + peak_bytes=conversion_peak, + limit_bytes=safety_cap, + ) + if not accepted: + self._reclaim_failed_turboquant_conversion( + request, + prompt_cache, + ) + raise PrefillMemoryExceededError( + message=( + f"TurboQuant {conversion_label} conversion would " + f"exceed the prefill safety cap at " + f"{processed_tokens} tokens" + ), + request_id=request.request_id, + estimated_bytes=int(estimated_peak), + limit_bytes=int(safety_cap), + ) from None + + if conversion_label == "mid-prefill": + logger.info( + "TurboQuant mid-prefill trigger for %s at %d tokens " + "(usage=%.3fGiB, conversion_peak=%.3fGiB, " + "safety_cap=%.3fGiB)", + request.request_id, + processed_tokens, + before / 1024**3, + conversion_peak / 1024**3, + safety_cap / 1024**3, + ) + + def _check_cancelled() -> None: + self._raise_if_prefill_cancelled( + request.request_id, + processed_tokens, + ) + + started = time.perf_counter() + stats: TurboQuantConversionStats | None = None + completed_at: float | None = None + after = before + failure_traceback: str | None = None + failure_message: str | None = None + aborted_uids: list[int] | None = None + aborted_tokens = processed_tokens + try: + stats = self._apply_turboquant_kv_convert_sliced( + prompt_cache, + check_cancelled=_check_cancelled, + log_result=log_result, + ) + converted_phase, _ = self._classify_prefill_cache(prompt_cache) + if converted_phase is not _PrefillKVPhase.TURBOQUANT: + raise RuntimeError( + f"conversion ended in incomplete cache phase " + f"{converted_phase.value}" + ) + gc.collect() + _sync_and_clear_cache(self._stream) + completed_at = time.perf_counter() + _conversion_coordinator.release_reservation(conversion_owner) + after = self._current_usage_bytes() + if safety_cap > 0 and after > safety_cap: + failure_message = ( + f"TurboQuant {conversion_label} converted cache " + "remained above the prefill safety cap" + ) + except _PrefillAbortedError as exc: + aborted_uids = list(exc.aborted_uids) + aborted_tokens = int(exc.processed_tokens) + _format_and_detach_exception(exc) + del exc + except Exception as exc: + failure_traceback = _format_and_detach_exception(exc) + failure_message = f"TurboQuant {conversion_label} conversion failed" + del exc + + if aborted_uids is not None: + stats = None + completed_at = None + self._reclaim_failed_turboquant_conversion( + request, + prompt_cache, + ) + raise _PrefillAbortedError( + aborted_uids, + aborted_tokens, + ) from None + + if failure_message is not None: + if failure_traceback is not None: + logger.error( + "TurboQuant %s conversion failed for %s:\n%s", + conversion_label, + request.request_id, + failure_traceback, + ) + stats = None + completed_at = None + self._reclaim_failed_turboquant_conversion( + request, + prompt_cache, + ) + raise PrefillMemoryExceededError( + message=(f"{failure_message}; the request cache was discarded"), + request_id=request.request_id, + estimated_bytes=int(estimated_peak), + limit_bytes=safety_cap or None, + ) from None + + if stats is None or completed_at is None: + raise RuntimeError("TurboQuant conversion produced no result") + result = _GuardedTurboQuantConversion( + stats=stats, + before_bytes=before, + after_bytes=after, + estimated_peak_bytes=int(estimated_peak), + conversion_seconds=completed_at - started, + completed_at=completed_at, + ) + + if result is None: + raise RuntimeError("TurboQuant conversion gate produced no result") + return result + + def _attempt_mid_prefill_conversion( + self, + *, + request: "Request", + prompt_cache: list[Any], + context: _PrefillContext, + processed_tokens: int, + safety_cap: int, + ) -> None: + """Run the request's sole pressure conversion or fail the request.""" + request.turboquant_mid_prefill_attempted = True + context.conversion_attempted = True + context.mid_triggered = True + context.trigger_tokens = processed_tokens + result = self._run_guarded_turboquant_conversion( + request=request, + prompt_cache=prompt_cache, + processed_tokens=processed_tokens, + safety_cap=safety_cap, + conversion_label="mid-prefill", + log_result=False, + ) + context.phase = _PrefillKVPhase.TURBOQUANT + context.conversion_eligible = False + context.converted_layers = result.stats.converted_layers + context.conversion_slices = result.stats.slices + context.skipped_dense_layers = result.stats.skipped_dense_layers + context.memory_before_bytes = result.before_bytes + context.memory_after_bytes = result.after_bytes + context.conversion_seconds = result.conversion_seconds + context.conversion_completed_at = result.completed_at + + def _mid_prefill_conversion_available( + self, + context: _PrefillContext | None, + prompt_cache: list[Any] | None, + request: "Request | None", + ) -> bool: + """Return whether this dense request still owns its one trigger.""" + if context is None: + return False + return bool( + getattr(self, "_turboquant_mid_prefill", False) + and getattr(self, "_turboquant_kv_bits", None) is not None + and context.phase is _PrefillKVPhase.DENSE + and context.conversion_eligible + and not context.conversion_attempted + and _conversion_coordinator.process_exclusive( + getattr(self, "_metal_process_owner", None) + ) + and request is not None + and not request.turboquant_mid_prefill_attempted + and ( + prompt_cache is None + or self._has_populated_turboquant_target(prompt_cache) + ) + ) + + def _can_defer_mid_prefill_preflight( + self, + phase: _PrefillKVPhase, + *, + request: "Request | None" = None, + conversion_eligible: bool | None = None, + ) -> bool: + """Let an unspent eligible conversion decide dense peak rejection.""" + return bool( + getattr(self, "_turboquant_mid_prefill", False) + and getattr(self, "_turboquant_kv_bits", None) is not None + and getattr(self, "_prefill_tq_kv_dtype_size", None) is not None + and _conversion_coordinator.process_exclusive( + getattr(self, "_metal_process_owner", None) + ) + and phase is _PrefillKVPhase.DENSE + and conversion_eligible is True + and (request is None or not request.turboquant_mid_prefill_attempted) + ) + + def _finalize_turboquant_prefill_cache( + self, + request: "Request", + prompt_cache: list[Any], + *, + processed_tokens: int, + context: _PrefillContext | None = None, + ) -> TurboQuantConversionStats | None: + """Complete ordinary post-prefill TurboQuant conversion.""" + if getattr(self, "_turboquant_kv_bits", None) is None: + return None + if not getattr(self, "_turboquant_mid_prefill", False): + if self._turboquant_eligible(prompt_cache): + self._apply_turboquant_kv_convert(prompt_cache) + return None + phase, _ = self._classify_prefill_cache(prompt_cache) + if phase is _PrefillKVPhase.UNSUPPORTED: + return None + if phase is _PrefillKVPhase.TURBOQUANT: + if context is not None: + context.phase = phase + return None + if phase is _PrefillKVPhase.INVALID_PARTIAL: + self._discard_failed_prefill_cache(request, prompt_cache) + raise PrefillMemoryExceededError( + message=( + "TurboQuant prefill cache was partially converted; the " + "request cache was discarded" + ), + request_id=request.request_id, + estimated_bytes=self._current_usage_bytes(), + limit_bytes=self._prefill_abort_cap() or None, + ) + + result = self._run_guarded_turboquant_conversion( + request=request, + prompt_cache=prompt_cache, + processed_tokens=processed_tokens, + safety_cap=self._prefill_abort_cap(), + conversion_label="post-prefill", + log_result=True, + ) + if context is not None: + context.phase = _PrefillKVPhase.TURBOQUANT + context.conversion_eligible = False + return result.stats + + def _log_mid_prefill_summary(self, context: _PrefillContext) -> None: + """Emit one completion record for an actual pressure conversion.""" + if not context.mid_triggered or context.summary_logged: + return + context.summary_logged = True + completed_at = time.perf_counter() + total_wall = ( + completed_at - context.started_at if context.started_at is not None else 0.0 + ) + post_trigger_wall = ( + completed_at - context.conversion_completed_at + if context.conversion_completed_at is not None + else 0.0 + ) + throughput = ( + context.post_trigger_tokens / post_trigger_wall + if post_trigger_wall > 0 + else 0.0 + ) + logger.info( + "TurboQuant mid-prefill complete for %s " + "(trigger_tokens=%d, conversion_pause=%.3fs, " + "post_trigger_tokens=%d, post_trigger_prefill_tps=%.2f tok/s, " + "total_wall=%.3fs, memory_before=%.3fGiB, memory_after=%.3fGiB, " + "layers=%d, slices=%d, skipped_dense=%d)", + context.request_id, + context.trigger_tokens or 0, + context.conversion_seconds, + context.post_trigger_tokens, + throughput, + total_wall, + context.memory_before_bytes / 1024**3, + context.memory_after_bytes / 1024**3, + context.converted_layers, + context.conversion_slices, + context.skipped_dense_layers, + ) + def _do_external_prefill( self, request: "Request", @@ -3095,7 +3796,11 @@ def _do_external_prefill( if existing_cache is None: self._apply_turboquant_kv_empty(cache) else: - self._apply_turboquant_kv_convert(cache) + self._finalize_turboquant_prefill_cache( + request, + cache, + processed_tokens=0, + ) return cache, tokens # Create or reuse cache @@ -3103,6 +3808,11 @@ def _do_external_prefill( prompt_cache = existing_cache else: prompt_cache = make_prompt_cache(self.model) + prefill_context = self._new_prefill_context( + request, + prompt_cache, + loop_label="external", + ) # Fresh TurboQuant requests run fp16 during the cold prefill loop and # are quantized once at the end. Restored TurboQuant prefix caches stay @@ -3214,6 +3924,7 @@ def _do_external_prefill( request_id=request.request_id, loop_label="external", kv_len=base_size + processed_tokens, + prefill_context=prefill_context, ) # Pre-chunk safety guard: NEVER submit a chunk whose predicted peak @@ -3228,6 +3939,9 @@ def _do_external_prefill( progress=processed_tokens, loop_label="external", request_id=request.request_id, + request=request, + prompt_cache=prompt_cache, + prefill_context=prefill_context, ) _throttle_pre = get_phys_footprint() @@ -3249,6 +3963,10 @@ def _do_external_prefill( model_kwargs["vlm_extra_kwargs"] = _slice_vlm_extra( extra_kwargs, n_to_process ) + measure_post_trigger = ( + prefill_context.mid_triggered + and prefill_context.phase is _PrefillKVPhase.TURBOQUANT + ) self.model( input_arr[:, :n_to_process], cache=prompt_cache, @@ -3269,7 +3987,10 @@ def _do_external_prefill( loop_label="external", kv_len=base_size + processed_tokens, requested_step=prefill_step_size, + phase=prefill_context.phase, ) + if measure_post_trigger: + prefill_context.post_trigger_tokens += n_to_process self._maybe_record_fixed_state_bytes(prompt_cache) # Enforcer-requested hard-pressure drain. The flag's normal # consumption point is the end-of-step cleanup, but this loop @@ -3353,9 +4074,7 @@ def _do_external_prefill( # band by design — the per-chunk notice is DEBUG there, # not a warning about an unexpected state. _log = ( - logger.debug - if self._prefill_speed_priority - else logger.warning + logger.debug if self._prefill_speed_priority else logger.warning ) _log( f"Prefill above max_bytes at " @@ -3419,14 +4138,17 @@ def _do_external_prefill( # format is unchanged. _merge_caches() then builds a # BatchTurboQuantKVCache when this request is inserted. Gated to dense # KVCache models — chunked/rotating caches stay fp16. - if self._turboquant_kv_bits is not None and self._turboquant_eligible( - prompt_cache - ): - self._apply_turboquant_kv_convert(prompt_cache) + self._finalize_turboquant_prefill_cache( + request, + prompt_cache, + processed_tokens=processed_tokens, + context=prefill_context, + ) if getattr(request, "cached_tokens", 0) > 0: with mx.stream(self._stream): _materialize_cache_storage(prompt_cache) + self._log_mid_prefill_summary(prefill_context) return prompt_cache, last_token @@ -3460,7 +4182,13 @@ def _do_external_prefill( _MEMORY_ADMISSION_STALL_TIMEOUT_S: float = 60.0 _STORE_CACHE_ADMISSION_STALL_TIMEOUT_S: float = 60.0 - def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float: + def _predicted_chunk_transient( + self, + n_tokens: int, + kv_len: int, + *, + phase: _PrefillKVPhase = _PrefillKVPhase.DENSE, + ) -> float: """Conservative predicted Metal peak growth for one prefill chunk. The per-chunk SDPA/MoE transient scales with ``query_len * kv_len``, so @@ -3478,7 +4206,17 @@ def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float: if n_tokens <= 0: return 0.0 per_token = 0.0 - tracker = self._prefill_transient_tracker + tracker_selector = getattr( + self, + "_prefill_transient_tracker_for_phase", + None, + ) + if callable(tracker_selector): + tracker = tracker_selector(phase) + elif phase is _PrefillKVPhase.TURBOQUANT: + tracker = getattr(self, "_prefill_tq_transient_tracker", None) + else: + tracker = getattr(self, "_prefill_transient_tracker", None) if tracker is not None: if tracker.last_n_tokens > 0 and tracker.last_delta_bytes > 0: per_token = max( @@ -3490,11 +4228,42 @@ def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float: static = self.memory_monitor.estimate_chunk_transient_bytes( n_tokens, kv_len + n_tokens ) - static += self.memory_monitor.estimate_prompt_kv_bytes(n_tokens) + if phase is _PrefillKVPhase.TURBOQUANT: + raw_bits = getattr(self, "_turboquant_kv_bits", None) + tq_workspace_estimator = getattr( + self.memory_monitor, + "estimate_turboquant_prefill_attention_bytes", + None, + ) + if ( + isinstance(raw_bits, (int, float)) + and not isinstance(raw_bits, bool) + and callable(tq_workspace_estimator) + ): + static = max( + static, + tq_workspace_estimator( + n_tokens, + kv_len + n_tokens, + bits=float(raw_bits), + ), + ) + width_selector = getattr(self, "_prefill_phase_dtype_size", None) + dtype_size = width_selector(phase) if callable(width_selector) else None + static += self.memory_monitor.estimate_prompt_kv_bytes( + n_tokens, + dtype_size=dtype_size, + ) per_token = max(per_token, float(static) / n_tokens) return per_token * n_tokens * self._PREFILL_TRANSIENT_SAFETY - def _admission_transient_bound(self, n_tokens: int, kv_len: int) -> float: + def _admission_transient_bound( + self, + n_tokens: int, + kv_len: int, + *, + phase: _PrefillKVPhase = _PrefillKVPhase.DENSE, + ) -> float: """Transient charge for admission and the guard's pass/abort gates. The largest FLOOR-SIZE chunk transient observed this session is a @@ -3514,8 +4283,22 @@ def _admission_transient_bound(self, n_tokens: int, kv_len: int) -> float: shrink arithmetic stay on ``_predicted_chunk_transient``; a flat size-invariant bound would zero out their proportional response. """ - bound = self._predicted_chunk_transient(n_tokens, kv_len) - tracker = self._prefill_transient_tracker + bound = self._predicted_chunk_transient( + n_tokens, + kv_len, + phase=phase, + ) + tracker_selector = getattr( + self, + "_prefill_transient_tracker_for_phase", + None, + ) + if callable(tracker_selector): + tracker = tracker_selector(phase) + elif phase is _PrefillKVPhase.TURBOQUANT: + tracker = getattr(self, "_prefill_tq_transient_tracker", None) + else: + tracker = getattr(self, "_prefill_transient_tracker", None) if tracker is not None: bound = max(bound, float(tracker.observed_max_bytes)) return bound @@ -3530,6 +4313,20 @@ def _prefill_abort_cap(self) -> int: cap = self._memory_abort_limit_bytes or self._memory_hard_limit_bytes return int(cap * self._prefill_abort_margin) if cap > 0 else 0 + def _prefill_sizing_target(self) -> int: + """Return the existing reserve-aware chunk-sizing pressure boundary.""" + hard_cap = self._memory_hard_limit_bytes + if hard_cap <= 0: + return 0 + headroom_safety = getattr( + self, + "_prefill_headroom_safety", + self._PREFILL_HEADROOM_SAFETY, + ) + target = int(hard_cap * headroom_safety) + abort_cap = self._prefill_abort_cap() + return min(target, abort_cap) if abort_cap > 0 else target + def _admission_limit_bytes(self) -> int: """The line admission estimates must stay under. @@ -3581,13 +4378,7 @@ def _sdpa256_unfused_headroom(self) -> int: ) return _SDPA256_UNBOUNDED_HEADROOM return -1 - headroom_safety = getattr( - self, "_prefill_headroom_safety", self._PREFILL_HEADROOM_SAFETY - ) - target = int(hard_cap * headroom_safety) - abort_cap = self._prefill_abort_cap() - if abort_cap > 0: - target = min(target, abort_cap) + target = self._prefill_sizing_target() return target - self._current_usage_bytes() _MAX_PREFILL_EVICTION_RETRIES = 1 @@ -3601,11 +4392,14 @@ def _raise_prefill_eviction_if_available( predicted_transient: int, requested_tokens: int, reason: str, + processed_tokens: int = 0, ) -> None: """Pause a request once so EngineCore can evict idle LRU models.""" request = self.requests.get(request_id) if request is None: return + if getattr(self, "_prefill_eviction_callback_configured", None) is False: + return max_retries = getattr( self, "_MAX_PREFILL_EVICTION_RETRIES", @@ -3626,6 +4420,7 @@ def _raise_prefill_eviction_if_available( predicted_transient_bytes=int(predicted_transient), requested_tokens=int(requested_tokens), reason=reason, + processed_tokens=max(0, int(processed_tokens)), ) logger.info( "Request %s needs prefill headroom before throttling " @@ -3646,34 +4441,125 @@ def _guard_prefill_chunk( progress: int, loop_label: str, request_id: str | None = None, + request: "Request | None" = None, + prompt_cache: list[Any] | None = None, + prefill_context: _PrefillContext | None = None, ) -> int: - """Clamp/abort a prefill chunk so its predicted peak can never reach - the physical Metal cap (the uncatchable async OOM crash). - - Returns a chunk size whose predicted peak fits under the margined cap - (possibly shrunk from ``n_tokens``). If even the minimum chunk would - not fit after a reclaim, raises a clean RuntimeError — the context is - genuinely too large for available memory. That message intentionally - does NOT contain "Memory limit exceeded", so ``_requeue_or_fail_prefill`` - fails it fast with a clear error rather than looping a doomed retry. - """ + """Return a chunk whose predicted peak stays below the physical cap.""" base_cap, cap, margin = self._prefill_abort_description() if cap <= 0: return n_tokens - # Speed priority: no shrinking — the floor IS the full chunk, so the - # abort gate below charges the full-step transient and the shrink - # math degenerates to returning n_tokens unchanged. + trigger_target = self._prefill_sizing_target() or cap if self._prefill_speed_priority: min_chunk = n_tokens else: min_chunk = max(1, self._prefill_min_chunk_tokens) + phase = ( + prefill_context.phase + if prefill_context is not None + else _PrefillKVPhase.DENSE + ) current = self._current_usage_bytes() - if current + self._admission_transient_bound(n_tokens, kv_len) <= cap: + full_transient = self._admission_transient_bound( + n_tokens, + kv_len, + phase=phase, + ) + if current + full_transient <= trigger_target: return n_tokens - # Predicted to breach — reclaim transients and re-measure once. + # Reclaim before either conversion or a final-cap decision. current = self._reclaim_prefill_headroom() - min_transient = self._admission_transient_bound(min_chunk, kv_len) + if current + full_transient <= trigger_target: + return n_tokens + + request_obj = request + if request_obj is None and request_id is not None: + request_obj = getattr(self, "requests", {}).get(request_id) + conversion_available = getattr( + self, + "_mid_prefill_conversion_available", + None, + ) + if callable(conversion_available) and conversion_available( + prefill_context, + prompt_cache, + request_obj, + ): + maybe_raise_eviction = getattr( + self, "_raise_prefill_eviction_if_available", None + ) + if request_id is not None and callable(maybe_raise_eviction): + maybe_raise_eviction( + request_id=request_id, + current=current, + target_cap=trigger_target, + predicted_transient=int(full_transient), + requested_tokens=n_tokens, + reason="turboquant_mid_prefill", + processed_tokens=progress, + ) + if ( + request_obj is not None + and prompt_cache is not None + and prefill_context is not None + ): + self._attempt_mid_prefill_conversion( + request=request_obj, + prompt_cache=prompt_cache, + context=prefill_context, + processed_tokens=progress, + safety_cap=cap, + ) + resized = self._adaptive_chunk_size( + n_tokens, + request_id=request_id or prefill_context.request_id, + loop_label=loop_label, + kv_len=kv_len, + prefill_context=prefill_context, + ) + return self._guard_prefill_chunk( + resized, + kv_len=kv_len, + progress=progress, + loop_label=loop_label, + request_id=request_id, + request=request_obj, + prompt_cache=prompt_cache, + prefill_context=prefill_context, + ) + + if current + full_transient <= cap: + # Adaptive sizing may have preserved this candidate solely so the + # guard could convert it. If the live cache is still empty or no + # longer eligible, restore ordinary target-based sizing. + resized = self._adaptive_chunk_size( + n_tokens, + request_id=request_id + or (prefill_context.request_id if prefill_context is not None else ""), + loop_label=loop_label, + kv_len=kv_len, + prefill_context=prefill_context, + defer_mid_prefill_conversion=False, + ) + if resized < n_tokens: + return self._guard_prefill_chunk( + resized, + kv_len=kv_len, + progress=progress, + loop_label=loop_label, + request_id=request_id, + request=request_obj, + prompt_cache=prompt_cache, + prefill_context=prefill_context, + ) + return n_tokens + + min_transient = self._admission_transient_bound( + min_chunk, + kv_len, + phase=phase, + ) if current + min_transient > cap: maybe_raise_eviction = getattr( self, "_raise_prefill_eviction_if_available", None @@ -3701,8 +4587,6 @@ def _guard_prefill_chunk( ) binding_str, advice = describe_ceiling_binding( static=self._memory_static_ceiling_bytes, - # See _preflight_safety_rejection: the abort limit this cap - # derives from is min(static, metal_cap) by design. dynamic=( 0 if self._memory_abort_limit_bytes @@ -3729,13 +4613,16 @@ def _guard_prefill_chunk( limit_bytes=int(cap), ) - # The floor fits — pick the largest chunk that still fits under the cap. - per_token = self._predicted_chunk_transient(n_tokens, kv_len) / n_tokens + per_token = ( + self._predicted_chunk_transient( + n_tokens, + kv_len, + phase=phase, + ) + / n_tokens + ) safe_n = int((cap - current) / per_token) if per_token > 0 else n_tokens n_fit = max(min_chunk, min(n_tokens, safe_n)) - # Same quantization as the adaptive throttle: an off-grid size here - # would reintroduce the near-miss buffers _snap_chunk_size exists to - # avoid. n_fit = self._snap_chunk_size(n_fit, n_tokens) if n_fit < n_tokens: logger.debug( @@ -3785,6 +4672,8 @@ def _adaptive_chunk_size( request_id: str, loop_label: str, kv_len: int = 0, + prefill_context: _PrefillContext | None = None, + defer_mid_prefill_conversion: bool = True, ) -> int: """Size the next prefill chunk so its predicted peak stays under a safety margin below the hard cap. @@ -3841,18 +4730,24 @@ def _adaptive_chunk_size( # safety) — see _predicted_chunk_transient. Anchored on the most recent # measurement so it tracks growth with kv_len instead of lagging behind # a long-run average. - per_token = self._predicted_chunk_transient(requested, kv_len) / requested + phase = ( + prefill_context.phase + if prefill_context is not None + else _PrefillKVPhase.DENSE + ) + per_token = ( + self._predicted_chunk_transient( + requested, + kv_len, + phase=phase, + ) + / requested + ) predictor = "measured" if per_token > 0 else "none" - # Keep each chunk's predicted peak under the LOWER of the dynamic - # throttle target and the prefill safety cap, so the peak can never - # reach the Metal wall (the uncatchable async OOM). - headroom_safety = getattr( - self, "_prefill_headroom_safety", self._PREFILL_HEADROOM_SAFETY - ) - safe_target = int(hard_cap * headroom_safety) - abort_cap = self._prefill_abort_cap() - target = min(safe_target, abort_cap) if abort_cap > 0 else safe_target + # Use the same reserve-aware target as the mid-prefill trigger. The + # stable abort cap remains the final pass/reject boundary. + target = self._prefill_sizing_target() soft_watermark = int(soft_base * self._prefill_safe_zone_ratio) if per_token <= 0: @@ -3884,6 +4779,24 @@ def _adaptive_chunk_size( requested_tokens=requested, reason="adaptive_prefill_throttle", ) + conversion_available = getattr( + self, + "_mid_prefill_conversion_available", + None, + ) + if ( + defer_mid_prefill_conversion + and callable(conversion_available) + and conversion_available( + prefill_context, + None, + getattr(self, "requests", {}).get(request_id), + ) + ): + # The one-shot LRU call above has had its chance. Preserve the + # full candidate so the common guard can remeasure, require + # populated history, and run the request's sole conversion. + return requested if self._prefill_speed_priority: # Speed priority: never shrink. Idle-model eviction above # still gets its chance to free memory; beyond that, the @@ -3956,7 +4869,7 @@ def _adaptive_chunk_size( predictor, per_token / 1024, current / 1024**3, - safe_target / 1024**3, + target / 1024**3, hard_cap / 1024**3, kv_len, band_ratio, @@ -3998,6 +4911,8 @@ def _current_usage_bytes(self, *, refresh_mlx_active: bool = True) -> int: Scheduler steps run on the MLX executor thread, so they can refresh mx.get_active_memory() safely. Event-loop callers such as early preflight use the cached executor sample and phys_footprint instead. + Every caller also sees a process-wide conversion peak reservation; + the conversion holder releases its reservation before its post sample. """ active = self._last_mlx_active_memory_bytes if refresh_mlx_active: @@ -4009,7 +4924,7 @@ def _current_usage_bytes(self, *, refresh_mlx_active: bool = True) -> int: else: hot_cache_bytes = Scheduler._hot_cache_cpu_bytes(self) phys = max(0, int(get_phys_footprint()) - hot_cache_bytes) - return max(active, phys) + return max(active, phys) + _conversion_coordinator.outstanding_bytes() def get_active_hot_cache_block_hashes(self) -> set[bytes]: """Return hot-cache block hashes owned by active in-flight requests.""" @@ -4216,6 +5131,7 @@ def _record_chunk_transient( loop_label: str, kv_len: int = 0, requested_step: int | None = None, + phase: _PrefillKVPhase = _PrefillKVPhase.DENSE, ) -> None: """Feed one chunk's measured transient into the EWMA tracker. @@ -4274,9 +5190,20 @@ def _record_chunk_transient( requested_step, ) return - self._prefill_transient_tracker.update( - n_tokens, delta, floor_sample=n_tokens <= min_chunk + tracker_selector = getattr( + self, + "_prefill_transient_tracker_for_phase", + None, ) + if callable(tracker_selector): + tracker = tracker_selector(phase) + elif phase is _PrefillKVPhase.TURBOQUANT: + tracker = getattr(self, "_prefill_tq_transient_tracker", None) + else: + tracker = getattr(self, "_prefill_transient_tracker", None) + if tracker is None: + return + tracker.update(n_tokens, delta, floor_sample=n_tokens <= min_chunk) logger.debug( "[throttle:%s] measure rid=%s n=%d kv_len=%d transient=%.2fMB per_token=%.1fKB ewma=%.1fKB observed_max=%.1fMB samples=%d", loop_label, @@ -4285,9 +5212,9 @@ def _record_chunk_transient( kv_len, delta / 1024**2, (delta / max(n_tokens, 1)) / 1024, - self._prefill_transient_tracker.bytes_per_token / 1024, - self._prefill_transient_tracker.observed_max_bytes / 1024**2, - self._prefill_transient_tracker.samples, + tracker.bytes_per_token / 1024, + tracker.observed_max_bytes / 1024**2, + tracker.samples, ) def _maybe_record_fixed_state_bytes(self, cache_list: Any) -> None: @@ -4322,6 +5249,18 @@ def _walk(c: Any) -> int: self._fixed_state_recorded = True if total > 0: self.memory_monitor.set_fixed_state_bytes(total) + manager = getattr(self, "paged_ssd_cache_manager", None) + resize_writer_queue = getattr( + manager, + "set_expected_block_payload_bytes", + None, + ) + if callable(resize_writer_queue): + block_tokens = max(1, int(self.config.paged_cache_block_size)) + payload = self.memory_monitor.estimate_paged_writer_block_memory( + block_tokens + ) + resize_writer_queue(math.ceil(payload)) logger.debug( "Fixed recurrent state measured: %.1fMB per sequence", total / 1024**2, @@ -4404,6 +5343,11 @@ def _begin_prefill( if existing_cache is not None else make_prompt_cache(self.model) ) + prefill_context = self._new_prefill_context( + request, + prompt_cache, + loop_label="chunked_step", + ) block_size = self.config.paged_cache_block_size boundary_enabled = ( @@ -4444,6 +5388,7 @@ def _begin_prefill( boundary_enabled=boundary_enabled, block_size=block_size, total_length=len(tokens), + prefill_context=prefill_context, ) def _step_prefill_chunk(self, state: _PrefillState) -> bool: @@ -4489,6 +5434,7 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: request_id=state.request.request_id, loop_label="chunked_step", kv_len=state.base_size + state.tokens_processed, + prefill_context=state.prefill_context, ) # Pre-chunk safety guard (mirrors the external loop): never submit a @@ -4499,6 +5445,9 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: progress=state.tokens_processed, loop_label="chunked_step", request_id=state.request.request_id, + request=state.request, + prompt_cache=state.cache, + prefill_context=state.prefill_context, ) _throttle_pre = get_phys_footprint() @@ -4509,18 +5458,37 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: with mx.stream(self._stream): chunk = state.tokens_remaining[:, :n] state.tokens_remaining = state.tokens_remaining[:, n:] + measure_post_trigger = ( + state.prefill_context is not None + and state.prefill_context.mid_triggered + and state.prefill_context.phase is _PrefillKVPhase.TURBOQUANT + ) self.model(chunk, cache=state.cache) mx.eval([c.state for c in state.cache]) _throttle_post = get_phys_footprint() - self._record_chunk_transient( - n, - _throttle_pre, - _throttle_post, - request_id=state.request.request_id, - loop_label="chunked_step", - kv_len=state.base_size + state.tokens_processed, - requested_step=prefill_step_size, - ) + if state.prefill_context is None: + self._record_chunk_transient( + n, + _throttle_pre, + _throttle_post, + request_id=state.request.request_id, + loop_label="chunked_step", + kv_len=state.base_size + state.tokens_processed, + requested_step=prefill_step_size, + ) + else: + self._record_chunk_transient( + n, + _throttle_pre, + _throttle_post, + request_id=state.request.request_id, + loop_label="chunked_step", + kv_len=state.base_size + state.tokens_processed, + requested_step=prefill_step_size, + phase=state.prefill_context.phase, + ) + if measure_post_trigger and state.prefill_context is not None: + state.prefill_context.post_trigger_tokens += n self._maybe_record_fixed_state_bytes(state.cache) state.tokens_processed += n @@ -4546,11 +5514,7 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: state.request.request_id, state.tokens_processed, state.total_length - 1, - ( - self.config.model_name - if self.config.model_name - else "" - ), + (self.config.model_name if self.config.model_name else ""), ) # Memory monitoring — use max(active, phys_footprint) so MLX cache @@ -4603,9 +5567,7 @@ def _step_prefill_chunk(self, state: _PrefillState) -> bool: # Speed priority runs full chunks through this caution band # by design — the per-chunk notice is DEBUG there, not a # warning about an unexpected state. - _log = ( - logger.debug if self._prefill_speed_priority else logger.warning - ) + _log = logger.debug if self._prefill_speed_priority else logger.warning _log( f"Chunked prefill above max_bytes at " f"{state.tokens_processed} tokens: " @@ -4634,15 +5596,22 @@ def _emit_final_boundary_if_needed(self, state: _PrefillState) -> None: ) def _finalize_chunked_prefill_cache_for_insert( - self, request: "Request", prompt_cache: list[Any] | None + self, + request: "Request", + prompt_cache: list[Any] | None, + *, + processed_tokens: int = 0, + prefill_context: _PrefillContext | None = None, ) -> None: - """Mirror external prefill's post-prefill cache epilogue.""" - if not prompt_cache or self._turboquant_kv_bits is None: - return - if not self._turboquant_eligible(prompt_cache): + """Mirror external prefill's typed post-prefill cache epilogue.""" + if not prompt_cache: return - - self._apply_turboquant_kv_convert(prompt_cache) + self._finalize_turboquant_prefill_cache( + request, + prompt_cache, + processed_tokens=processed_tokens, + context=prefill_context, + ) if getattr(request, "cached_tokens", 0) > 0: with mx.stream(self._stream): _materialize_cache_storage(prompt_cache) @@ -4699,9 +5668,18 @@ def _insert_prefilled_request( request.request_id, vlm_mtp_uid, ) + if state.prefill_context is not None: + self._log_mid_prefill_summary(state.prefill_context) return - self._finalize_chunked_prefill_cache_for_insert(request, state.cache) + self._finalize_chunked_prefill_cache_for_insert( + request, + state.cache, + processed_tokens=state.tokens_processed, + prefill_context=state.prefill_context, + ) + if state.prefill_context is not None: + self._log_mid_prefill_summary(state.prefill_context) per_row_lps = state.per_row_lps if state.per_row_lps is not None else [] # insert() merges the prompt cache into the batch KV caches with lazy @@ -4794,7 +5772,7 @@ def _advance_chunked_prefills( still_prefilling.append(request) still_prefilling.extend(pending_prefills[index + 1 :]) logger.info( - "Paused chunked prefill request %s for LRU eviction " "(reason=%s)", + "Paused chunked prefill request %s for LRU eviction (reason=%s)", rid, e.request.reason, ) @@ -4864,7 +5842,24 @@ def _advance_chunked_prefills( # Clean up the prefill-progress tracker entry. get_prefill_tracker().remove(rid) - self._insert_prefilled_request(request, state, scheduled) + try: + self._insert_prefilled_request(request, state, scheduled) + except _PrefillAbortedError: + _sync_and_clear_cache(self._stream) + self._cleanup_prefill_abort_request(request) + continue + except PrefillMemoryExceededError as exc: + logger.error( + "Chunked prefill completion capacity rejected for %s: %s", + rid, + exc, + ) + self._release_paged_cache_for_request(rid) + self.requests.pop(rid, None) + self._clear_request_admission_bookkeeping(rid) + _sync_and_clear_cache(self._stream) + rejected.append(_prefill_memory_exception_output(rid, exc)) + continue self.prefilling = still_prefilling @@ -4934,8 +5929,7 @@ def _buffer_stop_sequence_output( pending_tokens = tuple(token for token, _ in pending) reported_match = tuple( - int(token) - for token in (getattr(response, "match_sequence", None) or ()) + int(token) for token in (getattr(response, "match_sequence", None) or ()) ) matched_sequence = ( reported_match @@ -4967,7 +5961,7 @@ def strip_matched_prefix(text: str) -> str: suppressed_stream_text ): terminal_output.output_text = terminal_output.output_text[ - :-len(suppressed_stream_text) + : -len(suppressed_stream_text) ] else: matched_prefix_tokens = matched_sequence[:-1] @@ -4980,7 +5974,7 @@ def strip_matched_prefix(text: str) -> str: and request.request_id not in self._output_parser_sessions ): terminal_output.output_text = self.tokenizer.decode( - output_token_ids[:-len(matched_prefix_tokens)] + output_token_ids[: -len(matched_prefix_tokens)] ) request.output_text = terminal_output.output_text @@ -5657,6 +6651,7 @@ def _extract_prefill_snapshot_states( pre-extracted marker alongside raw decode-path snapshots (those are already decoupled copies from ``extract_cache``). """ + def _copy_containers(value: Any) -> Any: # ArraysCache.state returns its live slot LIST (not a copy); # the model rebinds slots in place, so container structure @@ -6329,9 +7324,9 @@ def _validate_cache(self, cache: Any) -> bool: expected_cache = make_prompt_cache(self.model) except Exception: expected_cache = None - if isinstance(expected_cache, (list, tuple)) and len( - expected_cache - ) == len(cache): + if isinstance(expected_cache, (list, tuple)) and len(expected_cache) == len( + cache + ): arrays_names = {"ArraysCache", "SizedArraysCache"} for layer_cache, expected_layer in zip(cache, expected_cache): if ( @@ -7246,7 +8241,9 @@ def set_specprefill_draft_model( draft_cache_list, model_name=name, ) - draft_layer_cache_types = draft_model_cache_config.get_type_names() + draft_layer_cache_types = ( + draft_model_cache_config.get_type_names() + ) except Exception as e: logger.debug( "Could not infer SpecPrefill draft cache layout: %s", e @@ -7383,9 +8380,7 @@ def _route_to_vlm_mtp( "logits processors without vlm_mtp support (%s); falling " "back to BatchGenerator", request.request_id, - ", ".join( - type(proc).__name__ for proc in unsupported_processors - ), + ", ".join(type(proc).__name__ for proc in unsupported_processors), ) return None @@ -8052,8 +9047,7 @@ def _process_pending_reclaim(self) -> None: logger.warning(f"Idle reclaim failed: {e}") return logger.info( - "Idle reclaim: trimmed Metal transients between turns " - "(%.1fGB -> %.1fGB)", + "Idle reclaim: trimmed Metal transients between turns (%.1fGB -> %.1fGB)", before / 1024**3, after / 1024**3, ) @@ -8211,9 +9205,7 @@ def has_pending_route_preflight_cleanup(self) -> bool: Active requests are intentionally not included: their memory is live and must remain charged to a concurrent admission. """ - return bool( - self._pending_async_removes or self._deferred_clear_at is not None - ) + return bool(self._pending_async_removes or self._deferred_clear_at is not None) def refresh_route_preflight_usage(self) -> int: """Publish a fresh MLX memory sample for route-level retry. @@ -8387,12 +9379,23 @@ def _preflight_memory_check( prompt_tokens = request.num_prompt_tokens cached_tokens = request.cached_tokens or 0 + phase = _PrefillKVPhase.DENSE + conversion_eligible = self._turboquant_preflight_conversion_eligible + if getattr(self, "_turboquant_mid_prefill", False): + prompt_cache = getattr(request, "prompt_cache", None) + if isinstance(prompt_cache, list): + cache_phase, conversion_eligible = self._classify_prefill_cache( + prompt_cache + ) + if cache_phase is _PrefillKVPhase.TURBOQUANT: + phase = cache_phase current = self._current_usage_bytes() est = self._admission_estimate( num_prompt_tokens=prompt_tokens, cached_tokens=cached_tokens, current=current, + phase=phase, ) if est is None: return None # can't estimate, skip @@ -8415,6 +9418,12 @@ def _preflight_memory_check( requested_tokens=est.floor_chunk, reason="prefill_preflight", ) + if self._can_defer_mid_prefill_preflight( + phase, + request=request, + conversion_eligible=conversion_eligible, + ): + return None message = self._format_rejection_message( estimated=est.estimated, @@ -8442,6 +9451,12 @@ def _preflight_memory_check( requested_tokens=est.floor_chunk, reason="prefill_safety_cap", ) + if self._can_defer_mid_prefill_preflight( + phase, + request=request, + conversion_eligible=conversion_eligible, + ): + return None return safety_rejection return None @@ -8451,6 +9466,7 @@ def _admission_estimate( num_prompt_tokens: int, cached_tokens: int, current: int, + phase: _PrefillKVPhase = _PrefillKVPhase.DENSE, ) -> _AdmissionEstimate | None: """Deterministic admission estimate shared by every preflight path. @@ -8501,9 +9517,19 @@ def _admission_estimate( floor_chunk = min(charge_tokens, prefill_tokens) kv_len = max(int(num_prompt_tokens) - 1 - floor_chunk, 0) kv_exact = int( - monitor.estimate_resident_kv_bytes(new_tokens, chunk_tokens=floor_chunk) + monitor.estimate_resident_kv_bytes( + new_tokens, + chunk_tokens=floor_chunk, + dtype_size=self._prefill_phase_dtype_size(phase), + ) + ) + transient = int( + self._admission_transient_bound( + floor_chunk, + kv_len, + phase=phase, + ) ) - transient = int(self._admission_transient_bound(floor_chunk, kv_len)) if kv_exact <= 0 and transient <= 0: return None return _AdmissionEstimate( @@ -8598,6 +9624,11 @@ def preflight_or_raise( admission_limit = self._admission_limit_bytes() if est.estimated > admission_limit: + if self._can_defer_mid_prefill_preflight( + _PrefillKVPhase.DENSE, + conversion_eligible=self._turboquant_preflight_conversion_eligible, + ): + return message = self._format_rejection_message( estimated=est.estimated, current=current, @@ -8625,10 +9656,14 @@ def preflight_or_raise( ) if safety_rejection is None: return + if self._can_defer_mid_prefill_preflight( + _PrefillKVPhase.DENSE, + conversion_eligible=self._turboquant_preflight_conversion_eligible, + ): + return logger.warning( - "Preflight safety-cap rejected (%d tokens, cached=%d, " - "request_id=%s): %s", + "Preflight safety-cap rejected (%d tokens, cached=%d, request_id=%s): %s", num_prompt_tokens, cached_tokens, request_id, @@ -9150,7 +10185,9 @@ def _report_sparse_progress(processed: int, total: int) -> None: check_abort=_check_specprefill_abort, report_system_progress=_report_system_progress, report_sparse_progress=_report_sparse_progress, - sync_and_clear_cache=lambda: _sync_and_clear_cache(self._stream), + sync_and_clear_cache=lambda: _sync_and_clear_cache( + self._stream + ), log=logger, extract_cache_states=self._extract_cache_states, # Preserve an ordinary cache hit that already extends @@ -9237,9 +10274,24 @@ def _report_sparse_progress(processed: int, total: int) -> None: ): sm = self._build_state_machine(request) per_row_lps = list(logits_processors) if logits_processors else [] - state = self._begin_prefill( - request, tokens_to_process, cache_to_use - ) + try: + state = self._begin_prefill( + request, tokens_to_process, cache_to_use + ) + except PrefillMemoryExceededError as e: + logger.error( + "Chunked prefill setup capacity rejected for %s: %s", + request.request_id, + e, + ) + self._release_paged_cache_for_request(request.request_id) + self.requests.pop(request.request_id, None) + self._clear_request_admission_bookkeeping(request.request_id) + get_prefill_tracker().remove(request.request_id) + rejected_outputs.append( + _prefill_memory_exception_output(request.request_id, e) + ) + continue state.sampler = sampler state.sm = sm state.per_row_lps = per_row_lps @@ -9313,7 +10365,31 @@ def _report_sparse_progress(processed: int, total: int) -> None: self._emit_final_boundary_if_needed(state) _sync_and_clear_cache(self._stream) get_prefill_tracker().remove(request.request_id) - self._insert_prefilled_request(request, state, scheduled) + try: + self._insert_prefilled_request(request, state, scheduled) + except _PrefillAbortedError: + _sync_and_clear_cache(self._stream) + self._cleanup_prefill_abort_request(request) + continue + except PrefillMemoryExceededError as e: + logger.error( + "Chunked prefill completion capacity rejected " + "for %s: %s", + request.request_id, + e, + ) + self._release_paged_cache_for_request(request.request_id) + self.requests.pop(request.request_id, None) + self._clear_request_admission_bookkeeping( + request.request_id + ) + rejected_outputs.append( + _prefill_memory_exception_output( + request.request_id, + e, + ) + ) + continue else: self.prefilling.append(request) self._prefill_states[request.request_id] = state @@ -9342,6 +10418,9 @@ def _report_sparse_progress(processed: int, total: int) -> None: self.request_id_to_uid.pop(request.request_id, None) self._release_paged_cache_for_request(request.request_id) get_prefill_tracker().remove(request.request_id) + self._reset_partial_external_prefill_for_eviction( + request, e.request + ) self._pause_for_prefill_eviction(request, e.request) break except PrefillMemoryExceededError as e: @@ -9972,11 +11051,7 @@ def _cleanup_finished(self, finished_ids: set[str]) -> None: ) ) if intermediate_snapshots is not None: - for ( - snapshot_cache - ) in ( - intermediate_snapshots.iter_in_memory_extracted() - ): + for snapshot_cache in intermediate_snapshots.iter_in_memory_extracted(): pre_eval_arrays.extend( self._collect_arrays_from_extracted_cache( snapshot_cache @@ -10600,6 +11675,38 @@ def _requeue_or_fail_prefill(self, request: "Request", error: Exception) -> bool ) return True + def _reset_partial_external_prefill_for_eviction( + self, + request: "Request", + eviction: PrefillEvictionRequest, + ) -> None: + """Drop a restored prefix mutated before an external-prefill pause.""" + if eviction.processed_tokens <= 0 or request.prompt_cache is None: + return + + saved = getattr(request, "_prefill_saved_rope_deltas", None) + if saved is not None: + lm = getattr(self.model, "_language_model", None) + if lm is not None and hasattr(lm, "_rope_deltas"): + lm._rope_deltas = saved + request._prefill_saved_rope_deltas = None + + request.prompt_cache = None + request.cached_tokens = 0 + request.remaining_tokens = request.prompt_token_ids + request.block_table = None + request.shared_prefix_blocks = 0 + request._extracted_cache = None + request._model_cache_config = None + self._prefix_cache_prepared.discard(request.request_id) + _sync_and_clear_cache(self._stream) + logger.info( + "Discarded partially advanced prefix for %s before prefill retry " + "(processed_tokens=%d)", + request.request_id, + eviction.processed_tokens, + ) + def _pause_for_prefill_eviction( self, request: "Request", @@ -10607,11 +11714,9 @@ def _pause_for_prefill_eviction( ) -> None: """Hold a request until EngineCore can evict idle models asynchronously. - The request's prefix-cache state (prompt_cache, block_table, - cached_tokens, remaining_tokens) is deliberately left untouched so - a reconstructed prefix survives the pause and the retry prefills - only the uncached suffix instead of recomputing the prompt cold - (#2180). + An unmodified reconstructed prefix survives a pause before the first + forward. The external-prefill catch resets any prefix already extended + by completed chunks so retry cannot replay the same suffix twice. """ self._pending_prefill_eviction_request = eviction request.status = RequestStatus.WAITING @@ -10666,8 +11771,15 @@ def step(self) -> SchedulerOutput: if self.prefilling: self._advance_chunked_prefills(chunked_scheduled, chunked_rejected) - # Schedule waiting requests - scheduled, rejected = self._schedule_waiting() + # A chunked prefill may have paused for async LRU eviction. Deliver + # that request before inspecting the waiting queue: admitting a + # second pressured request here could overwrite the single pending + # hand-off before EngineCore observes it. The pending request is + # consumed below, so waiting admission resumes on the next step. + if self._pending_prefill_eviction_request is None: + scheduled, rejected = self._schedule_waiting() + else: + scheduled, rejected = [], [] # Merge chunked-prefill completions into the scheduled list. if chunked_scheduled: scheduled = chunked_scheduled + scheduled @@ -10791,8 +11903,7 @@ def step(self) -> SchedulerOutput: finished=True, finish_reason="error", error=( - f"Cache corruption not recoverable " - f"after retries: {e}" + f"Cache corruption not recoverable after retries: {e}" ), ) ) @@ -10835,7 +11946,7 @@ def step(self) -> SchedulerOutput: import traceback logger.error( - f"Error in batch generation step: {e}\n" f"{traceback.format_exc()}" + f"Error in batch generation step: {e}\n{traceback.format_exc()}" ) raise @@ -11133,6 +12244,30 @@ def _set_model_info_for_monitor(self) -> None: if self.memory_monitor is None: return + cache_list_for_tq: list[Any] | None = None + self._turboquant_preflight_conversion_eligible = None + try: + cache_factory = getattr(self.model, "make_cache", None) + except Exception: + cache_factory = None + cache_factory_available = callable(cache_factory) + if cache_factory_available: + try: + cache_candidate = cache_factory() + except Exception: + pass + else: + if isinstance(cache_candidate, list): + cache_list_for_tq = cache_candidate + try: + self._turboquant_preflight_conversion_eligible = ( + self._confirm_turboquant_preflight_conversion_eligibility( + cache_candidate + ) + ) + except Exception: + self._turboquant_preflight_conversion_eligible = None + try: # Try to get model config config = None @@ -11195,6 +12330,8 @@ def _cfg_get(obj: Any, key: str, default: Any = None) -> Any: elif self.model.dtype == mx.bfloat16: base_dtype_size = 2 dtype_size = base_dtype_size + self._prefill_dense_kv_dtype_size = base_dtype_size + self._prefill_tq_kv_dtype_size = None # Extract num_attention_heads (query heads) for SDPA peak estimation num_attention_heads = ( @@ -11204,19 +12341,17 @@ def _cfg_get(obj: Any, key: str, default: Any = None) -> Any: ) # Classify layer cache types for hybrid models - cache_list_for_tq = None + # Reuse the one-shot make_cache probe performed above. actual_kv_cache_layers = None num_kv_cache_layers = num_layers rotating_layer_specs: list[tuple[int, int]] = [] arrays_cache_layers = 0 - if not hasattr(self.model, "make_cache"): + if not cache_factory_available: actual_kv_cache_layers = num_layers - elif collect_kv_layer_specs is not None: + elif cache_list_for_tq is not None and collect_kv_layer_specs is not None: try: - cache_list = self.model.make_cache() - cache_list_for_tq = cache_list full_layers, rotating_layer_specs, arrays_cache_layers = ( - collect_kv_layer_specs(cache_list) + collect_kv_layer_specs(cache_list_for_tq) ) actual_kv_cache_layers = full_layers num_kv_cache_layers = full_layers @@ -11249,7 +12384,10 @@ def _cfg_get(obj: Any, key: str, default: Any = None) -> Any: ) ) ): - tq_dtype_size = float(self._turboquant_kv_bits) / 8.0 + (2.0 / head_dim) + tq_dtype_size = turboquant_mse_bytes_per_element( + head_dim, + float(self._turboquant_kv_bits), + ) if ( self._turboquant_skip_last and not isinstance(actual_kv_cache_layers, bool) @@ -11260,6 +12398,7 @@ def _cfg_get(obj: Any, key: str, default: Any = None) -> Any: ) / actual_kv_cache_layers else: dtype_size = tq_dtype_size + self._prefill_tq_kv_dtype_size = dtype_size kv_bytes_per_token = ( estimate_mla_kv_bytes_per_token( @@ -11519,12 +12658,15 @@ def _init_tiered_cache(self) -> bool: # happy path here is ``has_model_info() is True``; this # else branch only fires for skeletal test fixtures. if self.memory_monitor is not None and self.memory_monitor.has_model_info(): - # ``estimate_block_memory(1)`` returns all-layers K+V - # bytes for a single token at the dtype the monitor was - # configured with — exactly the per-token cost the - # queue cap needs to weigh. - expected_kv_bytes_per_token = self.memory_monitor.estimate_block_memory( - 1 + # SSD boundary blocks also retain rotating windows and fixed + # recurrent snapshots. Convert that whole-block estimate back + # to the manager's effective per-token constructor input. + block_tokens = max(1, int(self.config.paged_cache_block_size)) + block_payload = self.memory_monitor.estimate_paged_writer_block_memory( + block_tokens + ) + expected_kv_bytes_per_token = max( + 1, math.ceil(block_payload / block_tokens) ) else: expected_kv_bytes_per_token = 200_000 # PagedSSDCacheManager default diff --git a/omlx/server.py b/omlx/server.py index abafcfb0c..f8a966500 100644 --- a/omlx/server.py +++ b/omlx/server.py @@ -1466,6 +1466,43 @@ async def _ensure_tokenizer_for_system_probe( await engine.start() +async def _preflight_chat_with_eviction_budget( + engine: BaseEngine, + messages: list, + chat_kwargs: dict, + *, + request_id: str | None, +) -> None: + """Run route preflight and carry a real callback attempt to admission. + + The marker lives only in this HTTP request's kwargs, so concurrent routes + cannot consume each other's retry and cancellation leaves no keyed state. + """ + callback_attempted = await engine.preflight_chat( + messages, + request_id=request_id, + **chat_kwargs, + ) + if callback_attempted is True: + chat_kwargs["prefill_eviction_callback_attempted"] = True + + +async def _preflight_completion_with_eviction_budget( + engine: BaseEngine, + prompt: str, + generation_kwargs: dict[str, bool], + *, + request_id: str | None, +) -> None: + """Carry a completed route eviction attempt to this prompt's admission.""" + callback_attempted = await engine.preflight_completion( + prompt, + request_id=request_id, + ) + if callback_attempted is True: + generation_kwargs["prefill_eviction_callback_attempted"] = True + + def _unsupported_mid_system_policy() -> str: settings = _server_state.global_settings preserve_cache = True @@ -2960,7 +2997,7 @@ async def create_completion( request: CompletionRequest, http_request: FastAPIRequest, _: bool = Depends(verify_api_key), -): +) -> StreamingResponse: """Create a text completion.""" if _server_state.oq_manager and _server_state.oq_manager.is_quantizing: raise HTTPException( @@ -2992,9 +3029,17 @@ async def create_completion( # log line and the FastAPI handler trace correlate with whatever # the client is using on its side. upstream_request_id = http_request.headers.get("x-request-id") + completion_preflight_kwargs_by_prompt: list[dict[str, bool]] = [ + {} for _ in prompts + ] await _raise_if_llm_lease_abort_requested(lease) - for prompt in prompts: - await engine.preflight_completion(prompt, request_id=upstream_request_id) + for i, prompt in enumerate(prompts): + await _preflight_completion_with_eviction_budget( + engine, + prompt, + completion_preflight_kwargs_by_prompt[i], + request_id=upstream_request_id, + ) await _raise_if_llm_lease_abort_requested(lease) if request.stream: @@ -3013,6 +3058,12 @@ async def create_completion( prompt_token_ids=prompt_token_ids_by_prompt[0], resolved_model=resolved_model, response_id=response_id, + prefill_eviction_callback_attempted=( + completion_preflight_kwargs_by_prompt[0].get( + "prefill_eviction_callback_attempted", + False, + ) + ), ), http_request=http_request, keepalive_chunk=keepalive, @@ -3024,7 +3075,7 @@ async def create_completion( ) # Non-streaming response with keepalive during prefill - async def _build_completion(): + async def _build_completion() -> str: await _raise_if_llm_lease_abort_requested(lease) start_time = time.perf_counter() choices = [] @@ -3081,6 +3132,7 @@ async def _build_completion(): xtc_threshold=xtc_threshold, stop=request.stop, seed=request.seed, + **completion_preflight_kwargs_by_prompt[i], **gen_kwargs, ) if i == 0: @@ -3506,10 +3558,11 @@ async def create_chat_completion( # an incomplete chunked read. Running the check here lets # prefill_memory_exceeded_handler return a clean HTTP 400. await _raise_if_llm_lease_abort_requested(lease) - await engine.preflight_chat( + await _preflight_chat_with_eviction_budget( + engine, messages, + chat_kwargs, request_id=http_request.headers.get("x-request-id"), - **chat_kwargs, ) await _raise_if_llm_lease_abort_requested(lease) @@ -4097,6 +4150,7 @@ async def stream_completion( prompt_token_ids: list[int] | None = None, resolved_model: str | None = None, response_id: str | None = None, + prefill_eviction_callback_attempted: bool = False, ) -> AsyncIterator[str]: """Stream completion response.""" response_id = response_id or f"cmpl-{uuid.uuid4().hex[:8]}" @@ -4138,6 +4192,8 @@ async def stream_completion( thinking_budget = _resolve_thinking_budget(request, request.model) if thinking_budget is not None: gen_kwargs["thinking_budget"] = thinking_budget + if prefill_eviction_callback_attempted: + gen_kwargs["prefill_eviction_callback_attempted"] = True try: async for output in engine.stream_generate( prompt=prompt, @@ -5419,10 +5475,11 @@ async def create_anthropic_message( # Pre-flight prefill memory guard — must precede any StreamingResponse # return so PrefillMemoryExceededError can be mapped to HTTP 400. await _raise_if_llm_lease_abort_requested(lease) - await engine.preflight_chat( + await _preflight_chat_with_eviction_budget( + engine, messages, + chat_kwargs, request_id=http_request.headers.get("x-request-id"), - **chat_kwargs, ) await _raise_if_llm_lease_abort_requested(lease) @@ -5901,10 +5958,11 @@ async def create_response( # Pre-flight prefill memory guard — must precede any StreamingResponse # return so PrefillMemoryExceededError can be mapped to HTTP 400. await _raise_if_llm_lease_abort_requested(lease) - await engine.preflight_chat( + await _preflight_chat_with_eviction_budget( + engine, messages, + chat_kwargs, request_id=http_request.headers.get("x-request-id"), - **chat_kwargs, ) await _raise_if_llm_lease_abort_requested(lease) diff --git a/omlx/turboquant_kv.py b/omlx/turboquant_kv.py index 088443c80..d2b01eace 100644 --- a/omlx/turboquant_kv.py +++ b/omlx/turboquant_kv.py @@ -11,7 +11,9 @@ import logging import math -from typing import List, Optional +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any import mlx.core as mx from mlx_lm.models.cache import ( @@ -45,12 +47,508 @@ logger = logging.getLogger(__name__) __all__ = [ + "TURBOQUANT_CONVERSION_SLICE_TOKENS", + "TURBOQUANT_PREFILL_KEY_CHUNK_TOKENS", + "TURBOQUANT_PREFILL_QUERY_BLOCK_TOKENS", + "TurboQuantConversionStats", "TurboQuantKVCache", "BatchTurboQuantKVCache", + "convert_kv_cache_sliced", + "estimate_turboquant_conversion_peak_bytes", + "estimate_turboquant_prefill_attention_workspace_bytes", + "turboquant_mse_bytes_per_element", "turboquant_enabled", ] +TURBOQUANT_CONVERSION_SLICE_TOKENS = 8192 +_CONVERSION_WORKSPACE_ARRAYS_PER_SOURCE = 4 +TURBOQUANT_PREFILL_QUERY_BLOCK_TOKENS = 256 +TURBOQUANT_PREFILL_KEY_CHUNK_TOKENS = 16384 + + +@dataclass(frozen=True, slots=True) +class TurboQuantConversionStats: + """Observed shape and byte counts for one cache-list conversion.""" + + converted_layers: int + already_quantized_layers: int + skipped_dense_layers: int + slices: int + source_bytes: int + converted_bytes: int + + +def _turboquant_family_indices(cache_list: list[Any]) -> list[int]: + """Return full-attention layer indices used by the skip-last rule.""" + return [ + index + for index, cache_obj in enumerate(cache_list) + if isinstance(cache_obj, (KVCache, TurboQuantKVCache)) + ] + + +def _turboquant_target_indices( + cache_list: list[Any], *, skip_last: bool +) -> tuple[set[int], int | None]: + """Return conversion targets and the optional dense skip-last layer.""" + family_indices = _turboquant_family_indices(cache_list) + skipped_index = ( + family_indices[-1] if skip_last and len(family_indices) > 1 else None + ) + targets = set(family_indices) + if skipped_index is not None: + targets.remove(skipped_index) + return targets, skipped_index + + +def _turboquant_mse_bit_widths(bits: float) -> tuple[int, int]: + """Return the integer key/value widths used by the MSE codec.""" + validated_bits = float(_validate_bits(bits)) + if math.isclose(validated_bits, round(validated_bits), abs_tol=1e-6): + width = int(round(validated_bits)) + return width, width + return int(math.floor(validated_bits)), int(math.ceil(validated_bits)) + + +def _quantized_mse_vector_bytes(head_dim: int, bits: int) -> int: + """Return one packed MSE vector's norm and uint32 index bytes.""" + packed_words = (head_dim * bits + 31) // 32 + return mx.float16.size + packed_words * mx.uint32.size + + +def turboquant_mse_bytes_per_element(head_dim: int, bits: float) -> float: + """Return the average packed MSE K/V resident width per element.""" + if not isinstance(head_dim, int) or isinstance(head_dim, bool) or head_dim <= 0: + raise ValueError("head_dim must be a positive integer") + key_bits, value_bits = _turboquant_mse_bit_widths(bits) + key_vector_bytes = _quantized_mse_vector_bytes(head_dim, key_bits) + value_vector_bytes = _quantized_mse_vector_bytes(head_dim, value_bits) + return (key_vector_bytes + value_vector_bytes) / (2 * head_dim) + +def estimate_turboquant_prefill_attention_workspace_bytes( + *, + query_tokens: int, + kv_len: int, + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + bits: float, + compute_dtype_size: float = 2, + causal: bool = True, +) -> int: + """Bound the first chunked Q8-style TurboQuant prefill attention call. + + The long-prefill route retains all completed query-block outputs while it + evaluates one 256-query by 16384-key block at a time. This structural + bound prices those retained outputs, the active score/softmax tensors, + unpack/cast/codebook tensors for K and V, packed state slices, and the + caller-owned query input. It does not rely on allocator fusion or on a + prior TurboQuant transient sample. + """ + dimensions = ( + query_tokens, + kv_len, + num_query_heads, + num_kv_heads, + head_dim, + ) + if any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in dimensions + ): + return 0 + if ( + not isinstance(bits, (int, float)) + or isinstance(bits, bool) + or not math.isfinite(float(bits)) + or float(bits) <= 0 + ): + return 0 + if ( + not isinstance(compute_dtype_size, (int, float)) + or isinstance(compute_dtype_size, bool) + or not math.isfinite(float(compute_dtype_size)) + or float(compute_dtype_size) <= 0 + ): + return 0 + + q_block = min(query_tokens, TURBOQUANT_PREFILL_QUERY_BLOCK_TOKENS) + k_block = min(kv_len, TURBOQUANT_PREFILL_KEY_CHUNK_TOKENS) + key_bits, value_bits = _turboquant_mse_bit_widths(float(bits)) + key_words = (head_dim * key_bits + 31) // 32 + value_words = (head_dim * value_bits + 31) // 32 + compute_bytes = float(compute_dtype_size) + + # Full-query buffers: caller input, scaled queries, final compute cast, + # retained float32 blocks, and the concatenated float32 result. + total = (3 * compute_bytes + 8) * num_query_heads * query_tokens * head_dim + # Nine active float32 query/value/accumulator stages. + total += 36 * num_query_heads * q_block * head_dim + # Dots, scaled scores, softmax subtraction, and weights. + total += 16 * num_query_heads * q_block * k_block + # K and V uint32 unpack, int32 cast, and float32 codebook-take tensors. + total += 24 * num_kv_heads * k_block * head_dim + # K/V norm casts plus packed state-slice materialization. + total += 8 * num_kv_heads * k_block + total += ( + num_kv_heads + * k_block + * (4 * (key_words + value_words) + 4) + ) + # Per-query max/denominator and online-softmax state. + total += 48 * num_query_heads * q_block + + if causal: + # One additional masked score result, the causal bool tile, and its + # query/key index vectors. + total += 4 * num_query_heads * q_block * k_block + total += q_block * k_block + total += 8 * (q_block + k_block) + + return int(math.ceil(total)) + + +def _quantized_state_shape_bytes( + *, + batch_size: int, + num_heads: int, + num_tokens: int, + head_dim: int, + bits: int, +) -> int: + """Return packed MSE-state bytes for one key or value tensor.""" + vectors = batch_size * num_heads * num_tokens + return vectors * _quantized_mse_vector_bytes(head_dim, bits) + + +def _layer_conversion_peak_bytes( + keys: mx.array, + values: mx.array, + *, + bits: float, + slice_tokens: int, +) -> int: + """Bound incremental bytes while converting one dense cache layer. + + After the first slice establishes the state type, the destination reserves + final capacity before later slices run. The peak therefore holds the final + state, one quantized slice, and bounded quantization workspace. The + workspace charges four fp32-sized arrays per key/value source tensor. + """ + batch_size = int(keys.shape[0]) + num_heads = int(keys.shape[1]) + num_tokens = int(keys.shape[2]) + key_dim = int(keys.shape[3]) + value_dim = int(values.shape[3]) + key_bits, value_bits = _turboquant_mse_bit_widths(bits) + bounded_tokens = min(num_tokens, slice_tokens) + + def _state_bytes(tokens: int) -> int: + return _quantized_state_shape_bytes( + batch_size=batch_size, + num_heads=num_heads, + num_tokens=tokens, + head_dim=key_dim, + bits=key_bits, + ) + _quantized_state_shape_bytes( + batch_size=batch_size, + num_heads=num_heads, + num_tokens=tokens, + head_dim=value_dim, + bits=value_bits, + ) + + final_state = _state_bytes(num_tokens) + slice_state = _state_bytes(bounded_tokens) + source_elements = batch_size * num_heads * bounded_tokens * (key_dim + value_dim) + workspace = ( + source_elements * mx.float32.size * _CONVERSION_WORKSPACE_ARRAYS_PER_SOURCE + ) + codec_tables = 2 * (key_dim * key_dim + value_dim * value_dim) * mx.float32.size + return final_state + slice_state + workspace + codec_tables + + +def _validate_dense_kv_state(keys: mx.array, values: mx.array) -> int: + """Validate the dense source shape and return its logical token count.""" + if keys.ndim != 4 or values.ndim != 4: + raise ValueError("TurboQuant conversion requires 4-D K/V state") + if keys.shape[:3] != values.shape[:3]: + raise ValueError("TurboQuant conversion requires matching K/V batch shapes") + return int(keys.shape[2]) + + +def _append_turboquant_slice( + cache_obj: TurboQuantKVCache, + keys: mx.array, + values: mx.array, + *, + start: int, + end: int, + stream: Any | None, +) -> None: + """Quantize and materialize one bounded token slice. + + ``TurboQuantKVCache.state`` caches a lazy prefix slice. The converter + ignores the returned proxies, so clear its private prefix immediately + before the dependency constructs each slice update and again after eval. + No converter-owned proxy then survives to the next donation decision. + """ + cache_obj._cached_state = None + cache_obj._cached_state_offset = -1 + try: + if stream is None: + cache_obj.update_and_fetch( + keys[:, :, start:end, :], + values[:, :, start:end, :], + ) + mx.eval(cache_obj.keys, cache_obj.values) + return + with mx.stream(stream): + cache_obj.update_and_fetch( + keys[:, :, start:end, :], + values[:, :, start:end, :], + ) + mx.eval(cache_obj.keys, cache_obj.values) + finally: + cache_obj._cached_state = None + cache_obj._cached_state_offset = -1 + + +def _validate_converted_layer( + cache_obj: TurboQuantKVCache, + *, + expected_tokens: int, + expected_bits: float, +) -> None: + """Reject an incomplete or incompatible converted cache candidate.""" + if cache_obj.offset != expected_tokens: + raise RuntimeError( + f"TurboQuant candidate converted {cache_obj.offset} " + f"of {expected_tokens} tokens" + ) + if not math.isclose(cache_obj.bits, expected_bits, abs_tol=1e-6): + raise RuntimeError( + f"TurboQuant candidate uses {cache_obj.bits} bits, " + f"expected {expected_bits}" + ) + key_state, value_state = cache_obj.state + for label, state in (("key", key_state), ("value", value_state)): + if not isinstance(state, TurboQuantMSEState): + raise RuntimeError( + f"TurboQuant {label} candidate has unsupported " + f"{type(state).__name__} state" + ) + if state.norms.dtype != mx.float16 or state.indices.dtype != mx.uint32: + raise RuntimeError(f"TurboQuant {label} candidate has invalid state dtypes") + if int(state.norms.shape[2]) != expected_tokens: + raise RuntimeError( + f"TurboQuant {label} candidate has incomplete logical state" + ) + + +def estimate_turboquant_conversion_peak_bytes( + cache_list: list[Any], + *, + bits: float, + skip_last: bool, + slice_tokens: int = TURBOQUANT_CONVERSION_SLICE_TOKENS, +) -> int: + """Return a conservative incremental peak for layer-wise conversion.""" + if slice_tokens <= 0: + raise ValueError("slice_tokens must be positive") + _validate_bits(bits) + target_indices, _ = _turboquant_target_indices(cache_list, skip_last=skip_last) + peak = 0 + for index in target_indices: + cache_obj = cache_list[index] + if isinstance(cache_obj, TurboQuantKVCache): + if not math.isclose(cache_obj.bits, bits, abs_tol=1e-6): + raise ValueError( + f"TurboQuant layer {index} uses {cache_obj.bits} bits, " + f"expected {bits}" + ) + continue + if not isinstance(cache_obj, KVCache) or cache_obj.empty(): + continue + keys, values = cache_obj.state + _validate_dense_kv_state(keys, values) + peak = max( + peak, + _layer_conversion_peak_bytes( + keys, + values, + bits=bits, + slice_tokens=slice_tokens, + ), + ) + return peak + + +def convert_kv_cache_sliced( + cache_list: list[Any], + *, + bits: float, + skip_last: bool, + slice_tokens: int = TURBOQUANT_CONVERSION_SLICE_TOKENS, + stream: Any | None = None, + check_cancelled: Callable[[], None] | None = None, +) -> TurboQuantConversionStats: + """Convert dense full-attention layers in bounded token slices. + + Each destination layer is fully evaluated before replacing its dense + source. Prior converted layers are released before the next layer starts. + If a callback raises, the current dense layer remains in place; callers + must discard the whole cache because earlier layers may already be swapped. + """ + from .utils.metal_sync import _sync_and_clear_cache + + if slice_tokens <= 0: + raise ValueError("slice_tokens must be positive") + validated_bits = float(_validate_bits(bits)) + target_indices, skipped_index = _turboquant_target_indices( + cache_list, skip_last=skip_last + ) + converted_layers = 0 + already_quantized_layers = 0 + skipped_dense_layers = 0 + slices = 0 + source_bytes = 0 + converted_bytes = 0 + + for index in _turboquant_family_indices(cache_list): + cache_obj: Any | None = cache_list[index] + keys: mx.array | None = None + values: mx.array | None = None + turbo_cache: TurboQuantKVCache | None = None + converted_current_layer = False + try: + if index == skipped_index: + if isinstance(cache_obj, KVCache): + skipped_dense_layers += 1 + continue + if index not in target_indices: + continue + if isinstance(cache_obj, TurboQuantKVCache): + if not math.isclose(cache_obj.bits, validated_bits, abs_tol=1e-6): + raise ValueError( + f"TurboQuant layer {index} uses {cache_obj.bits} bits, " + f"expected {validated_bits}" + ) + already_quantized_layers += 1 + converted_bytes += int(cache_obj.nbytes) + continue + if not isinstance(cache_obj, KVCache): + continue + if check_cancelled is not None: + check_cancelled() + + turbo_cache = TurboQuantKVCache(bits=validated_bits) + if cache_obj.empty(): + cache_list[index] = turbo_cache + converted_layers += 1 + continue + + keys, values = cache_obj.state + num_tokens = _validate_dense_kv_state(keys, values) + source_bytes += int(keys.nbytes + values.nbytes) + + first_end = min(slice_tokens, num_tokens) + _append_turboquant_slice( + turbo_cache, + keys, + values, + start=0, + end=first_end, + stream=stream, + ) + slices += 1 + _sync_and_clear_cache(stream) + if check_cancelled is not None: + check_cancelled() + + if first_end < num_tokens: + if stream is None: + turbo_cache.keys = _reserve_state_capacity( + turbo_cache.keys, + turbo_cache.offset, + num_tokens, + num_tokens, + ) + turbo_cache.values = _reserve_state_capacity( + turbo_cache.values, + turbo_cache.offset, + num_tokens, + num_tokens, + ) + turbo_cache._cached_state = None + turbo_cache._cached_state_offset = -1 + mx.eval(turbo_cache.keys, turbo_cache.values) + else: + with mx.stream(stream): + turbo_cache.keys = _reserve_state_capacity( + turbo_cache.keys, + turbo_cache.offset, + num_tokens, + num_tokens, + ) + turbo_cache.values = _reserve_state_capacity( + turbo_cache.values, + turbo_cache.offset, + num_tokens, + num_tokens, + ) + turbo_cache._cached_state = None + turbo_cache._cached_state_offset = -1 + mx.eval(turbo_cache.keys, turbo_cache.values) + _sync_and_clear_cache(stream) + if check_cancelled is not None: + check_cancelled() + + for start in range(first_end, num_tokens, slice_tokens): + if check_cancelled is not None: + check_cancelled() + end = min(start + slice_tokens, num_tokens) + _append_turboquant_slice( + turbo_cache, + keys, + values, + start=start, + end=end, + stream=stream, + ) + slices += 1 + _sync_and_clear_cache(stream) + if check_cancelled is not None: + check_cancelled() + + _validate_converted_layer( + turbo_cache, + expected_tokens=num_tokens, + expected_bits=validated_bits, + ) + cache_list[index] = turbo_cache + converted_layers += 1 + converted_bytes += int(turbo_cache.nbytes) + converted_current_layer = True + finally: + cache_obj = None + keys = None + values = None + turbo_cache = None + if converted_current_layer: + _sync_and_clear_cache(stream) + + return TurboQuantConversionStats( + converted_layers=converted_layers, + already_quantized_layers=already_quantized_layers, + skipped_dense_layers=skipped_dense_layers, + slices=slices, + source_bytes=source_bytes, + converted_bytes=converted_bytes, + ) + + # --------------------------------------------------------------------------- # Codec rebuild for SSD cache reconstruction # --------------------------------------------------------------------------- @@ -234,7 +732,9 @@ class BatchTurboQuantKVCache(TurboQuantKVCache): overrides make_mask for per-request left_padding support. """ - def __init__(self, left_padding: List[int], bits: float = 4.0, seed: int = 0): + def __init__( + self, left_padding: list[int], bits: float = 4.0, seed: int = 0 + ) -> None: super().__init__(bits=bits, seed=seed) self.group_size = 0 self.left_padding = mx.array(left_padding) @@ -304,8 +804,8 @@ def make_mask( self, N: int, return_array: bool = False, - window_size: Optional[int] = None, - ): + window_size: int | None = None, + ) -> str | mx.array | None: offset = self.offset if isinstance(offset, int): return create_attention_mask(N, offset, return_array, window_size) @@ -481,7 +981,7 @@ def extract(self, idx: int) -> TurboQuantKVCache: return tq @classmethod - def merge(cls, caches: List[TurboQuantKVCache]) -> "BatchTurboQuantKVCache": + def merge(cls, caches: list[TurboQuantKVCache]) -> BatchTurboQuantKVCache: for cache in caches: if not isinstance(cache, TurboQuantKVCache): raise TypeError( diff --git a/omlx/utils/metal_sync.py b/omlx/utils/metal_sync.py index 1c129b764..bba0fc70f 100644 --- a/omlx/utils/metal_sync.py +++ b/omlx/utils/metal_sync.py @@ -16,10 +16,16 @@ """ import threading +import weakref +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any import mlx.core as mx from mlx_lm.generate import generation_stream +from ..exceptions import TurboQuantProcessExclusiveError + # Module-level alias so callers can fall back to mlx-lm's default stream # when no per-engine stream is provided. _default_generation_stream = generation_stream @@ -33,7 +39,218 @@ _mx_buffer_access_lock = threading.RLock() -def _sync_and_clear_cache(stream=None): +class _ConversionCoordinator: + """Own process-exclusive mid-prefill conversion and peak reservations. + + A mid-prefill engine must be the process's sole ``EngineCore`` and claims + the capability only after the global MLX executor has drained. While that + engine lives, new engines and independent Metal workers fail closed. This + keeps cache-clearing conversion away from streams it cannot drain. + """ + + def __init__(self) -> None: + self._condition = threading.Condition(threading.Lock()) + self._registered_engines: weakref.WeakSet[Any] = weakref.WeakSet() + self._exclusive_owner: weakref.ReferenceType[Any] | None = None + self._background_metal_operations = 0 + self._waiting_conversions = 0 + self._conversion_owner: object | None = None + self._reservation_owner: object | None = None + self._outstanding_bytes = 0 + + def _exclusive_owner_unlocked(self) -> Any | None: + owner_ref = self._exclusive_owner + if owner_ref is None: + return None + owner = owner_ref() + if owner is None: + self._exclusive_owner = None + return owner + + def register_engine(self, owner: Any) -> None: + """Register an EngineCore before it creates a Metal executor.""" + with self._condition: + exclusive_owner = self._exclusive_owner_unlocked() + if exclusive_owner is not None and exclusive_owner is not owner: + raise TurboQuantProcessExclusiveError( + "TurboQuant mid-prefill requires process-exclusive Metal " + "access; unload the mid-prefill model before loading " + "another engine" + ) + self._registered_engines.add(owner) + + def unregister_engine(self, owner: Any) -> None: + """Release an engine after its owning-thread stream has drained.""" + with self._condition: + if self._exclusive_owner_unlocked() is owner: + self._exclusive_owner = None + self._registered_engines.discard(owner) + self._condition.notify_all() + + def claim_process_exclusive(self, owner: Any) -> None: + """Claim the mid-prefill capability for the process's sole engine.""" + with self._condition: + if owner not in self._registered_engines: + raise RuntimeError("mid-prefill owner is not a registered engine") + exclusive_owner = self._exclusive_owner_unlocked() + if exclusive_owner is not None and exclusive_owner is not owner: + raise TurboQuantProcessExclusiveError( + "another engine already owns process-exclusive Metal access" + ) + other_engines = [ + engine for engine in self._registered_engines if engine is not owner + ] + if other_engines: + raise TurboQuantProcessExclusiveError( + "TurboQuant mid-prefill requires process-exclusive Metal " + "access; unload all other engines before enabling it" + ) + if self._background_metal_operations: + raise TurboQuantProcessExclusiveError( + "TurboQuant mid-prefill cannot start while an independent " + "Metal operation is active" + ) + self._exclusive_owner = weakref.ref(owner) + + def process_exclusive(self, owner: Any | None) -> bool: + """Return whether a registered ``owner`` holds the process capability.""" + if owner is None: + return False + with self._condition: + if owner not in self._registered_engines: + return False + return self._exclusive_owner_unlocked() is owner + + def assert_background_metal_allowed(self) -> None: + """Reject a global-executor task while a mid-prefill engine is live.""" + with self._condition: + if self._exclusive_owner_unlocked() is not None: + raise TurboQuantProcessExclusiveError( + "Independent Metal work is unavailable while a " + "TurboQuant mid-prefill engine owns the process" + ) + + @contextmanager + def background_metal_operation(self) -> Iterator[None]: + """Track a non-executor Metal worker such as oQ quantization.""" + with self._condition: + if self._exclusive_owner_unlocked() is not None: + raise TurboQuantProcessExclusiveError( + "Independent Metal work is unavailable while a " + "TurboQuant mid-prefill engine owns the process" + ) + self._background_metal_operations += 1 + try: + yield + finally: + with self._condition: + self._background_metal_operations -= 1 + self._condition.notify_all() + + @contextmanager + def conversion(self, *, process_owner: Any | None = None) -> Iterator[object]: + """Serialize a bounded conversion owned by the exclusive engine.""" + owner = object() + acquired = False + with self._condition: + if ( + process_owner is None + or process_owner not in self._registered_engines + or self._exclusive_owner_unlocked() is not process_owner + ): + raise RuntimeError( + "TurboQuant mid-prefill conversion lacks process-exclusive " + "Metal ownership" + ) + self._waiting_conversions += 1 + try: + while self._conversion_owner is not None: + self._condition.wait() + self._conversion_owner = owner + acquired = True + finally: + self._waiting_conversions -= 1 + if not acquired: + self._condition.notify_all() + try: + yield owner + finally: + with self._condition: + if self._reservation_owner is owner: + self._reservation_owner = None + self._outstanding_bytes = 0 + if self._conversion_owner is not owner: + raise RuntimeError("TurboQuant conversion gate ownership was lost") + self._conversion_owner = None + self._condition.notify_all() + + def try_reserve( + self, + owner: object, + *, + current_bytes: int, + peak_bytes: int, + limit_bytes: int, + ) -> tuple[bool, int]: + """Atomically check headroom and publish an accepted conversion peak.""" + if current_bytes < 0 or peak_bytes < 0 or limit_bytes < 0: + raise ValueError("conversion memory values must be non-negative") + with self._condition: + if self._conversion_owner is not owner: + raise RuntimeError( + "conversion reservation requires exclusive ownership" + ) + if ( + self._reservation_owner is not None + and self._reservation_owner is not owner + ): + raise RuntimeError("another conversion reservation is active") + prior_outstanding = ( + 0 if self._reservation_owner is owner else self._outstanding_bytes + ) + estimated_bytes = current_bytes + prior_outstanding + peak_bytes + if limit_bytes > 0 and estimated_bytes > limit_bytes: + return False, estimated_bytes + self._reservation_owner = owner + self._outstanding_bytes = peak_bytes + self._condition.notify_all() + return True, estimated_bytes + + def release_reservation(self, owner: object) -> None: + """Release the holder's peak before its post-conversion sample.""" + with self._condition: + if self._conversion_owner is not owner: + raise RuntimeError("conversion reservation owner is not active") + if self._reservation_owner is None: + return + if self._reservation_owner is not owner: + raise RuntimeError("conversion reservation ownership was lost") + self._reservation_owner = None + self._outstanding_bytes = 0 + self._condition.notify_all() + + def outstanding_bytes(self, *, exclude_owner: object | None = None) -> int: + """Return bytes reserved by another conversion holder.""" + with self._condition: + if self._reservation_owner is exclude_owner: + return 0 + return self._outstanding_bytes + + def snapshot(self) -> tuple[int, int, bool, int]: + """Return background, waiting, active-writer, and reservation state.""" + with self._condition: + return ( + self._background_metal_operations, + self._waiting_conversions, + self._conversion_owner is not None, + self._outstanding_bytes, + ) + + +_conversion_coordinator = _ConversionCoordinator() + + +def _sync_and_clear_cache(stream: Any | None = None) -> None: """Synchronize in-flight GPU work before clearing the Metal buffer cache. Without synchronization, mx.clear_cache() can release Metal buffers that diff --git a/scripts/validate_turboquant_mid_prefill.py b/scripts/validate_turboquant_mid_prefill.py new file mode 100755 index 000000000..c6cc74402 --- /dev/null +++ b/scripts/validate_turboquant_mid_prefill.py @@ -0,0 +1,2074 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Reproducible, parent-supervised TurboQuant mid-prefill validation.""" + +from __future__ import annotations + +import argparse +import asyncio +import fcntl +import gc +import hashlib +import importlib.metadata +import json +import logging +import math +import os +import platform +import re +import statistics +import subprocess +import sys +import tempfile +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import asdict, dataclass +from difflib import SequenceMatcher +from functools import partial +from pathlib import Path +from types import SimpleNamespace +from typing import Any, TypeVar + +logger = logging.getLogger(__name__) + +DATASET_REPO = "openai/mrcr" +DATASET_REVISION = "f4c69fae7cf81f7ca26b9fee34b392a50f6b8a1d" +DATASET_FILE = "2needle/2needle_0.parquet" +DEFAULT_MATRIX_ROWS: tuple[int, ...] = (104, 109, 136, 301, 311, 328) +PRIMARY_PERFORMANCE_ROWS: Mapping[int, str] = { + 104: "primary_32k", + 311: "primary_64k", +} +ORGANIC_ROW_INDEX = 35 +ORGANIC_PROMPT_TOKENS = 131_071 +FIXED_REPLAY_TOKENS = 256 +ORGANIC_PREFILL_ABORT_MARGIN = 0.95 +ORGANIC_PREFILL_MIN_CHUNK_TOKENS = 32 +SCHEMA_VERSION = 1 +GIB = 1024**3 +TELEMETRY_EXIT_GRACE_SECONDS = 1.0 +_FORCE_ENV_PREFIX = "OMLX_FORCE_" +_MEMORY_TOTAL_RE = re.compile(r"The system has\s+(\d+)\s+\(") +_MEMORY_FREE_RE = re.compile(r"System-wide memory free percentage:\s*(\d+(?:\.\d+)?)%") +_T = TypeVar("_T") + + +class ValidationError(RuntimeError): + """Raised when the validation contract cannot be established.""" + + +class TelemetryError(ValidationError): + """Raised when mandatory supervisor telemetry is invalid.""" + + +@dataclass(frozen=True, slots=True) +class ValidationMode: + """One fixed matrix validation mode.""" + + name: str + conversion: str + bits: float | None + + +@dataclass(frozen=True, slots=True) +class MemoryPressureReading: + """Parsed kernel memory-pressure reading.""" + + total_bytes: int + free_percent: float + headroom_bytes: int + + +@dataclass(slots=True) +class SupervisedOutcome: + """One child result plus mandatory parent-side safety observations.""" + + result: dict[str, Any] | None + telemetry: dict[str, Any] + returncode: int + error: str | None + + +@dataclass(slots=True) +class OrganicPrefillPause: + """First organic prefill attempt paused for the production LRU callback.""" + + baseline: dict[str, int] + context: Any + eviction_request: Any + external_seconds: float + total_started: float + + +def matrix_modes() -> tuple[ValidationMode, ...]: + """Return the complete, ordered five-mode validation matrix.""" + return ( + ValidationMode("dense", "dense", None), + ValidationMode("ordinary-q8", "ordinary", 8.0), + ValidationMode("ordinary-q4", "ordinary", 4.0), + ValidationMode("mid-q8", "mid", 8.0), + ValidationMode("mid-q4", "mid", 4.0), + ) + + +def pinned_dataset_config() -> dict[str, Any]: + """Return immutable MRCR retrieval provenance as JSON-compatible data.""" + return { + "repo_id": DATASET_REPO, + "revision": DATASET_REVISION, + "filename": DATASET_FILE, + "matrix_row_indices": list(DEFAULT_MATRIX_ROWS), + "organic_row_index": ORGANIC_ROW_INDEX, + } + + +def select_rows( + rows: Sequence[Mapping[str, Any]], row_indices: Sequence[int] +) -> list[dict[str, Any]]: + """Select rows by stable positional index and annotate each selected row.""" + selected: list[dict[str, Any]] = [] + seen: set[int] = set() + for raw_index in row_indices: + if isinstance(raw_index, bool) or not isinstance(raw_index, int): + raise ValueError("row indices must be integers") + if raw_index < 0 or raw_index >= len(rows): + raise IndexError(f"row index {raw_index} is outside 0..{len(rows) - 1}") + if raw_index in seen: + raise ValueError(f"duplicate row index {raw_index}") + seen.add(raw_index) + row = dict(rows[raw_index]) + row["row_index"] = raw_index + selected.append(row) + return selected + + +def official_retrieval_score(response: str, answer: str, prefix: str) -> float: + """Apply MRCR's official prefix gate and SequenceMatcher ratio.""" + if not response.startswith(prefix): + return 0.0 + sampled = response.removeprefix(prefix) + expected = answer.removeprefix(prefix) + return float(SequenceMatcher(None, sampled, expected).ratio()) + + +def parse_memory_pressure_output(output: str) -> MemoryPressureReading: + """Parse ``memory_pressure -Q`` output or fail closed.""" + total_match = _MEMORY_TOTAL_RE.search(output) + free_match = _MEMORY_FREE_RE.search(output) + if total_match is None or free_match is None: + raise TelemetryError("memory_pressure output omitted required fields") + total_bytes = int(total_match.group(1)) + free_percent = float(free_match.group(1)) + if total_bytes <= 0: + raise TelemetryError("memory_pressure reported a non-positive total") + if not math.isfinite(free_percent) or not 0.0 <= free_percent <= 100.0: + raise TelemetryError("memory_pressure reported an invalid free percentage") + headroom_bytes = int(total_bytes * free_percent / 100.0) + return MemoryPressureReading(total_bytes, free_percent, headroom_bytes) + + +def safety_violation( + child_phys_bytes: int, + host_headroom_bytes: int, + *, + child_limit_bytes: int, + host_minimum_bytes: int, +) -> str | None: + """Return the fail-closed threshold violation, if any.""" + values = ( + child_phys_bytes, + host_headroom_bytes, + child_limit_bytes, + host_minimum_bytes, + ) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + return "invalid safety telemetry type" + if child_phys_bytes <= 0 or host_headroom_bytes < 0: + return "invalid safety telemetry value" + if child_limit_bytes <= 0 or host_minimum_bytes < 0: + return "invalid safety threshold" + if child_phys_bytes >= child_limit_bytes: + return "child physical footprint reached its safety limit" + if host_headroom_bytes < host_minimum_bytes: + return "host headroom fell below its safety minimum" + return None + + +def canonical_json_bytes(value: Any) -> bytes: + """Serialize JSON data deterministically for hashing and atomic output.""" + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + + +def stable_digest(value: Any) -> str: + """Return a SHA-256 digest of canonical JSON data.""" + return hashlib.sha256(canonical_json_bytes(value)).hexdigest() + + +def hash_text(value: str) -> str: + """Return the exact UTF-8 SHA-256 of a string.""" + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + +def hash_token_ids(token_ids: Sequence[int]) -> str: + """Return a stable SHA-256 over an exact token-id sequence.""" + return stable_digest(list(token_ids)) + + +def _enable_verified_no_cache(file_descriptor: int) -> None: + """Enable Darwin F_NOCACHE and require the setter's verified success result.""" + if sys.platform != "darwin": + raise ValidationError("F_NOCACHE validation requires Darwin") + command = getattr(fcntl, "F_NOCACHE", None) + if isinstance(command, bool) or not isinstance(command, int): + raise ValidationError("Python does not expose Darwin F_NOCACHE") + try: + result = fcntl.fcntl(file_descriptor, command, 1) + except OSError as exc: + raise ValidationError( + f"F_NOCACHE failed for source tensor read: {exc}" + ) from exc + if result != 0: + raise ValidationError(f"F_NOCACHE setter returned unverified result {result!r}") + + +def hash_file(path: Path, chunk_bytes: int = 8 * 1024 * 1024) -> str: + """Hash one file through a verified no-cache descriptor.""" + digest = hashlib.sha256() + with path.open("rb", buffering=0) as handle: + _enable_verified_no_cache(handle.fileno()) + while True: + chunk = handle.read(chunk_bytes) + if not chunk: + break + digest.update(chunk) + return digest.hexdigest() + + +def _relative_manifest_path(path: Path, *, root: Path) -> str: + """Return one manifest path relative to its model root.""" + return path.relative_to(root).as_posix() + + +def build_model_manifest(model_path: Path) -> list[dict[str, Any]]: + """Hash every source safetensors file below a local model directory.""" + root = model_path.expanduser().resolve() + if not root.is_dir(): + raise ValidationError(f"model path is not a directory: {root}") + files = sorted( + (path for path in root.rglob("*.safetensors") if path.is_file()), + key=partial(_relative_manifest_path, root=root), + ) + if not files: + raise ValidationError(f"no source .safetensors files found under {root}") + return [ + { + "path": path.relative_to(root).as_posix(), + "size_bytes": path.stat().st_size, + "sha256": hash_file(path), + } + for path in files + ] + + +def manifests_equal( + before: Sequence[Mapping[str, Any]], after: Sequence[Mapping[str, Any]] +) -> bool: + """Compare source tensor manifests independent of mapping key order.""" + return canonical_json_bytes(list(before)) == canonical_json_bytes(list(after)) + + +def atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None: + """Atomically replace a JSON result and fsync both file and directory.""" + destination = path.expanduser().resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "wb") as handle: + handle.write(canonical_json_bytes(dict(payload))) + handle.write(b"\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, destination) + directory_fd = os.open(destination.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if temporary.exists(): + temporary.unlink() + + +def shape_validation_result( + *, + kind: str, + provenance: Mapping[str, Any], + config: Mapping[str, Any], + before_manifest: Sequence[Mapping[str, Any]], + after_manifest: Sequence[Mapping[str, Any]], + results: Sequence[Mapping[str, Any]], + error: str | None, +) -> dict[str, Any]: + """Build the stable top-level result shape used by both public modes.""" + unchanged = manifests_equal(before_manifest, after_manifest) + effective_error = error + if not unchanged: + mutation_error = "source .safetensors manifest changed during validation" + effective_error = ( + mutation_error + if effective_error is None + else f"{effective_error}; {mutation_error}" + ) + return { + "schema_version": SCHEMA_VERSION, + "status": "ok" if effective_error is None else "failed", + "kind": kind, + "provenance": dict(provenance), + "config": dict(config), + "source_tensors": { + "before": list(before_manifest), + "after": list(after_manifest), + "unchanged": unchanged, + }, + "results": [dict(result) for result in results], + "error": effective_error, + } + + +def shape_validation_checkpoint( + *, + kind: str, + provenance: Mapping[str, Any], + config: Mapping[str, Any], + before_manifest: Sequence[Mapping[str, Any]], + results: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Build a restart-safe partial result without claiming final tensor proof.""" + return { + "schema_version": SCHEMA_VERSION, + "status": "running", + "kind": kind, + "provenance": dict(provenance), + "config": dict(config), + "source_tensors": { + "before": list(before_manifest), + "after": None, + "unchanged": None, + }, + "results": [dict(result) for result in results], + "error": None, + } + + +def child_environment(base_environment: Mapping[str, str]) -> dict[str, str]: + """Build a child environment with every forced diagnostic hook removed.""" + return { + key: value + for key, value in base_environment.items() + if not key.startswith(_FORCE_ENV_PREFIX) + } + + +def build_child_command( + script_path: Path, spec_path: Path, result_path: Path +) -> list[str]: + """Build the private child invocation without any trigger controls.""" + return [ + sys.executable, + str(script_path.expanduser().resolve()), + "__child", + "--spec", + str(spec_path.expanduser().resolve()), + "--result", + str(result_path.expanduser().resolve()), + ] + + +def matrix_shared_metadata( + *, + row_index: int, + chunk_size: int, + max_prompt_tokens: int, + greedy_token_limit: int, +) -> dict[str, Any]: + """Build metadata that must be byte-identical across a sample's modes.""" + return { + "dataset": { + "repo_id": DATASET_REPO, + "revision": DATASET_REVISION, + "filename": DATASET_FILE, + "row_index": row_index, + }, + "quality_cell": True, + "performance_cell": PRIMARY_PERFORMANCE_ROWS.get(row_index), + "thinking_enabled": False, + "chunk_size": chunk_size, + "max_prompt_tokens": max_prompt_tokens, + "greedy_token_limit": greedy_token_limit, + "teacher_forced_replay_token_limit": FIXED_REPLAY_TOKENS, + } + + +def validate_matrix_results( + results: Sequence[Mapping[str, Any]], + expected_rows: Sequence[int] | None = None, +) -> None: + """Reject incomplete modes, rows, or cross-mode metadata/input drift.""" + expected_modes = {mode.name for mode in matrix_modes()} + by_row: dict[int, list[Mapping[str, Any]]] = {} + for result in results: + metadata = result.get("metadata") + if not isinstance(metadata, Mapping): + raise ValidationError("matrix child result omitted metadata") + dataset = metadata.get("dataset") + if not isinstance(dataset, Mapping) or not isinstance( + dataset.get("row_index"), int + ): + raise ValidationError("matrix child result omitted its row index") + by_row.setdefault(int(dataset["row_index"]), []).append(result) + if expected_rows is not None and set(by_row) != set(expected_rows): + raise ValidationError( + f"matrix rows {sorted(by_row)} do not match {sorted(expected_rows)}" + ) + for row_index, group in by_row.items(): + modes = {str(item.get("mode")) for item in group} + if len(group) != len(expected_modes) or modes != expected_modes: + raise ValidationError( + f"row {row_index} has modes {sorted(modes)}, expected {sorted(expected_modes)}" + ) + metadata_digests = {stable_digest(item.get("metadata")) for item in group} + if len(metadata_digests) != 1: + raise ValidationError(f"row {row_index} metadata drifted across modes") + input_rows: list[tuple[Any, ...]] = [] + for item in group: + metrics = item.get("metrics") + if not isinstance(metrics, Mapping): + raise ValidationError(f"row {row_index} omitted input metrics") + values = ( + metrics.get("dataset_prompt_sha256"), + metrics.get("rendered_prompt_sha256"), + metrics.get("answer_sha256"), + metrics.get("prompt_token_ids_sha256"), + metrics.get("answer_token_ids_sha256"), + metrics.get("prompt_token_count"), + metrics.get("answer_token_count"), + ) + if any(value is None for value in values): + raise ValidationError(f"row {row_index} omitted exact input identity") + input_rows.append(values) + if len(set(input_rows)) != 1: + raise ValidationError(f"row {row_index} inputs drifted across modes") + + +def _run_text_command(command: Sequence[str], cwd: Path | None = None) -> str | None: + """Run a short provenance command and return stripped stdout on success.""" + try: + completed = subprocess.run( + list(command), + cwd=cwd, + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return completed.stdout.strip() + + +def _parse_optional_int(value: str | None) -> int | None: + """Parse an optional integer provenance value.""" + if value is None: + return None + try: + return int(value) + except ValueError: + return None + + +def collect_provenance(config: Mapping[str, Any]) -> dict[str, Any]: + """Collect git, package, hardware, and exact invocation configuration.""" + repo_root = Path(__file__).resolve().parents[1] + packages: dict[str, str | None] = {} + for package in ( + "omlx", + "mlx", + "mlx-lm", + "mlx-vlm", + "transformers", + "huggingface-hub", + "pyarrow", + ): + try: + packages[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + packages[package] = None + git_head = _run_text_command(("git", "rev-parse", "HEAD"), repo_root) + git_status = _run_text_command( + ("git", "status", "--porcelain", "--untracked-files=no"), repo_root + ) + physical_memory = _run_text_command(("/usr/sbin/sysctl", "-n", "hw.memsize")) + hardware = { + "platform": platform.platform(), + "machine": platform.machine(), + "processor": platform.processor(), + "macos_version": platform.mac_ver()[0], + "chip": _run_text_command( + ("/usr/sbin/sysctl", "-n", "machdep.cpu.brand_string") + ), + "physical_memory_bytes": _parse_optional_int(physical_memory), + } + return { + "git": { + "head": git_head, + "tracked_worktree_clean": ( + git_status == "" if git_status is not None else None + ), + }, + "packages": packages, + "hardware": hardware, + "config_sha256": stable_digest(config), + } + + +def fetch_mrcr_rows( + row_indices: Sequence[int], cache_dir: Path | None +) -> list[dict[str, Any]]: + """Fetch the pinned parquet revision and return validated positional rows.""" + import pyarrow.parquet as parquet + from huggingface_hub import hf_hub_download + + parquet_path = hf_hub_download( + repo_id=DATASET_REPO, + filename=DATASET_FILE, + revision=DATASET_REVISION, + repo_type="dataset", + cache_dir=str(cache_dir) if cache_dir is not None else None, + ) + table = parquet.read_table( + parquet_path, + columns=["prompt", "answer", "random_string_to_prepend"], + ) + selected: list[dict[str, Any]] = [] + seen: set[int] = set() + for index in row_indices: + if isinstance(index, bool) or not isinstance(index, int): + raise ValidationError("MRCR row indices must be integers") + if index < 0 or index >= table.num_rows: + raise ValidationError( + f"MRCR row {index} is outside pinned file length {table.num_rows}" + ) + if index in seen: + raise ValidationError(f"duplicate MRCR row {index}") + seen.add(index) + selected.append( + { + "row_index": index, + "prompt": table["prompt"][index].as_py(), + "answer": table["answer"][index].as_py(), + "random_string_to_prepend": table["random_string_to_prepend"][ + index + ].as_py(), + } + ) + for row in selected: + prompt = row.get("prompt") + answer = row.get("answer") + prefix = row.get("random_string_to_prepend") + if ( + not isinstance(prompt, str) + or not isinstance(answer, str) + or not isinstance(prefix, str) + ): + raise ValidationError(f"MRCR row {row['row_index']} has invalid fields") + messages = json.loads(prompt) + if not isinstance(messages, list) or not all( + isinstance(message, dict) for message in messages + ): + raise ValidationError( + f"MRCR row {row['row_index']} prompt is not a message list" + ) + row["messages"] = messages + row["dataset_prompt_sha256"] = hash_text(prompt) + del row["prompt"] + return selected + + +def probe_memory_pressure() -> MemoryPressureReading: + """Run the mandatory kernel memory-pressure probe.""" + try: + completed = subprocess.run( + ("/usr/bin/memory_pressure", "-Q"), + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise TelemetryError(f"memory_pressure probe failed: {exc}") from exc + if completed.returncode != 0: + raise TelemetryError( + f"memory_pressure exited with status {completed.returncode}" + ) + return parse_memory_pressure_output(completed.stdout) + + +def probe_child_footprint(pid: int) -> int: + """Read a child's kernel phys_footprint or fail closed.""" + from omlx.utils.proc_memory import get_phys_footprint + + try: + value = get_phys_footprint(pid) + except Exception as exc: + raise TelemetryError(f"phys_footprint probe failed: {exc}") from exc + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise TelemetryError("phys_footprint probe returned no valid data") + return value + + +def _terminate_child( + process: subprocess.Popen[bytes], grace_seconds: float = 5.0 +) -> None: + """Terminate a supervised child and escalate only after a bounded grace.""" + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=grace_seconds) + + +def supervise_child( + *, + command: Sequence[str], + environment: Mapping[str, str], + result_path: Path, + poll_interval_seconds: float, + child_limit_bytes: int, + host_minimum_bytes: int, +) -> SupervisedOutcome: + """Run one child while requiring valid safety telemetry on every poll.""" + if poll_interval_seconds <= 0: + raise ValueError("poll interval must be positive") + process = subprocess.Popen(list(command), env=dict(environment)) + peak_child = 0 + minimum_headroom: int | None = None + samples = 0 + error: str | None = None + try: + while True: + returncode = process.poll() + if returncode is not None: + if samples == 0: + error = "child exited before mandatory safety telemetry was sampled" + break + try: + child_phys = probe_child_footprint(process.pid) + pressure = probe_memory_pressure() + samples += 1 + peak_child = max(peak_child, child_phys) + minimum_headroom = ( + pressure.headroom_bytes + if minimum_headroom is None + else min(minimum_headroom, pressure.headroom_bytes) + ) + violation = safety_violation( + child_phys, + pressure.headroom_bytes, + child_limit_bytes=child_limit_bytes, + host_minimum_bytes=host_minimum_bytes, + ) + if violation is not None: + error = violation + _terminate_child(process) + break + except TelemetryError as exc: + # Darwin can stop serving proc_pid_rusage just before waitpid + # observes a normal exit. Allow only a bounded teardown grace. + try: + process.wait(timeout=TELEMETRY_EXIT_GRACE_SECONDS) + except subprocess.TimeoutExpired: + error = str(exc) + _terminate_child(process) + else: + if samples == 0: + error = ( + "child exited before mandatory safety telemetry was sampled" + ) + break + time.sleep(poll_interval_seconds) + + returncode = process.wait() + result: dict[str, Any] | None = None + if error is None: + if not result_path.is_file(): + error = "child exited without an atomic result" + else: + try: + decoded = json.loads(result_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + error = f"child result is invalid: {exc}" + else: + if not isinstance(decoded, dict): + error = "child result is not an object" + else: + result = decoded + if decoded.get("status") != "ok": + error = str( + decoded.get("error") or "child reported failure" + ) + if returncode != 0 and error is None: + error = f"child exited with status {returncode}" + telemetry = { + "samples": samples, + "peak_child_phys_footprint_bytes": peak_child, + "minimum_host_headroom_bytes": minimum_headroom, + "child_limit_bytes": child_limit_bytes, + "host_minimum_bytes": host_minimum_bytes, + "poll_interval_seconds": poll_interval_seconds, + } + return SupervisedOutcome(result, telemetry, returncode, error) + finally: + if process.poll() is None: + _terminate_child(process) + + +def _apply_chat_template(tokenizer: Any, messages: list[dict[str, Any]]) -> str: + """Render a model chat template while explicitly disabling thinking.""" + if not hasattr(tokenizer, "apply_chat_template"): + raise ValidationError("model tokenizer has no chat template") + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + if not isinstance(rendered, str): + raise ValidationError("chat template did not return text") + return rendered + + +def _encode_prompt(tokenizer: Any, text: str) -> list[int]: + """Encode already-templated prompt text using the model tokenizer.""" + encoded = tokenizer.encode(text) + return [int(token) for token in encoded] + + +def _encode_answer(tokenizer: Any, text: str) -> list[int]: + """Encode answer text without introducing a fresh sequence BOS token.""" + try: + encoded = tokenizer.encode(text, add_special_tokens=False) + except TypeError: + encoded = tokenizer.encode(text) + return [int(token) for token in encoded] + + +def _decode_tokens(tokenizer: Any, token_ids: Sequence[int]) -> str: + """Decode generated token IDs without retaining model outputs.""" + try: + decoded = tokenizer.decode(list(token_ids), skip_special_tokens=True) + except TypeError: + decoded = tokenizer.decode(list(token_ids)) + if not isinstance(decoded, str): + raise ValidationError("tokenizer.decode did not return text") + return decoded + + +def _extract_logits(output: Any) -> Any: + """Normalize mlx-lm array and model-output return forms.""" + return output.logits if hasattr(output, "logits") else output + + +def _timed_synchronized(mx: Any, operation: Callable[[], _T]) -> tuple[_T, float]: + """Fence one timed MLX phase on both sides.""" + mx.synchronize() + started = time.perf_counter() + result = operation() + mx.synchronize() + return result, time.perf_counter() - started + + +def _evaluate_cache(mx: Any, cache: list[Any]) -> None: + """Materialize cache state for one completed fixed chunk.""" + mx.eval([cache_object.state for cache_object in cache]) + + +def _prefill_chunk( + mx: Any, model: Any, cache: list[Any], token_ids: Sequence[int] +) -> None: + """Advance one prompt chunk while materializing only bounded cache state.""" + model(mx.array([list(token_ids)]), cache=cache) + _evaluate_cache(mx, cache) + + +def _model_chunk( + mx: Any, model: Any, cache: list[Any], token_ids: Sequence[int] +) -> Any: + """Run one bounded final/decode forward and materialize current logits.""" + output = model(mx.array([list(token_ids)]), cache=cache) + logits = _extract_logits(output) + mx.eval(logits, [cache_object.state for cache_object in cache]) + return logits + + +def _ordinary_convert(mx: Any, cache: list[Any], bits: float) -> dict[str, Any]: + """Use the production ordinary post-prefill conversion implementation.""" + from omlx.scheduler import Scheduler + + owner = SimpleNamespace(_turboquant_kv_bits=bits, _turboquant_skip_last=True) + Scheduler._apply_turboquant_kv_convert(owner, cache) + _evaluate_cache(mx, cache) + return {"path": "ordinary_from_cache"} + + +def _mid_convert(mx: Any, cache: list[Any], bits: float) -> dict[str, Any]: + """Use the bounded existing sliced conversion at the fixed midpoint.""" + from omlx.turboquant_kv import convert_kv_cache_sliced + + stats = convert_kv_cache_sliced(cache, bits=bits, skip_last=True) + _evaluate_cache(mx, cache) + return {"path": "sliced", **asdict(stats)} + + +def _run_direct_prefill( + mx: Any, + model: Any, + prompt_ids: Sequence[int], + mode: ValidationMode, + chunk_size: int, +) -> tuple[list[Any], Any, dict[str, Any]]: + """Run the direct fixed-chunk prompt loop with its exact conversion boundary.""" + from mlx_lm.models.cache import make_prompt_cache + + if len(prompt_ids) < 2: + raise ValidationError("rendered prompt must contain at least two tokens") + cache = make_prompt_cache(model) + cacheable = list(prompt_ids[:-1]) + midpoint = len(cacheable) // 2 + prefix_forward_seconds = 0.0 + suffix_forward_seconds = 0.0 + conversion_seconds = 0.0 + conversion: dict[str, Any] | None = None + + mx.synchronize() + total_started = time.perf_counter() + for start in range(0, midpoint, chunk_size): + end = min(start + chunk_size, midpoint) + _, elapsed = _timed_synchronized( + mx, + partial(_prefill_chunk, mx, model, cache, cacheable[start:end]), + ) + prefix_forward_seconds += elapsed + mx.clear_cache() + + if mode.conversion == "mid": + if mode.bits is None: + raise ValidationError("mid mode omitted its bit width") + conversion, conversion_seconds = _timed_synchronized( + mx, partial(_mid_convert, mx, cache, mode.bits) + ) + + mx.synchronize() + suffix_wall_started = time.perf_counter() + for start in range(midpoint, len(cacheable), chunk_size): + end = min(start + chunk_size, len(cacheable)) + _, elapsed = _timed_synchronized( + mx, + partial(_prefill_chunk, mx, model, cache, cacheable[start:end]), + ) + suffix_forward_seconds += elapsed + mx.clear_cache() + mx.synchronize() + suffix_wall_seconds = time.perf_counter() - suffix_wall_started + + # Production external prefill converts the resident N-1 prompt cache + # before BatchGenerator processes the held final prompt token. + if mode.conversion == "ordinary": + if mode.bits is None: + raise ValidationError("ordinary mode omitted its bit width") + conversion, conversion_seconds = _timed_synchronized( + mx, partial(_ordinary_convert, mx, cache, mode.bits) + ) + + logits, final_seconds = _timed_synchronized( + mx, partial(_model_chunk, mx, model, cache, prompt_ids[-1:]) + ) + suffix_forward_seconds += final_seconds + suffix_wall_seconds += final_seconds + mx.synchronize() + total_seconds = time.perf_counter() - total_started + suffix_tokens = len(prompt_ids) - midpoint + return ( + cache, + logits, + { + "prompt_tokens": len(prompt_ids), + "held_final_prompt_tokens": 1, + "midpoint_token_index": midpoint, + "same_boundary_suffix_tokens": suffix_tokens, + "prefix_forward_seconds": prefix_forward_seconds, + "same_boundary_suffix_forward_seconds": suffix_forward_seconds, + "same_boundary_suffix_seconds": suffix_wall_seconds, + "same_boundary_suffix_tokens_per_second": ( + suffix_tokens / suffix_wall_seconds if suffix_wall_seconds > 0 else 0.0 + ), + "conversion_seconds": conversion_seconds, + "conversion": conversion, + "total_prefill_seconds": total_seconds, + "total_prefill_tokens_per_second": ( + len(prompt_ids) / total_seconds if total_seconds > 0 else 0.0 + ), + }, + ) + + +def _percentile_nearest_rank(values: Sequence[float], percentile: float) -> float: + """Return a deterministic nearest-rank percentile.""" + if not values: + return 0.0 + ordered = sorted(values) + rank = max(1, math.ceil(percentile * len(ordered))) + return float(ordered[rank - 1]) + + +def _model_token_chunk( + mx: Any, + model: Any, + cache: list[Any], + token_id: int, + stream: Any | None, +) -> Any: + """Run one teacher-forced token on the scheduler stream when provided.""" + if stream is None: + return _model_chunk(mx, model, cache, (token_id,)) + with mx.stream(stream): + return _model_chunk(mx, model, cache, (token_id,)) + + +def _teacher_forced_replay( + mx: Any, + model: Any, + cache: list[Any], + initial_logits: Any, + replay_ids: Sequence[int], + *, + stream: Any | None = None, +) -> dict[str, Any]: + """Score and time bounded one-token teacher-forced forwards.""" + if not replay_ids: + raise ValidationError("teacher-forced replay is empty") + logits = initial_logits + latencies: list[float] = [] + nll_values: list[float] = [] + for token_id in replay_ids: + row = logits[0, -1].astype(mx.float32) + nll = mx.logsumexp(row) - row[int(token_id)] + mx.eval(nll) + nll_values.append(float(nll.item())) + + logits, elapsed = _timed_synchronized( + mx, + partial( + _model_token_chunk, + mx, + model, + cache, + int(token_id), + stream, + ), + ) + latencies.append(elapsed) + total_nll = math.fsum(nll_values) + mean_nll = total_nll / len(nll_values) + remaining = nll_values[1:] + remaining_mean = math.fsum(remaining) / len(remaining) if remaining else None + total_seconds = math.fsum(latencies) + return { + "token_count": len(replay_ids), + "token_ids_sha256": hash_token_ids(replay_ids), + "nll_sum": total_nll, + "mean_nll": mean_nll, + "perplexity": math.exp(mean_nll), + "first_token_nll": nll_values[0], + "tokens_2_to_n_count": len(remaining), + "tokens_2_to_n_mean_nll": remaining_mean, + "tokens_2_to_n_perplexity": ( + math.exp(remaining_mean) if remaining_mean is not None else None + ), + "decode_seconds": total_seconds, + "decode_tokens_per_second": ( + len(replay_ids) / total_seconds if total_seconds > 0 else 0.0 + ), + "latency_median_seconds": float(statistics.median(latencies)), + "latency_p95_seconds": _percentile_nearest_rank(latencies, 0.95), + } + + +def _eos_token_ids(tokenizer: Any) -> set[int]: + """Normalize a tokenizer's EOS declaration.""" + raw = getattr(tokenizer, "eos_token_id", None) + if raw is None: + return set() + if isinstance(raw, int) and not isinstance(raw, bool): + return {raw} + if isinstance(raw, Sequence) and not isinstance(raw, str | bytes): + return {int(value) for value in raw} + return set() + + +def _greedy_response( + mx: Any, + model: Any, + cache: list[Any], + initial_logits: Any, + tokenizer: Any, + token_limit: int, +) -> tuple[list[int], str]: + """Generate a separate untimed greedy response for official MRCR grading.""" + logits = initial_logits + generated: list[int] = [] + eos_ids = _eos_token_ids(tokenizer) + for _ in range(token_limit): + token = mx.argmax(logits[0, -1]) + mx.eval(token) + token_id = int(token.item()) + if token_id in eos_ids: + break + generated.append(token_id) + logits = _model_chunk(mx, model, cache, (token_id,)) + return generated, _decode_tokens(tokenizer, generated) + + +def _reset_mlx_measurement(mx: Any) -> dict[str, int]: + """Clear transients, capture separate baselines, and reset allocator peak.""" + mx.synchronize() + mx.clear_cache() + mx.synchronize() + baseline = { + "mlx_active_baseline_bytes": int(mx.get_active_memory()), + "mlx_cache_baseline_bytes": int(mx.get_cache_memory()), + } + mx.reset_peak_memory() + return baseline + + +def _finish_mlx_measurement(mx: Any, baseline: Mapping[str, int]) -> dict[str, int]: + """Capture MLX peak/final counters without conflating phys_footprint.""" + mx.synchronize() + return { + **dict(baseline), + "mlx_peak_bytes": int(mx.get_peak_memory()), + "mlx_active_final_bytes": int(mx.get_active_memory()), + "mlx_cache_final_bytes": int(mx.get_cache_memory()), + } + + +def _load_direct_model( + model_path: str, mode: ValidationMode, trust: bool +) -> tuple[Any, Any]: + """Load one direct-loop model with production compatibility transforms.""" + from omlx.model_settings import ModelSettings + from omlx.utils.model_loading import ( + apply_post_load_transforms, + lm_load_compat, + materialize_lazy_state, + maybe_apply_pre_load_patches, + ) + from omlx.utils.tokenizer import get_tokenizer_config + + settings = ModelSettings( + enable_thinking=False, + turboquant_kv_enabled=mode.bits is not None, + turboquant_mid_prefill=mode.conversion == "mid", + turboquant_kv_bits=mode.bits or 4.0, + turboquant_skip_last=True, + ) + maybe_apply_pre_load_patches(model_path, model_settings=settings) + tokenizer_config = get_tokenizer_config(model_path, trust_remote_code=trust) + model, tokenizer = lm_load_compat( + model_path, + tokenizer_config=tokenizer_config, + trust_remote_code=trust, + ) + model = apply_post_load_transforms(model, settings) + materialize_lazy_state(model) + if mode.bits is not None: + from omlx.patches.turboquant_attention import apply_turboquant_attention_patch + + apply_turboquant_attention_patch() + return model, tokenizer + + +def run_matrix_child(spec: Mapping[str, Any]) -> dict[str, Any]: + """Execute one fresh real-model matrix cell.""" + import mlx.core as mx + + mode_data = spec["mode"] + mode = ValidationMode( + str(mode_data["name"]), + str(mode_data["conversion"]), + float(mode_data["bits"]) if mode_data["bits"] is not None else None, + ) + sample = spec["sample"] + metadata = dict(spec["metadata"]) + model: Any | None = None + tokenizer: Any | None = None + primary_cache: list[Any] | None = None + greedy_cache: list[Any] | None = None + try: + model, tokenizer = _load_direct_model( + str(spec["model_path"]), mode, bool(spec.get("trust_remote_code", False)) + ) + rendered = _apply_chat_template(tokenizer, list(sample["messages"])) + prompt_ids = _encode_prompt(tokenizer, rendered) + if len(prompt_ids) > int(metadata["max_prompt_tokens"]): + raise ValidationError( + f"rendered prompt has {len(prompt_ids)} tokens, above safety limit " + f"{metadata['max_prompt_tokens']}" + ) + answer = str(sample["answer"]) + answer_ids = _encode_answer(tokenizer, answer) + replay_ids = answer_ids[:FIXED_REPLAY_TOKENS] + baseline = _reset_mlx_measurement(mx) + + primary_cache, primary_logits, prefill = _run_direct_prefill( + mx, model, prompt_ids, mode, int(metadata["chunk_size"]) + ) + replay = _teacher_forced_replay( + mx, model, primary_cache, primary_logits, replay_ids + ) + primary_cache = None + primary_logits = None + gc.collect() + mx.synchronize() + mx.clear_cache() + + greedy_cache, greedy_logits, _ = _run_direct_prefill( + mx, model, prompt_ids, mode, int(metadata["chunk_size"]) + ) + generated_ids, response = _greedy_response( + mx, + model, + greedy_cache, + greedy_logits, + tokenizer, + int(metadata["greedy_token_limit"]), + ) + score = official_retrieval_score( + response, + answer, + str(sample["random_string_to_prepend"]), + ) + memory = _finish_mlx_measurement(mx, baseline) + return { + "schema_version": SCHEMA_VERSION, + "status": "ok", + "kind": "matrix", + "mode": mode.name, + "metadata": metadata, + "metrics": { + "dataset_prompt_sha256": sample["dataset_prompt_sha256"], + "rendered_prompt_sha256": hash_text(rendered), + "prompt_token_ids_sha256": hash_token_ids(prompt_ids), + "answer_sha256": hash_text(answer), + "answer_token_ids_sha256": hash_token_ids(answer_ids), + "prompt_token_count": len(prompt_ids), + "answer_token_count": len(answer_ids), + "teacher_forced_replay_truncated": len(answer_ids) > len(replay_ids), + "prefill": prefill, + "teacher_forced": replay, + "greedy_retrieval": { + "response": response, + "response_sha256": hash_text(response), + "token_count": len(generated_ids), + "token_ids_sha256": hash_token_ids(generated_ids), + "required_prefix": sample["random_string_to_prepend"], + "official_score": score, + }, + "memory": memory, + }, + } + finally: + primary_cache = None + greedy_cache = None + model = None + tokenizer = None + gc.collect() + try: + mx.synchronize() + mx.clear_cache() + except Exception: + logger.exception("failed to clear MLX cache after matrix child") + + +def _find_pool_model_id(pool: Any, model_path: Path) -> str: + """Resolve the discovered EnginePool ID for an exact local model path.""" + expected = model_path.expanduser().resolve() + matches = [ + model_id + for model_id in pool.get_model_ids() + if Path(pool.get_entry(model_id).model_path).expanduser().resolve() == expected + ] + if len(matches) != 1: + raise ValidationError( + f"EnginePool discovered {len(matches)} entries for exact model path {expected}" + ) + return matches[0] + + +def _capture_initial_prefill_pause( + mx: Any, + scheduler: Any, + request: Any, + prompt_ids: Sequence[int], +) -> OrganicPrefillPause: + """Run the first organic attempt and require the production LRU pause.""" + from omlx.scheduler import _PrefillEvictionNeeded + + captured: list[Any] = [] + original = scheduler._new_prefill_context + + def _capture( + inner_request: Any, + prompt_cache: list[Any], + *, + loop_label: str, + ) -> Any: + context = original(inner_request, prompt_cache, loop_label=loop_label) + captured.append(context) + return context + + scheduler._new_prefill_context = _capture + baseline = _reset_mlx_measurement(mx) + mx.synchronize() + total_started = time.perf_counter() + attempt_started = time.perf_counter() + try: + try: + completed = scheduler._do_external_prefill( + request, + list(prompt_ids), + existing_cache=None, + ) + except _PrefillEvictionNeeded as exc: + mx.synchronize() + external_seconds = time.perf_counter() - attempt_started + eviction_request = exc.request + else: + del completed + raise ValidationError( + "organic first prefill attempt completed without the production LRU pause" + ) + if len(captured) != 1: + raise ValidationError( + f"organic first attempt captured {len(captured)} contexts instead of one" + ) + context = captured[0] + if context.mid_triggered: + raise ValidationError( + "organic first attempt converted before the LRU pause" + ) + return OrganicPrefillPause( + baseline=baseline, + context=context, + eviction_request=eviction_request, + external_seconds=external_seconds, + total_started=total_started, + ) + finally: + scheduler._new_prefill_context = original + gc.collect() + mx.synchronize() + mx.clear_cache() + + +def _capture_external_prefill_after_pause( + mx: Any, + scheduler: Any, + request: Any, + prompt_ids: Sequence[int], + replay_ids: Sequence[int], + pause: OrganicPrefillPause, + eviction_callback_result: bool, + eviction_pause_seconds: float, +) -> dict[str, Any]: + """Retry organic prefill after LRU handling and require one conversion.""" + captured: list[Any] = [] + original = scheduler._new_prefill_context + + def _capture( + inner_request: Any, + prompt_cache: list[Any], + *, + loop_label: str, + ) -> Any: + context = original(inner_request, prompt_cache, loop_label=loop_label) + captured.append(context) + return context + + scheduler._new_prefill_context = _capture + try: + (cache, last_token), retry_external_seconds = _timed_synchronized( + mx, + partial( + scheduler._do_external_prefill, + request, + list(prompt_ids), + existing_cache=None, + ), + ) + external_completed = time.perf_counter() + with mx.stream(scheduler._stream): + logits, final_seconds = _timed_synchronized( + mx, + partial( + _model_chunk, + mx, + scheduler.model, + cache, + tuple(last_token), + ), + ) + total_seconds = time.perf_counter() - pause.total_started + replay = _teacher_forced_replay( + mx, + scheduler.model, + cache, + logits, + replay_ids, + stream=scheduler._stream, + ) + if len(captured) != 1: + raise ValidationError( + f"organic retry captured {len(captured)} contexts instead of one" + ) + context = captured[0] + if context is pause.context: + raise ValidationError("organic retry reused the first attempt's context") + if pause.context.mid_triggered: + raise ValidationError("organic first attempt converted before its pause") + if not context.mid_triggered: + raise ValidationError( + "organic retry completed without mid-prefill conversion" + ) + if int(request.prefill_eviction_retries) != 1: + raise ValidationError( + "organic request did not retain exactly one production LRU retry" + ) + post_trigger_seconds = ( + external_completed - context.conversion_completed_at + if context.conversion_completed_at is not None + else 0.0 + ) + memory = _finish_mlx_measurement(mx, pause.baseline) + cache = None + logits = None + external_attempt_seconds = [ + pause.external_seconds, + retry_external_seconds, + ] + return { + "actual_mid_prefill_trigger": True, + "prefill_attempt_count": 2, + "eviction_pause_count": 1, + "prefill_eviction_retries": int(request.prefill_eviction_retries), + "eviction_callback_reclaimed": bool(eviction_callback_result), + "eviction_pause_seconds": eviction_pause_seconds, + "trigger_tokens": context.trigger_tokens, + "conversion_seconds": float(context.conversion_seconds), + "converted_layers": int(context.converted_layers), + "conversion_slices": int(context.conversion_slices), + "skipped_dense_layers": int(context.skipped_dense_layers), + "external_prefill_attempt_seconds": external_attempt_seconds, + "external_prefill_seconds": sum(external_attempt_seconds), + "held_final_prompt_token_seconds": final_seconds, + "total_prefill_seconds": total_seconds, + "total_prefill_tokens_per_second": ( + len(prompt_ids) / total_seconds if total_seconds > 0 else 0.0 + ), + "post_trigger_tokens": int(context.post_trigger_tokens), + "post_trigger_seconds": post_trigger_seconds, + "post_trigger_tokens_per_second": ( + context.post_trigger_tokens / post_trigger_seconds + if post_trigger_seconds > 0 + else 0.0 + ), + "teacher_forced": replay, + "memory": memory, + } + finally: + scheduler._new_prefill_context = original + gc.collect() + mx.synchronize() + mx.clear_cache() + + +def _configure_organic_pressure( + scheduler: Any, + *, + soft_limit_bytes: int, + hard_limit_bytes: int, + prefill_abort_margin: float, + prefill_min_chunk_tokens: int, +) -> None: + """Apply normal propagated scheduler pressure limits without a trigger hook.""" + if soft_limit_bytes <= 0 or hard_limit_bytes <= 0: + raise ValidationError("organic scheduler limits must be positive") + if soft_limit_bytes > hard_limit_bytes: + raise ValidationError("organic soft limit exceeds hard limit") + if not 0.0 < prefill_abort_margin <= 1.0: + raise ValidationError("organic prefill abort margin must be in (0, 1]") + if prefill_min_chunk_tokens <= 0: + raise ValidationError("organic prefill minimum chunk must be positive") + scheduler._memory_limit_bytes = soft_limit_bytes + scheduler._memory_hard_limit_bytes = hard_limit_bytes + scheduler._memory_hard_watermark_bytes = hard_limit_bytes + scheduler._memory_abort_limit_bytes = hard_limit_bytes + scheduler._memory_static_ceiling_bytes = hard_limit_bytes + scheduler._memory_dynamic_ceiling_bytes = hard_limit_bytes + scheduler._memory_metal_cap_bytes = hard_limit_bytes + scheduler._memory_guard_tier = "custom" + scheduler._prefill_abort_margin = prefill_abort_margin + scheduler._prefill_min_chunk_tokens = prefill_min_chunk_tokens + scheduler._prefill_memory_guard = True + scheduler._memory_limits_propagated = True + + +async def _run_organic_child_async(spec: Mapping[str, Any]) -> dict[str, Any]: + """Load via EnginePool and run the production external-prefill pressure path.""" + import mlx.core as mx + + from omlx.engine_pool import EnginePool + from omlx.model_settings import ModelSettings + from omlx.request import Request, SamplingParams + from omlx.scheduler import SchedulerConfig + + model_path = Path(str(spec["model_path"])).expanduser().resolve() + chunk_size = int(spec["chunk_size"]) + scheduler_config = SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=8192, + completion_batch_size=1, + prefill_step_size=chunk_size, + chunked_prefill=False, + prefill_speed_priority=False, + paged_ssd_cache_dir=None, + hot_cache_max_size=0, + ) + pool = EnginePool(scheduler_config=scheduler_config) + engine: Any | None = None + scheduler: Any | None = None + request: Any | None = None + try: + pool.discover_models(str(model_path.parent)) + model_id = _find_pool_model_id(pool, model_path) + settings = ModelSettings( + enable_thinking=False, + trust_remote_code=bool(spec.get("trust_remote_code", False)), + turboquant_kv_enabled=True, + turboquant_mid_prefill=True, + turboquant_kv_bits=8.0, + turboquant_skip_last=True, + ) + engine = await pool.get_engine( + model_id, + force_lm=True, + runtime_settings=settings, + ) + if pool.loaded_model_count != 1: + raise ValidationError( + "organic validation did not retain exclusive model count" + ) + rendered = _apply_chat_template( + engine.tokenizer, + list(spec["sample"]["messages"]), + ) + rendered_ids = _encode_prompt(engine.tokenizer, rendered) + required = ORGANIC_PROMPT_TOKENS + FIXED_REPLAY_TOKENS + if len(rendered_ids) < required: + raise ValidationError( + f"organic row rendered to {len(rendered_ids)} tokens, requires {required}" + ) + prompt_ids = rendered_ids[:ORGANIC_PROMPT_TOKENS] + replay_ids = rendered_ids[ORGANIC_PROMPT_TOKENS:required] + core = engine._engine.engine + await core.stop() + scheduler = core.scheduler + if scheduler._prefill_eviction_callback_configured is not True: + raise ValidationError("organic engine has no production prefill callback") + _configure_organic_pressure( + scheduler, + soft_limit_bytes=int(spec["scheduler_soft_limit_bytes"]), + hard_limit_bytes=int(spec["scheduler_hard_limit_bytes"]), + prefill_abort_margin=float(spec["prefill_abort_margin"]), + prefill_min_chunk_tokens=int(spec["prefill_min_chunk_tokens"]), + ) + request = Request( + request_id="turboquant-organic-validation", + prompt=list(prompt_ids), + sampling_params=SamplingParams( + max_tokens=FIXED_REPLAY_TOKENS, + temperature=0.0, + ), + prompt_token_ids=list(prompt_ids), + num_prompt_tokens=len(prompt_ids), + skip_cache_store=True, + ) + scheduler.requests[request.request_id] = request + loop = asyncio.get_running_loop() + pause = await loop.run_in_executor( + core._mlx_executor, + _capture_initial_prefill_pause, + mx, + scheduler, + request, + prompt_ids, + ) + eviction_started = time.perf_counter() + eviction_callback_result = await pool._evict_idle_lru_for_prefill( + exclude_model_id=model_id, + eviction_request=pause.eviction_request, + ) + eviction_pause_seconds = time.perf_counter() - eviction_started + metrics = await loop.run_in_executor( + core._mlx_executor, + _capture_external_prefill_after_pause, + mx, + scheduler, + request, + prompt_ids, + replay_ids, + pause, + eviction_callback_result, + eviction_pause_seconds, + ) + return { + "schema_version": SCHEMA_VERSION, + "status": "ok", + "kind": "organic", + "mode": "organic-mid-q8", + "metadata": dict(spec["metadata"]), + "metrics": { + "dataset_prompt_sha256": spec["sample"]["dataset_prompt_sha256"], + "full_rendered_sha256": hash_text(rendered), + "full_rendered_token_count": len(rendered_ids), + "prompt_token_ids_sha256": hash_token_ids(prompt_ids), + "prompt_token_count": len(prompt_ids), + "replay_token_ids_sha256": hash_token_ids(replay_ids), + "replay_token_count": len(replay_ids), + **metrics, + }, + } + finally: + if scheduler is not None and request is not None: + scheduler.requests.pop(request.request_id, None) + await pool.shutdown() + engine = None + gc.collect() + try: + mx.synchronize() + mx.clear_cache() + except Exception: + logger.exception("failed to clear MLX cache after organic child") + + +def run_organic_child(spec: Mapping[str, Any]) -> dict[str, Any]: + """Run the asynchronous organic child entry point.""" + return asyncio.run(_run_organic_child_async(spec)) + + +def run_child(spec_path: Path, result_path: Path) -> int: + """Execute a private child spec and always publish an atomic result.""" + try: + decoded = json.loads(spec_path.read_text(encoding="utf-8")) + if not isinstance(decoded, dict): + raise ValidationError("child spec is not an object") + kind = decoded.get("kind") + if kind == "matrix": + result = run_matrix_child(decoded) + elif kind == "organic": + result = run_organic_child(decoded) + else: + raise ValidationError(f"unknown child kind: {kind!r}") + atomic_write_json(result_path, result) + return 0 + except BaseException as exc: + logger.exception("validation child failed") + atomic_write_json( + result_path, + { + "schema_version": SCHEMA_VERSION, + "status": "failed", + "error": f"{type(exc).__name__}: {exc}", + }, + ) + return 1 + + +def _ensure_outside_model(model_path: Path, candidate: Path, label: str) -> None: + """Reject every configured write location below the read-only model path.""" + model_root = model_path.expanduser().resolve() + resolved = candidate.expanduser().resolve() + if resolved == model_root or resolved.is_relative_to(model_root): + raise ValidationError(f"{label} must not be under the model path") + + +def _validate_runtime_write_locations( + model_path: Path, + *, + output_path: Path, + dataset_cache_dir: Path | None, +) -> None: + """Ensure every known cache/temp/result location is outside model sources.""" + _ensure_outside_model(model_path, output_path, "output") + if dataset_cache_dir is not None: + _ensure_outside_model(model_path, dataset_cache_dir, "dataset cache") + _ensure_outside_model( + model_path, + Path(tempfile.gettempdir()), + "temporary directory", + ) + for variable in ( + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "TRANSFORMERS_CACHE", + "XDG_CACHE_HOME", + "TMPDIR", + "TMP", + "TEMP", + ): + value = os.environ.get(variable) + if value: + _ensure_outside_model( + model_path, + Path(value), + f"{variable} environment location", + ) + + +def _run_supervised_spec( + *, + spec: Mapping[str, Any], + workspace: Path, + sequence: int, + args: argparse.Namespace, +) -> SupervisedOutcome: + """Persist one spec and supervise its fresh child.""" + spec_path = workspace / f"spec-{sequence:03d}.json" + result_path = workspace / f"result-{sequence:03d}.json" + atomic_write_json(spec_path, spec) + command = build_child_command(Path(__file__), spec_path, result_path) + return supervise_child( + command=command, + environment=child_environment(os.environ), + result_path=result_path, + poll_interval_seconds=float(args.poll_interval), + child_limit_bytes=int(float(args.child_memory_gib) * GIB), + host_minimum_bytes=int(float(args.host_headroom_gib) * GIB), + ) + + +def _matrix_config(args: argparse.Namespace) -> dict[str, Any]: + """Build JSON-compatible matrix configuration provenance.""" + return { + "model_path": str(Path(args.model).expanduser().resolve()), + "dataset_cache_dir": ( + str(Path(args.dataset_cache_dir).expanduser().resolve()) + if args.dataset_cache_dir + else None + ), + "dataset": pinned_dataset_config(), + "rows": list(args.rows), + "modes": [asdict(mode) for mode in matrix_modes()], + "chunk_size": args.chunk_size, + "max_prompt_tokens": args.max_prompt_tokens, + "greedy_token_limit": args.greedy_token_limit, + "teacher_forced_replay_token_limit": FIXED_REPLAY_TOKENS, + "thinking_enabled": False, + "child_memory_limit_bytes": int(float(args.child_memory_gib) * GIB), + "host_headroom_minimum_bytes": int(float(args.host_headroom_gib) * GIB), + "poll_interval_seconds": args.poll_interval, + "trust_remote_code": args.trust_remote_code, + } + + +def _organic_config(args: argparse.Namespace) -> dict[str, Any]: + """Build JSON-compatible organic configuration provenance.""" + return { + "model_path": str(Path(args.model).expanduser().resolve()), + "dataset_cache_dir": ( + str(Path(args.dataset_cache_dir).expanduser().resolve()) + if args.dataset_cache_dir + else None + ), + "dataset": { + **pinned_dataset_config(), + "row_index": ORGANIC_ROW_INDEX, + }, + "mode": "organic-mid-q8", + "prompt_tokens": ORGANIC_PROMPT_TOKENS, + "teacher_forced_replay_tokens": FIXED_REPLAY_TOKENS, + "chunk_size": args.chunk_size, + "scheduler": { + "max_num_seqs": 1, + "max_num_batched_tokens": 8192, + "completion_batch_size": 1, + "prefill_step_size": args.chunk_size, + "chunked_prefill": False, + "prefill_speed_priority": False, + "prefill_abort_margin": ORGANIC_PREFILL_ABORT_MARGIN, + "prefill_min_chunk_tokens": ORGANIC_PREFILL_MIN_CHUNK_TOKENS, + }, + "turboquant_bits": 8.0, + "turboquant_skip_last": True, + "exclusive_ownership": True, + "no_cache": True, + "forced_trigger": False, + "thinking_enabled": False, + "scheduler_soft_limit_bytes": int(float(args.scheduler_soft_limit_gib) * GIB), + "scheduler_hard_limit_bytes": int(float(args.child_memory_gib) * GIB), + "prefill_abort_margin": ORGANIC_PREFILL_ABORT_MARGIN, + "prefill_min_chunk_tokens": ORGANIC_PREFILL_MIN_CHUNK_TOKENS, + "child_memory_limit_bytes": int(float(args.child_memory_gib) * GIB), + "host_headroom_minimum_bytes": int(float(args.host_headroom_gib) * GIB), + "poll_interval_seconds": args.poll_interval, + "trust_remote_code": args.trust_remote_code, + } + + +def run_parent(args: argparse.Namespace) -> int: + """Run all public work under fresh-child supervision and tensor verification.""" + model_path = Path(args.model).expanduser().resolve() + output_path = Path(args.output).expanduser().resolve() + cache_dir = ( + Path(args.dataset_cache_dir).expanduser().resolve() + if args.dataset_cache_dir + else None + ) + _validate_runtime_write_locations( + model_path, + output_path=output_path, + dataset_cache_dir=cache_dir, + ) + config = _matrix_config(args) if args.command == "matrix" else _organic_config(args) + provenance = collect_provenance(config) + before = build_model_manifest(model_path) + results: list[dict[str, Any]] = [] + error: str | None = None + atomic_write_json( + output_path, + shape_validation_checkpoint( + kind=str(args.command), + provenance=provenance, + config=config, + before_manifest=before, + results=results, + ), + ) + try: + row_indices = ( + list(args.rows) if args.command == "matrix" else [ORGANIC_ROW_INDEX] + ) + rows = fetch_mrcr_rows(row_indices, cache_dir) + with tempfile.TemporaryDirectory( + prefix="omlx-turboquant-validation-" + ) as temporary: + workspace = Path(temporary) + sequence = 0 + if args.command == "matrix": + for sample in rows: + metadata = matrix_shared_metadata( + row_index=int(sample["row_index"]), + chunk_size=int(args.chunk_size), + max_prompt_tokens=int(args.max_prompt_tokens), + greedy_token_limit=int(args.greedy_token_limit), + ) + for mode in matrix_modes(): + spec = { + "schema_version": SCHEMA_VERSION, + "kind": "matrix", + "model_path": str(model_path), + "trust_remote_code": bool(args.trust_remote_code), + "mode": asdict(mode), + "sample": sample, + "metadata": metadata, + } + outcome = _run_supervised_spec( + spec=spec, + workspace=workspace, + sequence=sequence, + args=args, + ) + sequence += 1 + child_result = ( + dict(outcome.result) + if outcome.result is not None + else { + "schema_version": SCHEMA_VERSION, + "status": "failed", + "kind": "matrix", + "mode": mode.name, + "metadata": metadata, + "error": outcome.error, + } + ) + child_result["supervision"] = outcome.telemetry + results.append(child_result) + atomic_write_json( + output_path, + shape_validation_checkpoint( + kind=str(args.command), + provenance=provenance, + config=config, + before_manifest=before, + results=results, + ), + ) + if outcome.error is not None: + raise ValidationError(outcome.error) + validate_matrix_results(results, expected_rows=args.rows) + else: + sample = rows[0] + metadata = { + "dataset": { + "repo_id": DATASET_REPO, + "revision": DATASET_REVISION, + "filename": DATASET_FILE, + "row_index": ORGANIC_ROW_INDEX, + }, + "thinking_enabled": False, + "prompt_tokens": ORGANIC_PROMPT_TOKENS, + "teacher_forced_replay_tokens": FIXED_REPLAY_TOKENS, + "chunk_size": int(args.chunk_size), + "no_cache": True, + "exclusive_ownership": True, + "forced_trigger": False, + "prefill_abort_margin": ORGANIC_PREFILL_ABORT_MARGIN, + "prefill_min_chunk_tokens": ORGANIC_PREFILL_MIN_CHUNK_TOKENS, + } + spec = { + "schema_version": SCHEMA_VERSION, + "kind": "organic", + "model_path": str(model_path), + "trust_remote_code": bool(args.trust_remote_code), + "sample": sample, + "metadata": metadata, + "chunk_size": int(args.chunk_size), + "scheduler_soft_limit_bytes": int( + float(args.scheduler_soft_limit_gib) * GIB + ), + "scheduler_hard_limit_bytes": int( + float(args.child_memory_gib) * GIB + ), + "prefill_abort_margin": ORGANIC_PREFILL_ABORT_MARGIN, + "prefill_min_chunk_tokens": ORGANIC_PREFILL_MIN_CHUNK_TOKENS, + } + outcome = _run_supervised_spec( + spec=spec, + workspace=workspace, + sequence=sequence, + args=args, + ) + child_result = ( + dict(outcome.result) + if outcome.result is not None + else { + "schema_version": SCHEMA_VERSION, + "status": "failed", + "kind": "organic", + "mode": "organic-mid-q8", + "metadata": metadata, + "error": outcome.error, + } + ) + child_result["supervision"] = outcome.telemetry + results.append(child_result) + atomic_write_json( + output_path, + shape_validation_checkpoint( + kind=str(args.command), + provenance=provenance, + config=config, + before_manifest=before, + results=results, + ), + ) + if outcome.error is not None: + raise ValidationError(outcome.error) + except BaseException as exc: + logger.exception("validation parent failed") + error = f"{type(exc).__name__}: {exc}" + try: + after = build_model_manifest(model_path) + except BaseException as exc: + logger.exception("post-run tensor manifest failed") + after = [] + manifest_error = f"{type(exc).__name__}: {exc}" + error = manifest_error if error is None else f"{error}; {manifest_error}" + payload = shape_validation_result( + kind=str(args.command), + provenance=provenance, + config=config, + before_manifest=before, + after_manifest=after, + results=results, + error=error, + ) + atomic_write_json(output_path, payload) + return 0 if payload["status"] == "ok" else 1 + + +def _add_common_parent_arguments(parser: argparse.ArgumentParser) -> None: + """Add shared public supervisor arguments.""" + parser.add_argument( + "--model", required=True, help="Local read-only model directory" + ) + parser.add_argument("--output", required=True, help="Atomic provenance/result JSON") + parser.add_argument( + "--dataset-cache-dir", + help="Optional Hugging Face dataset cache (must be outside the model)", + ) + parser.add_argument( + "--chunk-size", + type=int, + default=2048, + help="Fixed prefill chunk size (default: 2048)", + ) + parser.add_argument( + "--child-memory-gib", + type=float, + default=36.0, + help="Terminate at this child phys_footprint in GiB (default: 36)", + ) + parser.add_argument( + "--host-headroom-gib", + type=float, + default=6.0, + help="Terminate below this kernel host headroom in GiB (default: 6)", + ) + parser.add_argument( + "--poll-interval", + type=float, + default=0.25, + help="Safety telemetry polling interval in seconds (default: 0.25)", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow model remote code during loading", + ) + + +def build_parser() -> argparse.ArgumentParser: + """Build the typed validation CLI parser.""" + parser = argparse.ArgumentParser( + prog="validate_turboquant_mid_prefill.py", + description=( + "Validate TurboQuant mid-prefill quality, throughput, and memory " + "with fail-closed parent supervision." + ), + ) + subparsers = parser.add_subparsers(dest="command", required=True) + matrix = subparsers.add_parser( + "matrix", + help="Run pinned MRCR rows in dense, ordinary Q8/Q4, and mid Q8/Q4 modes", + ) + _add_common_parent_arguments(matrix) + matrix.add_argument( + "--rows", + type=int, + nargs="+", + default=list(DEFAULT_MATRIX_ROWS), + help="Pinned parquet row indices (default: 104 109 136 301 311 328)", + ) + matrix.add_argument( + "--max-prompt-tokens", + type=int, + default=131072, + help="Reject, never truncate, rendered prompts above this count", + ) + matrix.add_argument( + "--greedy-token-limit", + type=int, + default=1024, + help="Maximum untimed greedy tokens for official MRCR scoring", + ) + + organic = subparsers.add_parser( + "organic", + help="Run the natural 131071-token production scheduler pressure path", + ) + _add_common_parent_arguments(organic) + organic.add_argument( + "--scheduler-soft-limit-gib", + type=float, + default=32.0, + help="Normal scheduler pressure target in GiB (default: 32)", + ) + + child = subparsers.add_parser("__child", help=argparse.SUPPRESS) + child.add_argument("--spec", required=True) + child.add_argument("--result", required=True) + return parser + + +def _validate_cli_arguments(args: argparse.Namespace) -> None: + """Reject nonsensical limits before hashing or launching children.""" + if args.command == "__child": + return + if args.chunk_size <= 0: + raise ValidationError("chunk size must be positive") + if args.child_memory_gib <= 0: + raise ValidationError("child memory limit must be positive") + if args.host_headroom_gib < 0: + raise ValidationError("host headroom minimum must be non-negative") + if args.poll_interval <= 0: + raise ValidationError("poll interval must be positive") + if args.command == "matrix": + if args.max_prompt_tokens <= 1: + raise ValidationError("max prompt tokens must exceed one") + if args.greedy_token_limit <= 0: + raise ValidationError("greedy token limit must be positive") + if tuple(args.rows) != DEFAULT_MATRIX_ROWS: + raise ValidationError("matrix rows are pinned to 104 109 136 301 311 328") + if args.command == "organic": + if args.scheduler_soft_limit_gib <= 0: + raise ValidationError("scheduler soft limit must be positive") + if args.scheduler_soft_limit_gib > args.child_memory_gib: + raise ValidationError("scheduler soft limit exceeds child hard limit") + + +def main(argv: Sequence[str] | None = None) -> int: + """CLI entry point.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + parser = build_parser() + args = parser.parse_args(argv) + try: + _validate_cli_arguments(args) + if args.command == "__child": + return run_child(Path(args.spec), Path(args.result)) + return run_parent(args) + except BaseException as exc: + logger.error("validation failed: %s: %s", type(exc).__name__, exc) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integration/test_server_endpoints.py b/tests/integration/test_server_endpoints.py index 7e5ef36ee..1be0bc100 100644 --- a/tests/integration/test_server_endpoints.py +++ b/tests/integration/test_server_endpoints.py @@ -7,10 +7,11 @@ """ import json +from collections.abc import AsyncIterator from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -852,6 +853,217 @@ def test_completion_with_list_prompt(self, client): data = response.json() assert "choices" in data + def test_completion_forwards_scalar_preflight_eviction_attempt( + self, + client: TestClient, + mock_llm_engine: MockBaseEngine, + ) -> None: + mock_llm_engine.preflight_completion = AsyncMock(return_value=True) + mock_llm_engine.generate = AsyncMock( + return_value=MockGenerationOutput(text="Generated response.") + ) + + response = client.post( + "/v1/completions", + json={"model": "test-model", "prompt": "First prompt"}, + ) + + assert response.status_code == 200 + assert ( + mock_llm_engine.generate.call_args.kwargs[ + "prefill_eviction_callback_attempted" + ] + is True + ) + + def test_completion_keeps_list_preflight_attempts_prompt_local( + self, + client: TestClient, + mock_llm_engine: MockBaseEngine, + ) -> None: + mock_llm_engine.preflight_completion = AsyncMock( + side_effect=[True, None], + ) + mock_llm_engine.generate = AsyncMock( + side_effect=[ + MockGenerationOutput(text="First response."), + MockGenerationOutput(text="Second response."), + ] + ) + + response = client.post( + "/v1/completions", + json={ + "model": "test-model", + "prompt": ["First prompt", "Second prompt"], + }, + ) + + assert response.status_code == 200 + first_kwargs = mock_llm_engine.generate.call_args_list[0].kwargs + second_kwargs = mock_llm_engine.generate.call_args_list[1].kwargs + assert first_kwargs["prefill_eviction_callback_attempted"] is True + assert "prefill_eviction_callback_attempted" not in second_kwargs + + def test_completion_stream_forwards_preflight_eviction_attempt( + self, + client: TestClient, + mock_llm_engine: MockBaseEngine, + ) -> None: + captured: dict[str, object] = {} + mock_llm_engine.preflight_completion = AsyncMock(return_value=True) + + async def _recording_stream_generate( + prompt: str, + **kwargs: object, + ) -> AsyncIterator[MockGenerationOutput]: + captured.update(kwargs) + yield MockGenerationOutput( + text="Hi", + new_text="Hi", + finished=True, + finish_reason="stop", + ) + + mock_llm_engine.stream_generate = _recording_stream_generate + + response = client.post( + "/v1/completions", + json={ + "model": "test-model", + "prompt": "Stream prompt", + "stream": True, + }, + ) + + assert response.status_code == 200 + assert captured["prefill_eviction_callback_attempted"] is True + + def test_completion_stream_rejects_unsupported_turboquant_layout_before_start( + self, + client: TestClient, + mock_llm_engine: MockBaseEngine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from mlx_lm.models.cache import RotatingKVCache + + import omlx.scheduler as scheduler_mod + from omlx.scheduler import Scheduler, SchedulerConfig + + model = MagicMock() + model.layers = [] + model.config = SimpleNamespace( + model_type="unit", + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + hidden_size=64, + head_dim=32, + ) + model.make_cache.return_value = [RotatingKVCache(max_size=128)] + tokenizer = MagicMock() + tokenizer.eos_token_id = 2 + scheduler = Scheduler( + model=model, + tokenizer=tokenizer, + config=SchedulerConfig( + max_num_seqs=1, + prefill_step_size=2048, + paged_cache_block_size=0, + ), + ) + assert scheduler._turboquant_preflight_conversion_eligible is False + + scheduler._prefill_memory_guard = True + scheduler._memory_hard_limit_bytes = 1 + scheduler._prefill_eviction_callback_configured = False + scheduler._turboquant_mid_prefill = True + scheduler._turboquant_kv_bits = 4.0 + scheduler._prefill_tq_kv_dtype_size = 1.0 + + def _zero_usage(*, refresh_mlx_active: bool = True) -> int: + del refresh_mlx_active + return 0 + + def _exclusive(_owner: object) -> bool: + return True + + monkeypatch.setattr(scheduler, "_current_usage_bytes", _zero_usage) + monkeypatch.setattr( + scheduler_mod._conversion_coordinator, + "process_exclusive", + _exclusive, + ) + + async def _preflight_completion( + prompt: str, + *, + request_id: str | None = None, + **kwargs: object, + ) -> bool | None: + del prompt, kwargs + scheduler.preflight_or_raise( + num_prompt_tokens=65536, + request_id=request_id, + ) + return None + + stream_started = False + + async def _stream_generate( + prompt: str, + **kwargs: object, + ) -> AsyncIterator[MockGenerationOutput]: + nonlocal stream_started + del prompt, kwargs + stream_started = True + yield MockGenerationOutput(text="must not stream") + + mock_llm_engine.preflight_completion = _preflight_completion + mock_llm_engine.stream_generate = _stream_generate + + response = client.post( + "/v1/completions", + json={ + "model": "test-model", + "prompt": "Rejected stream prompt", + "stream": True, + }, + headers={"x-request-id": "unsupported-layout"}, + ) + + assert response.status_code == 400 + assert response.headers["content-type"].startswith("application/json") + body = response.json() + assert body["error"]["code"] == "prefill_memory_exceeded" + assert body["error"]["omlx_code"] == "prefill_memory_exceeded" + assert body["error"]["estimated_bytes"] > body["error"]["limit_bytes"] + assert stream_started is False + + def test_completion_preflight_failure_does_not_generate( + self, + client: TestClient, + mock_llm_engine: MockBaseEngine, + ) -> None: + from omlx.exceptions import PrefillMemoryExceededError + + mock_llm_engine.preflight_completion = AsyncMock( + side_effect=PrefillMemoryExceededError( + message="rejected after callback", + request_id="failed-completion", + ) + ) + mock_llm_engine.generate = AsyncMock() + + response = client.post( + "/v1/completions", + json={"model": "test-model", "prompt": "Rejected prompt"}, + headers={"x-request-id": "failed-completion"}, + ) + + assert response.status_code == 400 + mock_llm_engine.generate.assert_not_awaited() + def test_completion_includes_cached_tokens_on_cache_hit( self, client, mock_llm_engine ): diff --git a/tests/test_admin_model_settings_template.py b/tests/test_admin_model_settings_template.py index e7c7ee14e..4bfe9c877 100644 --- a/tests/test_admin_model_settings_template.py +++ b/tests/test_admin_model_settings_template.py @@ -1,5 +1,6 @@ """Regression tests for admin model-settings UI gates.""" +import json from pathlib import Path @@ -14,7 +15,7 @@ def _section(html: str, start_marker: str, end_marker: str) -> str: return html.split(start_marker, 1)[1].split(end_marker, 1)[0] -def test_lightning_mtp_and_turboquant_are_not_ui_mutexed(): +def test_lightning_mtp_and_turboquant_are_not_ui_mutexed() -> None: html = _model_settings_template() turboquant = _section( @@ -32,7 +33,7 @@ def test_lightning_mtp_and_turboquant_are_not_ui_mutexed(): assert "modelSettings.turboquant_kv_enabled" not in lightning_mtp -def test_vlm_mtp_still_conflicts_with_turboquant(): +def test_vlm_mtp_still_conflicts_with_turboquant() -> None: html = _model_settings_template() vlm_mtp = _section( html, @@ -43,7 +44,83 @@ def test_vlm_mtp_still_conflicts_with_turboquant(): assert "modelSettings.turboquant_kv_enabled" in vlm_mtp -def test_reasoning_effort_offers_max_after_high(): +def test_mid_prefill_toggle_is_nested_under_turboquant_parent() -> None: + html = _model_settings_template() + turboquant = _section( + html, + "", + "", + ) + + parent_gate = 'x-show="modelSettings.turboquant_kv_enabled"' + child_key = "modelSettings.turboquant_mid_prefill" + assert parent_gate in turboquant + assert child_key in turboquant + assert turboquant.index(parent_gate) < turboquant.index(child_key) + assert 'id="turboquant-mid-prefill-label"' in turboquant + assert 'id="turboquant-kv-label"' in turboquant + assert 'id="turboquant-kv-hint"' in turboquant + assert 'aria-labelledby="turboquant-kv-label"' in turboquant + assert 'aria-describedby="turboquant-kv-hint"' in turboquant + assert ( + ":aria-checked=\"modelSettings.turboquant_kv_enabled ? 'true' : 'false'\"" + in turboquant + ) + assert 'for="turboquant-kv-bits"' in turboquant + assert 'id="turboquant-kv-bits"' in turboquant + assert 'role="switch"' in turboquant + assert 'aria-labelledby="turboquant-mid-prefill-label"' in turboquant + assert 'id="turboquant-mid-prefill-hint"' in turboquant + assert 'aria-describedby="turboquant-mid-prefill-hint"' in turboquant + assert ( + ":aria-checked=\"modelSettings.turboquant_mid_prefill ? 'true' : 'false'\"" + in turboquant + ) + + +def test_mid_prefill_dashboard_state_payload_profile_and_reset_strings() -> None: + root = Path(__file__).resolve().parents[1] + js = (root / "omlx/admin/static/js/dashboard.js").read_text() + + assert "turboquant_mid_prefill: s.turboquant_mid_prefill || false" in js + assert "out.turboquant_mid_prefill = !!ms.turboquant_mid_prefill;" in js + assert "turboquant_mid_prefill: !!this.modelSettings.turboquant_mid_prefill" in js + assert "out.turboquant_mid_prefill = !!ms.turboquant_kv_enabled" not in js + assert "turboquant_mid_prefill: this.modelSettings.turboquant_kv_enabled" not in js + assert "this.modelSettings.turboquant_mid_prefill = false;" in js + assert "'turboquant_mid_prefill'," in js + assert "this.profileFields.model_specific" in js + + +def test_locales_include_mid_prefill_copy_and_parent_semantics() -> None: + root = Path(__file__).resolve().parents[1] + locale_paths = sorted((root / "omlx/admin/i18n").glob("*.json")) + assert len(locale_paths) == 9 + + for path in locale_paths: + strings = json.loads(path.read_text()) + assert strings["modal.model_settings.turboquant_kv_hint"] + assert strings["modal.model_settings.turboquant_mid_prefill"] + assert strings["modal.model_settings.turboquant_mid_prefill_hint"] + + english = json.loads((root / "omlx/admin/i18n/en.json").read_text()) + assert english["modal.model_settings.turboquant_kv_hint"] == ( + "Compress the KV cache with vector quantization after prefill for generation. " + "Ordinary TurboQuant does not convert during cold prefill. Lower bits use " + "less memory; higher bits preserve more quality." + ) + assert ( + english["modal.model_settings.turboquant_mid_prefill"] + == "Convert under prefill pressure" + ) + assert english["modal.model_settings.turboquant_mid_prefill_hint"] == ( + "When a full prefill chunk cannot fit, convert the growing KV cache once " + "and continue with TurboQuant. Requires this to be the only loaded model. " + "Adds a one-time pause and may slow the rest of prefill." + ) + + +def test_reasoning_effort_offers_max_after_high() -> None: html = _model_settings_template() high_option = '