Skip to content

feat(runtime): mid-stream cancellation #161

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
john-rocky merged 1 commit into from
Apr 30, 2026
Merged

feat(runtime): mid-stream cancellation #161

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
19 changes: 18 additions & 1 deletion Sources/CoreMLLLM/CoreMLLLM.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1003,7 +1003,13 @@ public final class CoreMLLLM: @unchecked Sendable {
}

return AsyncStream { continuation in
Task(priority: decodePriority) {
// T4 cancellation: AsyncStream does NOT auto-cancel the producer
// task when the consumer breaks out of `for await`. Bind a
// handle so onTermination can call cancel(). Combined with
// Task.checkCancellation() at chunk and decode-step boundaries,
// a user-initiated cancel (or stream deinit) stops generation
// promptly instead of running to maxDecode.
let genTask = Task(priority: decodePriority) {
do {
let IMAGE_TOKEN_ID = 258880
let AUDIO_TOKEN_ID = 258881
Expand Down Expand Up @@ -1077,6 +1083,7 @@ public final class CoreMLLLM: @unchecked Sendable {
engine.currentPosition = prefillLen

for step in prefillLen..<tokens.count {
try Task.checkCancellation()
let tid = tokens[step]
try autoreleasepool {
if let emb = multimodalEmbedding(for: tid) {
Expand All @@ -1090,6 +1097,7 @@ public final class CoreMLLLM: @unchecked Sendable {
}
} else {
for (step, tid) in tokens.enumerated() {
try Task.checkCancellation()
try autoreleasepool {
if let emb = multimodalEmbedding(for: tid) {
nextID = try engine.predictStep(tokenID: 0, position: step,
Expand Down Expand Up @@ -1161,6 +1169,10 @@ public final class CoreMLLLM: @unchecked Sendable {
decodeLoop: while tokenCount < maxDecode {
if eosIDs.contains(Int(nid)) { break }
if engine.currentPosition >= ctxLimit { break }
// T4: stop the loop the moment the consumer
// unsubscribes — wired via continuation.onTermination
// → genTask.cancel() below.
try Task.checkCancellation()

let useEagle = (eagleSpec != nil) && didFirstDecode
&& engine.canSpeculate
Expand Down Expand Up @@ -1343,11 +1355,16 @@ public final class CoreMLLLM: @unchecked Sendable {
pos += 1
}
}
} catch is CancellationError {
// T4: user-cancelled. No log spam — this is expected.
} catch {
print("[CoreMLLLM] Error: \(error)")
}
continuation.finish()
}
continuation.onTermination = { @Sendable _ in
genTask.cancel()
}
}
}

Expand Down
Loading