From ea66df6fef20e2926f61cd41f735ee57ba98c45b Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Tue, 6 Oct 2026 18:23:32 +0200 Subject: [PATCH] feat: --thinking-budget caps reasoning tokens per request Adds --thinking-budget N (default off) and a per-request thinking_budget field. After N reasoning tokens the model's own is forced and it answers; 0 closes the block immediately. Built on mlx-swift-lm's ThinkingBudgetProcessor with an immediate transition, since Qwen3.5 declares none of its own. A budget the request cannot honour (max_tokens too small for the budget plus the answer) is an OpenAI-shaped 400 before any prefill. It applies to the plain, hybrid-cache and sliding-window paths; speculative, MTP and DFlash decoding log a warning and run without it. Also fixes the non-streaming response keeping a leading in the content when the reasoning was empty, which a zero budget hits every time. Co-Authored-By: Claude Sonnet 5.5 --- Sources/SwiftLM/Server.swift | 85 ++++++++++-- tests/SwiftLMTests/ThinkingBudgetTests.swift | 136 +++++++++++++++++++ 2 files changed, 213 insertions(+), 8 deletions(-) create mode 100644 tests/SwiftLMTests/ThinkingBudgetTests.swift diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index e5cd29a..927f75e 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -564,6 +564,9 @@ struct MLXServer: AsyncParsableCommand { @Flag(name: .long, help: "Enable thinking/reasoning mode (Qwen3.5 etc). Default: disabled") var thinking: Bool = false + @Option(name: .long, help: "Hard cap on reasoning tokens per request (overridable per request with thinking_budget). After N thinking tokens the server closes the thinking block and the model answers. 0 closes it immediately. Needs thinking enabled and --max-tokens (or the request's max_tokens) large enough to hold the budget plus the answer. Default: no cap") + var thinkingBudget: Int? + @Flag(name: .long, help: "Do not echo generated tokens to stdout as they stream. Request log lines are unaffected.") var noTokenEcho: Bool = false @@ -628,6 +631,9 @@ struct MLXServer: AsyncParsableCommand { var mtpAssistantModel: String? func validate() throws { + if let thinkingBudget, thinkingBudget < 0 { + throw ValidationError("--thinking-budget must be 0 or more.") + } if vision && noVision { throw ValidationError("--vision and --no-vision are mutually exclusive.") } @@ -1382,6 +1388,7 @@ struct MLXServer: AsyncParsableCommand { minP: self.minP, repeatPenalty: self.repeatPenalty, thinking: self.thinking, + thinkingBudget: self.thinkingBudget, maxPromptTokens: self.maxPromptTokens, promptCacheEntries: self.promptCacheEntries, tokenEcho: !self.noTokenEcho, @@ -1699,6 +1706,8 @@ struct ServerConfig: Sendable { let minP: Float? let repeatPenalty: Float? let thinking: Bool + /// `--thinking-budget`: default cap on reasoning tokens. nil = no cap. + let thinkingBudget: Int? /// `--max-prompt-tokens`: longer prompts are rejected before prefill. nil = no limit. let maxPromptTokens: Int? /// `--prompt-cache-entries`: LRU capacity of the prompt cache. @@ -2555,6 +2564,26 @@ func handleChatCompletion( } let promptTokens = lmInput.text.tokens.asArray(Int.self) + // ── Reasoning budget: per-request `thinking_budget` wins over `--thinking-budget` ── + let budgetComponents: GenerationComponents? + do { + budgetComponents = try thinkingBudgetComponents( + budget: chatReq.thinkingBudget ?? config.thinkingBudget, enableThinking: enableThinking, + reasoning: await container.configuration.reasoningConfig, + tokenizer: await container.tokenizer, parameters: params) + } catch let error as ThinkingBudgetError { + print("srv slot_reject: id 0 | thinking budget: \(error.errorDescription ?? "invalid")") + slot.release() + await stats.requestFinished(tokens: 0, duration: 0) + return Response( + status: .badRequest, headers: jsonHeaders(), + body: .init(byteBuffer: ByteBuffer(string: thinkingBudgetErrorBody(error)))) + } + let genComponents = budgetComponents ?? GenerationComponents() + if budgetComponents != nil, draftModelRef != nil || config.mtp || dflashModel != nil { + print("[SwiftLM] Warning: the thinking budget is not applied with speculative, MTP or DFlash decoding.") + } + // ── Issue #108: does the template leave the thinking block open? ── // Qwen3/3.5/3.6 templates append `\n` to the generation prompt when // enable_thinking is true, so the model emits reasoning with no opening tag of @@ -2729,7 +2758,7 @@ func handleChatCompletion( } stream = try MLXLMCommon.generate( input: LMInput(text: sliceText(lmInput.text, from: boundary)), - cache: cache, parameters: params, context: context) + cache: cache, parameters: params, context: context, components: genComponents) } else if let boundary = rotatingBoundary { var start = await promptCache.restore(newTokens: promptTokens, into: cache) ?? 0 if start >= promptLength { @@ -2753,7 +2782,7 @@ func handleChatCompletion( } stream = try MLXLMCommon.generate( input: LMInput(text: sliceText(lmInput.text, from: start)), - cache: cache, parameters: params, context: context) + cache: cache, parameters: params, context: context, components: genComponents) } else if let draftRef = draftModelRef { // Speculative decoding path: draft model generates candidates, main model verifies. // Bypass prompt cache to avoid draft/main KV drift on partial-match restores. @@ -2781,7 +2810,8 @@ func handleChatCompletion( ) } else { stream = try MLXLMCommon.generate( - input: trimmedInput, cache: cache, parameters: params, context: context + input: trimmedInput, cache: cache, parameters: params, context: context, + components: genComponents ) } } else { @@ -2792,7 +2822,8 @@ func handleChatCompletion( ) } else { stream = try MLXLMCommon.generate( - input: lmInput, cache: cache, parameters: params, context: context + input: lmInput, cache: cache, parameters: params, context: context, + components: genComponents ) } } @@ -3479,10 +3510,10 @@ func handleChatNonStreaming( print("srv debug: pre-extract fullText=\(fullText.prefix(40).debugDescription)") let (extracted, remaining) = extractThinkingBlock(from: fullText, alreadyOpen: thinkingPreOpened) print("srv debug: extracted=\(extracted != nil ? "true" : "false"), remaining_len=\(remaining.count)") - if let extracted { - reasoningContent = extracted - responseContent = remaining - } + // `remaining` is already `fullText` when no tag was found, and it drops the closing + // tag when the reasoning was empty (`` first, e.g. a zero thinking budget). + reasoningContent = extracted + responseContent = remaining } // ── JSON mode validation ── @@ -4142,6 +4173,41 @@ func promptTooLongBody(promptTokens: Int, limit: Int?) -> String? { return "{\"error\":{\"message\":\"This model's maximum context length is \(limit) tokens. However, your messages resulted in \(promptTokens) tokens.\",\"type\":\"invalid_request_error\",\"code\":\"context_length_exceeded\"}}" } +/// Components that cap reasoning at `budget` tokens, or nil when no cap applies (no budget, +/// thinking off, or a model that declares no reasoning protocol). Throws `ThinkingBudgetError` +/// for a budget the request cannot honour, e.g. a `max_tokens` too small for the budget plus +/// the answer. +/// +/// The transition is `.immediate`: the model's own `` closes the span with no extra +/// text. Qwen3.5 declares no validated hard-budget transition, so this is an opt-in. +func thinkingBudgetComponents( + budget: Int?, enableThinking: Bool, reasoning: ReasoningConfig?, + tokenizer: any MLXLMCommon.Tokenizer, parameters: GenerateParameters +) throws -> GenerationComponents? { + guard let budget, enableThinking, let reasoning else { return nil } + let configuration = try ThinkingBudgetConfiguration( + maximumTokenCount: budget, transitionOverride: .immediate) + let components = try GenerationComponents().applyingThinkingBudget( + configuration, reasoning: reasoning, tokenizer: tokenizer) + try components.validate(parameters: parameters) + return components +} + +/// OpenAI-shaped 400 body for a reasoning budget the request cannot honour. +func thinkingBudgetErrorBody(_ error: ThinkingBudgetError) -> String { + let payload: [String: Any] = ["error": [ + "message": error.errorDescription ?? "Invalid thinking budget.", + "type": "invalid_request_error", + "code": "invalid_thinking_budget", + ]] + guard let data = try? JSONSerialization.data(withJSONObject: payload), + let json = String(data: data, encoding: .utf8) + else { + return "{\"error\":{\"message\":\"Invalid thinking budget.\",\"type\":\"invalid_request_error\",\"code\":\"invalid_thinking_budget\"}}" + } + return json +} + /// Releases the slot and returns a 400 when the prompt is over `--max-prompt-tokens`, /// else nil. Runs before prefill, so a rejected prompt costs no GPU work. func rejectLongPrompt(promptTokens: Int, limit: Int?, slot: GenerationSlot, stats: ServerStats) async -> Response? { @@ -4556,6 +4622,8 @@ struct ChatCompletionRequest: Decodable { let responseFormat: ResponseFormat? /// Per-request Jinja template kwargs (e.g. {"enable_thinking": false} for Qwen3/Qwen3.5) let chatTemplateKwargs: [String: Bool]? + /// Per-request reasoning-token cap; overrides `--thinking-budget`. + let thinkingBudget: Int? /// Top-level thinking override emitted by Aegis-AI gateway let enableThinking: Bool? /// Number of bits for native MLX quantized KV cache (nil = no quantization). @@ -4577,6 +4645,7 @@ struct ChatCompletionRequest: Decodable { case responseFormat = "response_format" case chatTemplateKwargs = "chat_template_kwargs" case enableThinking = "enable_thinking" + case thinkingBudget = "thinking_budget" case kvBits = "kv_bits" } } diff --git a/tests/SwiftLMTests/ThinkingBudgetTests.swift b/tests/SwiftLMTests/ThinkingBudgetTests.swift new file mode 100644 index 0000000..82669d0 --- /dev/null +++ b/tests/SwiftLMTests/ThinkingBudgetTests.swift @@ -0,0 +1,136 @@ +import XCTest +import Foundation +import MLXLMCommon +@testable import SwiftLM + +/// `--thinking-budget` / `thinking_budget` cap a request's reasoning tokens. A budget the +/// request cannot honour is a 400 before any GPU work, and a request that does not think is +/// left alone. +final class ThinkingBudgetTests: XCTestCase { + + private let tokenizer = ByteTokenizer() + /// Qwen-style tags that the byte tokenizer can spell; Qwen3.5 declares no hard-budget + /// transition of its own, so the server supplies `.immediate`. + private let reasoning = ReasoningConfig( + startDelimiter: "", endDelimiter: "", + promptStrategy: .templateFlag(key: "enable_thinking", defaultOn: true), + isSpecialToken: false) + + private func components( + budget: Int?, thinking: Bool = true, reasoning: ReasoningConfig?? = .none, maxTokens: Int? = 4096 + ) throws -> GenerationComponents? { + try thinkingBudgetComponents( + budget: budget, enableThinking: thinking, reasoning: reasoning ?? self.reasoning, + tokenizer: tokenizer, parameters: GenerateParameters(maxTokens: maxTokens)) + } + + // MARK: - When a budget applies + + func testNoBudgetMeansNoComponents() throws { + XCTAssertNil(try components(budget: nil)) + } + + func testThinkingOffLeavesTheRequestAlone() throws { + XCTAssertNil(try components(budget: 64, thinking: false)) + } + + func testModelWithoutAReasoningProtocolIsLeftAlone() throws { + XCTAssertNil(try components(budget: 64, reasoning: .some(nil))) + } + + func testBudgetBuildsComponents() throws { + XCTAssertNotNil(try components(budget: 64)) + } + + func testZeroBudgetIsAllowed() throws { + XCTAssertNotNil(try components(budget: 0)) + } + + // MARK: - Budgets the request cannot honour + + func testMaxTokensTooSmallForBudgetAndAnswerIsRejected() { + XCTAssertThrowsError(try components(budget: 1000, maxTokens: 500)) { error in + guard case .insufficientGenerationTokenLimit(_, let actual) = error as? ThinkingBudgetError else { + return XCTFail("unexpected error \(error)") + } + XCTAssertEqual(actual, 500) + } + } + + func testUnlimitedMaxTokensAlwaysFits() throws { + XCTAssertNotNil(try components(budget: 100_000, maxTokens: nil)) + } + + func testNegativeBudgetIsRejected() { + XCTAssertThrowsError(try components(budget: -1)) { error in + XCTAssertEqual(error as? ThinkingBudgetError, .invalidMaximumTokenCount(-1)) + } + } + + // MARK: - 400 body + + func testErrorBodyIsOpenAIShaped() throws { + let body = thinkingBudgetErrorBody(.insufficientGenerationTokenLimit(required: 1100, actual: 500)) + let json = try XCTUnwrap(JSONSerialization.jsonObject(with: Data(body.utf8)) as? [String: Any]) + let error = try XCTUnwrap(json["error"] as? [String: Any]) + XCTAssertEqual(error["type"] as? String, "invalid_request_error") + XCTAssertEqual(error["code"] as? String, "invalid_thinking_budget") + XCTAssertFalse((error["message"] as? String ?? "").isEmpty) + } + + // MARK: - Empty reasoning (budget 0) + + /// A zero budget makes the model emit `` first. The handler keeps `remaining` + /// whether or not any reasoning was extracted, so this contract must hold: no reasoning, + /// and no closing tag left in the content. + func testEmptyReasoningDropsTheClosingTag() { + let (reasoning, content) = extractThinkingBlock(from: "\n\n9 sheep are left.", alreadyOpen: true) + XCTAssertNil(reasoning) + XCTAssertEqual(content, "9 sheep are left.") + } + + func testNoTagsLeavesTheTextUntouched() { + let (reasoning, content) = extractThinkingBlock(from: "plain answer", alreadyOpen: false) + XCTAssertNil(reasoning) + XCTAssertEqual(content, "plain answer") + } + + // MARK: - Flag and request parsing + + func testFlagParsesAndDefaultsToNil() throws { + XCTAssertEqual(try MLXServer.parse(["--model", "m", "--thinking-budget", "2048"]).thinkingBudget, 2048) + XCTAssertNil(try MLXServer.parse(["--model", "m"]).thinkingBudget) + } + + func testNegativeFlagIsRejected() { + XCTAssertThrowsError(try MLXServer.parse(["--model", "m", "--thinking-budget", "-5"])) + } + + func testRequestFieldDecodes() throws { + let json = #"{"model":"m","messages":[{"role":"user","content":"hi"}],"thinking_budget":128}"# + let request = try JSONDecoder().decode(ChatCompletionRequest.self, from: Data(json.utf8)) + XCTAssertEqual(request.thinkingBudget, 128) + } +} + +private struct ByteTokenizer: MLXLMCommon.Tokenizer { + func encode(text: String, addSpecialTokens: Bool) -> [Int] { text.utf8.map(Int.init) } + func decode(tokenIds: [Int], skipSpecialTokens: Bool) -> String { + String(decoding: tokenIds.map(UInt8.init), as: UTF8.self) + } + func convertTokenToId(_ token: String) -> Int? { + guard token.utf8.count == 1 else { return nil } + return token.utf8.first.map(Int.init) + } + func convertIdToToken(_ id: Int) -> String? { + guard (0 ... 255).contains(id) else { return nil } + return String(decoding: [UInt8(id)], as: UTF8.self) + } + var bosToken: String? { nil } + var eosToken: String? { nil } + var unknownToken: String? { nil } + func applyChatTemplate( + messages: [[String: any Sendable]], tools: [[String: any Sendable]]?, + additionalContext: [String: any Sendable]? + ) throws -> [Int] { [] } +}