Skip to content

Add multi-branch (beam) search to the constrained decoder #518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FuJacob merged 1 commit into from
Jun 2, 2026
Merged

Add multi-branch (beam) search to the constrained decoder #518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions Cotabby.xcodeproj/project.pbxproj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
279F017530A86AF62EB17918 /* EmojiSynonymCatalog.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0846DE4E0293AF13890620D3 /* EmojiSynonymCatalog.swift */; };
27D4F5CACADE171F142178B4 /* SettingsSidebarView.swift in Sources */ = {isa = PBXBuildFile; fileRef = BADB38D0160B47637572FC5E /* SettingsSidebarView.swift */; };
286B7022E2A2774275004447 /* WelcomeTemplateStepView.swift in Sources */ = {isa = PBXBuildFile; fileRef = A9199B9CEAB320982CA333B8 /* WelcomeTemplateStepView.swift */; };
29EC35D67D9B4C3C50222619 /* ConstrainedBeamSearchTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */; };
2C6159231472A849F15BD0AE /* ScreenFrameReader.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5484C8A04B9C00CF79D589EB /* ScreenFrameReader.swift */; };
2DF5A3826AAB99C279EBB8DE /* InputMonitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = B81DD30EB657368AACE9625A /* InputMonitor.swift */; };
2E3DEB7E89D0146274596F2E /* SettingsContainerView.swift in Sources */ = {isa = PBXBuildFile; fileRef = DB0CE9AB1286367BA2E82392 /* SettingsContainerView.swift */; };
Expand Down Expand Up @@ -200,6 +201,7 @@
C2C958D6E5F5FE1CCC414BCE /* SuggestionSubsystemContracts.swift in Sources */ = {isa = PBXBuildFile; fileRef = DEB16474A67CE1D210B944C9 /* SuggestionSubsystemContracts.swift */; };
C4C6734678797669055988E0 /* AppUpdateManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = AD9573F3504CAE6891DF9B7D /* AppUpdateManager.swift */; };
C618C5595DA9C57C806A3E03 /* SettingsAttentionEvaluatorTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2BC293F6125E2B14DCF05AD9 /* SettingsAttentionEvaluatorTests.swift */; };
C638466A2F373BABA6F60A6D /* ConstrainedBeamSearch.swift in Sources */ = {isa = PBXBuildFile; fileRef = A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */; };
C71B594433F3B411CAE5DE7E /* FocusCapabilityResolverTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D4F6D5F94B238F7B4BE7C247 /* FocusCapabilityResolverTests.swift */; };
C9B815652CED38966C53A5E8 /* EmojiVariantResolverTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = EE8BB19D8EC9A75CD3458A6B /* EmojiVariantResolverTests.swift */; };
CA5B2D226FBAA5419E78F14F /* SuggestionSessionReconciler.swift in Sources */ = {isa = PBXBuildFile; fileRef = CE0AA0503128B0FC3951D700 /* SuggestionSessionReconciler.swift */; };
Expand Down Expand Up @@ -423,6 +425,7 @@
A52D0B550E00EF173A5D157E /* LlamaRuntimeManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LlamaRuntimeManager.swift; sourceTree = "<group>"; };
A804F4DB6FD9BC8C27B2B65F /* LlamaRuntimeModels.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LlamaRuntimeModels.swift; sourceTree = "<group>"; };
A829F28F01FAE76CA7244BBC /* ModelFileValidatorTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelFileValidatorTests.swift; sourceTree = "<group>"; };
A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ConstrainedBeamSearch.swift; sourceTree = "<group>"; };
A854CAFB1F557BC4CAED8819 /* VisualContextCoordinator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = VisualContextCoordinator.swift; sourceTree = "<group>"; };
A863F41C0C03D7B4AC5DC002 /* MarkerSelectionSynthesizer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MarkerSelectionSynthesizer.swift; sourceTree = "<group>"; };
A9199B9CEAB320982CA333B8 /* WelcomeTemplateStepView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WelcomeTemplateStepView.swift; sourceTree = "<group>"; };
Expand Down Expand Up @@ -479,6 +482,7 @@
D5D6C2318E405AA717D1C256 /* WelcomePermissionStepView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WelcomePermissionStepView.swift; sourceTree = "<group>"; };
D814BBA41CF29E8DD9954651 /* OnboardingTemplateFeatureListTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OnboardingTemplateFeatureListTests.swift; sourceTree = "<group>"; };
D84D4528EEC9EFEB8AE8E318 /* ActivationIndicatorController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ActivationIndicatorController.swift; sourceTree = "<group>"; };
D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ConstrainedBeamSearchTests.swift; sourceTree = "<group>"; };
D9C1C921A1CDA2ADFC39EA01 /* AppsPaneView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppsPaneView.swift; sourceTree = "<group>"; };
DB0CE9AB1286367BA2E82392 /* SettingsContainerView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsContainerView.swift; sourceTree = "<group>"; };
DDE858CB1E687E3CEB8FDD5B /* SuggestionRequestFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SuggestionRequestFactory.swift; sourceTree = "<group>"; };
Expand Down Expand Up @@ -741,6 +745,7 @@
90B0D133AB77A2503FB08827 /* ClipboardRelevanceFilterTests.swift */,
D504BEB224E0C176F5FCFF6E /* CompletionRenderModePolicyTests.swift */,
06FF2B0A3094A952A8EBA9B5 /* ConfidenceSuppressionPolicyTests.swift */,
D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */,
6B0200AA5C4A644A3B52A3EC /* ConstrainedSamplerTests.swift */,
AF1E065C7FFB697FCEB2FA5C /* CotabbyTestFixtures.swift */,
AD752451330486FE270018B0 /* CustomRulesTests.swift */,
Expand Down Expand Up @@ -904,6 +909,7 @@
D3A2AC525DC664DB540D4F19 /* ClipboardRelevanceFilter.swift */,
53CF416511099C6818110F01 /* CompletionRenderModePolicy.swift */,
1BD71ECC2AE4821B643E0935 /* ConfidenceSuppressionPolicy.swift */,
A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */,
47CF010A66DA31989524FCD0 /* ConstrainedSampler.swift */,
C7B2D34A6F3AC9DFD61350F7 /* CotabbyDebugOptions.swift */,
29ED42C4BDD0C521101AF95E /* DeviceInfo.swift */,
Expand Down Expand Up @@ -1099,6 +1105,7 @@
7C94725B4837DEC9ECF1BC54 /* CompletionRenderMode.swift in Sources */,
3985F0F2B3178DBB945B1064 /* CompletionRenderModePolicy.swift in Sources */,
429CE592897D8A952F2916C3 /* ConfidenceSuppressionPolicy.swift in Sources */,
C638466A2F373BABA6F60A6D /* ConstrainedBeamSearch.swift in Sources */,
81FC391E375B7EEFF965FF1B /* ConstrainedSampler.swift in Sources */,
8B2DFC860803C0A7C4D34A36 /* ContextBuffer.swift in Sources */,
AA2E09FF7E430D66ECA8ECD5 /* CotabbyApp.swift in Sources */,
Expand Down Expand Up @@ -1262,6 +1269,7 @@
BFCA7FAFDAEBF586AB615567 /* ClipboardRelevanceFilterTests.swift in Sources */,
25F91CEF38400FD1ADB6B1AF /* CompletionRenderModePolicyTests.swift in Sources */,
91D8189EFCD1BA992EA6F038 /* ConfidenceSuppressionPolicyTests.swift in Sources */,
29EC35D67D9B4C3C50222619 /* ConstrainedBeamSearchTests.swift in Sources */,
1C00CE3D553B58723EAE5F92 /* ConstrainedSamplerTests.swift in Sources */,
5E10EFC426217CB7218A5847 /* CotabbyTestFixtures.swift in Sources */,
91D1F16B8C5DA281D4B7F699 /* CustomRulesTests.swift in Sources */,
Expand Down
6 changes: 6 additions & 0 deletions Cotabby/Models/LlamaRuntimeModels.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,12 @@ struct LlamaGenerationOptions: Equatable, Sendable {
/// validated on device. Changing it does not affect KV reuse, so it is intentionally excluded
/// from `SamplingFingerprint`.
var useConstrainedDecoder: Bool = false

/// Beam width for the constrained decoder. 1 keeps the single-path greedy decode; values > 1 run a
/// multi-branch beam search that explores several short continuations and keeps the highest-scoring
/// one. Only consulted when `useConstrainedDecoder` is true. Like `useConstrainedDecoder`, it does
/// not affect KV reuse, so it is excluded from `SamplingFingerprint`.
var beamWidth: Int = 1
}

/// The concrete runtime assets selected during bootstrap after checking available model files.
Expand Down
135 changes: 132 additions & 3 deletions Cotabby/Services/Runtime/LlamaRuntimeCore.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,9 +198,16 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable {
// The KV-trim defer above runs after whichever decoder returns. Both decoders share the
// prepared sequence and the same confidence-suppression contract; they differ only in how
// they pick each token (engine sampler vs. deterministic constrained selection).
return options.useConstrainedDecoder
? try runConstrainedDecode(sequenceID: sequenceID, options: options)
: runEngineSampledDecode(sequenceID: sequenceID, options: options)
guard options.useConstrainedDecoder else {
return runEngineSampledDecode(sequenceID: sequenceID, options: options)
}
return options.beamWidth > 1
? try runConstrainedBeamDecode(
sequenceID: sequenceID,
promptTokenCount: promptTokens.count,
options: options
)
: try runConstrainedDecode(sequenceID: sequenceID, options: options)
}

// MARK: - Decoders
Expand Down Expand Up @@ -406,6 +413,128 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable {
return generatedText
}

/// Multi-branch (beam) variant of the constrained decoder. Explores several short continuations
/// over the shared sequence — `EngineBeamStepper` syncs the KV cache to each branch's token path —
/// and returns the highest-scoring one. Reuses the same token profile, no-repeat-ngram guard,
/// sentence-boundary stop, and confidence suppression as the greedy path. The caller's KV-trim
/// defer restores the prompt-only state afterward.
private func runConstrainedBeamDecode(
sequenceID: Int32,
promptTokenCount: Int,
options: LlamaGenerationOptions
) throws -> String {
let profile = try autocompleteTokenProfile()
let vocabSize = profile.vocabSize
guard vocabSize > 0 else {
throw LlamaRuntimeError.generationFailed("Vocabulary unavailable for constrained decoding.")
}
let topK = options.topK > 0 ? options.topK : vocabSize
var currentPath: [Int] = []
var logitsBuffer = [Float](repeating: 0, count: vocabSize)
let candidates = ConstrainedBeamSearch.search(
nextLogits: { generatedTokens in
self.beamLogits(
forGeneratedTokens: generatedTokens,
sequenceID: sequenceID,
promptTokenCount: promptTokenCount,
currentPath: &currentPath,
logitsBuffer: &logitsBuffer
)
},
profile: profile,
configuration: BeamSearchConfiguration(
beamWidth: options.beamWidth,
maxTokens: options.maxPredictionTokens,
topK: topK,
noRepeatNgramSize: Self.noRepeatNgramSize
),
isSingleLine: options.singleLine
)
let best = candidates.first
CotabbyLogger.runtime.debug(
"Decode end",
metadata: [
"kind": .string("generate_beam"),
"beam_width": .stringConvertible(options.beamWidth),
"candidates": .stringConvertible(candidates.count),
"tokens_generated": .stringConvertible(best?.tokenIDs.count ?? 0)
]
)
guard let best else {
return ""
}
if Self.shouldSuppress(
sumLogprob: best.cumulativeLogprob,
tokensGenerated: best.tokenIDs.count,
options: options
) {
return ""
}
return best.text
}

/// Beam-search logits provider: syncs the shared sequence's KV to `generatedTokens`, then reads
/// the next-token logits. `currentPath` / `logitsBuffer` are owned by one beam run (the caller),
/// so this stays a plain method on the runtime where `engine` (a noncopyable C++ value) is mutable.
private func beamLogits(
forGeneratedTokens generatedTokens: [Int],
sequenceID: Int32,
promptTokenCount: Int,
currentPath: inout [Int],
logitsBuffer: inout [Float]
) -> [Float]? {
guard syncBeamSequence(
to: generatedTokens,
sequenceID: sequenceID,
promptTokenCount: promptTokenCount,
currentPath: &currentPath
) else {
return nil
}
let vocabSize = logitsBuffer.count
let written = logitsBuffer.withUnsafeMutableBufferPointer { buffer in
Int(engine.getNextTokenLogits(sequenceID, buffer.baseAddress, Int32(buffer.count)))
}
guard written == vocabSize else {
return nil
}
return logitsBuffer
}

/// Brings the sequence KV to exactly `target` tokens beyond the prompt: trim back to the longest
/// shared prefix with the current path, then accept the remaining target tokens. `currentPath` is
/// updated as tokens are accepted so it always reflects the real KV length, even on a mid-accept
/// failure (the caller treats a false return as "this branch cannot be extended").
private func syncBeamSequence(
to target: [Int],
sequenceID: Int32,
promptTokenCount: Int,
currentPath: inout [Int]
) -> Bool {
let shared = Self.commonPrefixLength(currentPath, target)
if currentPath.count > shared, !engine.trimKV(sequenceID, Int32(promptTokenCount + shared)) {
currentPath = []
return false
}
currentPath = Array(target[..<shared])
for index in shared ..< target.count {
guard engine.acceptToken(sequenceID, Int32(target[index])) == .ok else {
return false
}
currentPath.append(target[index])
}
return true
}

private static func commonPrefixLength(_ lhs: [Int], _ rhs: [Int]) -> Int {
var count = 0
let limit = min(lhs.count, rhs.count)
while count < limit, lhs[count] == rhs[count] {
count += 1
}
return count
}

/// Shared low-confidence gate for both decoders: drop completions the model itself was unsure
/// about. Disabled by default (confidenceFloor == -infinity). The KV-trim defer in `generate`
/// still runs because the caller returns "" rather than throwing.
Expand Down
12 changes: 11 additions & 1 deletion Cotabby/Services/Runtime/LlamaSuggestionEngine.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,15 @@ final class LlamaSuggestionEngine {
UserDefaults.standard.bool(forKey: constrainedDecoderDefaultsKey)
}

/// UserDefaults key (no UI) for the constrained decoder's beam width. Default 1 keeps the existing
/// single-path greedy decode; a value > 1 runs a multi-branch beam search. Paired with the
/// constrained-decoder flag as a hidden developer/dogfood knob until validated on device.
private static let constrainedBeamWidthDefaultsKey = "cotabbyConstrainedBeamWidth"
private static var constrainedBeamWidth: Int {
let stored = UserDefaults.standard.integer(forKey: constrainedBeamWidthDefaultsKey)
return stored > 0 ? stored : 1
}

init(runtimeManager: LlamaRuntimeGenerating) {
self.runtimeManager = runtimeManager
}
Expand Down Expand Up @@ -60,7 +69,8 @@ final class LlamaSuggestionEngine {
precedingText: request.context.precedingText,
trailingText: request.context.trailingText
),
useConstrainedDecoder: Self.isConstrainedDecoderEnabled
useConstrainedDecoder: Self.isConstrainedDecoderEnabled,
beamWidth: Self.constrainedBeamWidth
)
)
try Task.checkCancellation()
Expand Down
Loading