Skip to content

DFlash bridge keeps emitting tokens after a stop token in the same block #265

Closed

Description

@solderzzc

Found while reviewing #262 (merged as 7858bf5, fixes #261). Related to #224, which made the stop path cancel the DFlash stream.

Behaviour. The stop-token branch in dflashGenerationStream (Sources/SwiftLM/Server.swift, the .token case) calls onStopToken. The default handler returns .more, so the generation is not ended and the bridge keeps forwarding the committed tokens of the same DFlash block to onToken. The standard path (Libraries/MLXLMCommon/Evaluate.swift) ends on any stop id, whatever the handler's disposition. So under --dflash, text after the end-of-turn token can reach the client, and a <tool_call> written after the stop token can be parsed as a tool call.

Reproduction. On main at 7858bf5, a scratch test in DFlashStreamBridgeTests (not committed) feeds .token 5, then stop 99, then 7 and 8, with stopTokenIDs: [99] and the default SpellingHandler:

continuation.yield(.token(tokenID: 5, generatedTokens: 1, acceptanceRatio: 0, cyclesCompleted: 1))
continuation.yield(.token(tokenID: 99, generatedTokens: 2, acceptanceRatio: 0, cyclesCompleted: 1))
continuation.yield(.token(tokenID: 7, generatedTokens: 3, acceptanceRatio: 0, cyclesCompleted: 1))
continuation.yield(.token(tokenID: 8, generatedTokens: 4, acceptanceRatio: 0, cyclesCompleted: 1))
// ... summary, finish ...
XCTAssertEqual(chunks, ["t5 "], "tokens after the stop token were emitted")

Result: chunks == ["t5 ", "t7 ", "t8 "], so the test fails. The existing testStopTokenGoesToOnStopTokenNotOnToken ends on the stop token, so it cannot catch this.

Suggested fix. Treat a stop id as the end of the generation, unless the handler already cancelled:

let isStop = stopTokenIDs.contains(tokenID)
let disposition = isStop
    ? handler.onStopToken(tokenID, logProbabilities: nil, emit: emit)
    : handler.onToken(tokenID, logProbabilities: nil, emit: emit)
if case .more = disposition, !isStop { continue }
// stop ids fall through to the existing finish block: onGenerationEnd, .info with .stop, cancel

Add the regression test above to DFlashStreamBridgeTests. It should also check that the bridge yields exactly one .info with stopReason: .stop.

Severity: medium. Only tokens from the same DFlash block (up to blockTokens - 1) can leak, but the leak includes tool-call text.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions