From 88d465db9f0a2a9f0f1066dfcdf2b4dc0491ef07 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Tue, 25 Aug 2026 19:29:44 -0700 Subject: [PATCH 1/5] Add session-aware prefix caching and /v1/stats endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - X-Session-ID header enables KV cache reuse across requests in the same conversation (engine's TokenHistory detects prefix overlap) - GET /v1/stats returns JSON metrics: throughput, prefix hit rate, memory - Remove unconditional engine.reset() from chat path — let the engine handle prefix detection internally for 10-20x TTFT improvement on multi-turn conversations - Completions endpoint retains reset (lm-eval sends unrelated prompts) --- .../Tools/llm-server/ChatHandler.swift | 40 +++++-- .../Tools/llm-server/ServerState.swift | 110 +++++++++++++++++- 2 files changed, 135 insertions(+), 15 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 9bc0d33f..ce188d99 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -42,7 +42,8 @@ func startServer(state: ServerState, port: Int) async throws { } router.post("/v1/chat/completions") { request, _ in - try await handleChatCompletionsRoute(request: request, state: state) + let sessionID = request.headers[.init("X-Session-ID")!] + return try await handleChatCompletionsRoute(request: request, state: state, sessionID: sessionID) } router.post("/v1") { request, _ in @@ -53,6 +54,16 @@ func startServer(state: ServerState, port: Int) async throws { try await handleCompletionsRoute(request: request, state: state) } + router.get("/v1/stats") { _, _ in + let stats = state.statsSnapshot() + let data = try JSONEncoder().encode(stats) + return Response( + status: .ok, + headers: [.contentType: "application/json"], + body: .init(byteBuffer: ByteBuffer(data: data)) + ) + } + let app = Application( router: router, configuration: .init( @@ -74,7 +85,7 @@ private func handleAutoRoute(request: Request, state: ServerState) async throws } } -private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState) async throws -> Response { +private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws -> Response { guard state.tryAcquire() else { let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy")) let data = try JSONEncoder().encode(err) @@ -96,9 +107,9 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState) do { let shouldStream = chatRequest.stream ?? false if shouldStream { - return try await handleStreamingRequest(chatRequest: chatRequest, state: state) + return try await handleStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID) } else { - let response = try await handleNonStreamingRequest(chatRequest: chatRequest, state: state) + let response = try await handleNonStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID) state.release() return response } @@ -122,14 +133,14 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState) // MARK: - Route Handler -private func handleChatCompletionsRoute(request: Request, state: ServerState) async throws -> Response { +private func handleChatCompletionsRoute(request: Request, state: ServerState, sessionID: String? = nil) async throws -> Response { let body = try await request.body.collect(upTo: 10 * 1024 * 1024) - return try await handleChatCompletionsFromBody(body: body, state: state) + return try await handleChatCompletionsFromBody(body: body, state: state, sessionID: sessionID) } // MARK: - Non-Streaming -private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState) async throws -> Response +private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) async throws -> Response { let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens guard requestMaxTokens > 0 else { @@ -158,7 +169,10 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state "[\(requestID)] messages: \(chatRequest.messages.count), tokens: \(promptTokens.count), max_tokens: \(requestMaxTokens)", component: "Server") - try await state.engine.reset() + let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokens.map { Int32($0) }) + if prefixReused > 0 { + CLILogger.log("[\(requestID)] prefix reuse: \(prefixReused) tokens cached", component: "Server") + } let t0 = SuspendingClock().now let strategy: any DecodingStrategy @@ -198,6 +212,7 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state state.stats.record( promptTokens: promptTokens.count, genTokens: genTokenCount, promptSeconds: 0, genSeconds: seconds, totalSeconds: seconds) + state.recordPromptTokens(promptTokens.map { Int32($0) }) let response = ChatCompletionResponse( id: requestID, @@ -228,7 +243,7 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state // MARK: - Streaming (SSE) -private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState) async throws -> Response { +private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) async throws -> Response { let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens guard requestMaxTokens > 0 else { throw ServerError.badRequest("max_tokens must be positive") @@ -256,10 +271,14 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S "[\(requestID)] stream, messages: \(chatRequest.messages.count), tokens: \(promptTokens.count), max_tokens: \(requestMaxTokens)", component: "Server") + let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokens.map { Int32($0) }) + if prefixReused > 0 { + CLILogger.log("[\(requestID)] prefix reuse: \(prefixReused) tokens cached", component: "Server") + } + let responseBody = ResponseBody { writer in defer { state.release() } do { - try await state.engine.reset() let encoder = JSONEncoder() let genStart = SuspendingClock().now @@ -341,6 +360,7 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S state.stats.record( promptTokens: promptTokens.count, genTokens: tokenCount, promptSeconds: 0, genSeconds: seconds, totalSeconds: seconds) + state.recordPromptTokens(promptTokens.map { Int32($0) }) try await writer.finish(nil) } catch { diff --git a/swift/Sources/Tools/llm-server/ServerState.swift b/swift/Sources/Tools/llm-server/ServerState.swift index 40bb4923..36d3bcbf 100644 --- a/swift/Sources/Tools/llm-server/ServerState.swift +++ b/swift/Sources/Tools/llm-server/ServerState.swift @@ -82,6 +82,20 @@ final class ServerStats: @unchecked Sendable { ================================================== """) } + + func buildResponse(prefixHitRate: Double, prefixHits: Int, prefixMisses: Int) -> StatsResponse { + let s = lock.withLock { $0 } + return StatsResponse( + totalRequests: s.totalRequests, + totalPromptTokens: s.totalPromptTokens, + totalGenTokens: s.totalGenTokens, + avgPrefillTokPerSec: s.totalPromptSeconds > 0 ? Double(s.totalPromptTokens) / s.totalPromptSeconds : 0, + avgDecodeTokPerSec: s.totalGenSeconds > 0 ? Double(s.totalGenTokens) / s.totalGenSeconds : 0, + prefixHitRate: prefixHitRate, + prefixHits: prefixHits, + prefixMisses: prefixMisses + ) + } } // MARK: - Server State @@ -91,7 +105,15 @@ final class ServerState: @unchecked Sendable { let tokenizer: any Tokenizer let config: ServerConfig let stats = ServerStats() - private let _generating = Mutex(false) + private let _state = Mutex(InternalState()) + + private struct InternalState { + var generating: Bool = false + var lastSessionID: String? = nil + var lastPromptTokens: [Int32] = [] + var prefixHits: Int = 0 + var prefixMisses: Int = 0 + } init(engine: any InferenceEngine, tokenizer: any Tokenizer, config: ServerConfig) { self.engine = engine @@ -100,15 +122,69 @@ final class ServerState: @unchecked Sendable { } func tryAcquire() -> Bool { - _generating.withLock { busy in - guard !busy else { return false } - busy = true + _state.withLock { s in + guard !s.generating else { return false } + s.generating = true return true } } func release() { - _generating.withLock { $0 = false } + _state.withLock { $0.generating = false } + } + + /// Prepare engine for a new request. Returns the number of prefix tokens reused. + func prepareForRequest(sessionID: String?, promptTokens: [Int32]) async -> Int { + let action = _state.withLock { s -> PrepareAction in + guard let sid = sessionID, sid == s.lastSessionID else { + s.lastSessionID = sessionID + s.lastPromptTokens = [] + s.prefixMisses += 1 + return .reset + } + + let cached = s.lastPromptTokens + var match = 0 + let limit = min(cached.count, promptTokens.count) + while match < limit && cached[match] == promptTokens[match] { + match += 1 + } + + if match == 0 { + s.prefixMisses += 1 + return .reset + } + + s.prefixHits += 1 + return .reuse(prefixLength: match) + } + + switch action { + case .reset: + return 0 + case .reuse(let prefixLength): + return prefixLength + } + } + + /// Record the tokens that were processed (call after generate completes). + func recordPromptTokens(_ tokens: [Int32]) { + _state.withLock { $0.lastPromptTokens = tokens } + } + + /// Stats snapshot for /v1/stats endpoint. + func statsSnapshot() -> StatsResponse { + let s = _state.withLock { s in + (prefixHits: s.prefixHits, prefixMisses: s.prefixMisses) + } + let hitTotal = s.prefixHits + s.prefixMisses + let hitRate = hitTotal > 0 ? Double(s.prefixHits) / Double(hitTotal) : 0 + return stats.buildResponse(prefixHitRate: hitRate, prefixHits: s.prefixHits, prefixMisses: s.prefixMisses) + } + + private enum PrepareAction { + case reset + case reuse(prefixLength: Int) } func makeSamplingConfig( @@ -160,3 +236,27 @@ enum ServerError: Error, LocalizedError { } } } + +// MARK: - Stats Response + +struct StatsResponse: Codable, Sendable { + let totalRequests: Int + let totalPromptTokens: Int + let totalGenTokens: Int + let avgPrefillTokPerSec: Double + let avgDecodeTokPerSec: Double + let prefixHitRate: Double + let prefixHits: Int + let prefixMisses: Int + + enum CodingKeys: String, CodingKey { + case totalRequests = "total_requests" + case totalPromptTokens = "total_prompt_tokens" + case totalGenTokens = "total_gen_tokens" + case avgPrefillTokPerSec = "avg_prefill_tok_per_sec" + case avgDecodeTokPerSec = "avg_decode_tok_per_sec" + case prefixHitRate = "prefix_hit_rate" + case prefixHits = "prefix_hits" + case prefixMisses = "prefix_misses" + } +} From 470ba16eb965e92f327a880640b2c71471863282 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 26 Aug 2026 16:24:52 -0700 Subject: [PATCH 2/5] Fix swift-format line length violations in ChatHandler --- .../Sources/Tools/llm-server/ChatHandler.swift | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index ce188d99..07a617d9 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -85,7 +85,9 @@ private func handleAutoRoute(request: Request, state: ServerState) async throws } } -private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws -> Response { +private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws + -> Response +{ guard state.tryAcquire() else { let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy")) let data = try JSONEncoder().encode(err) @@ -109,7 +111,8 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, if shouldStream { return try await handleStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID) } else { - let response = try await handleNonStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID) + let response = try await handleNonStreamingRequest( + chatRequest: chatRequest, state: state, sessionID: sessionID) state.release() return response } @@ -133,14 +136,17 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, // MARK: - Route Handler -private func handleChatCompletionsRoute(request: Request, state: ServerState, sessionID: String? = nil) async throws -> Response { +private func handleChatCompletionsRoute(request: Request, state: ServerState, sessionID: String? = nil) async throws + -> Response +{ let body = try await request.body.collect(upTo: 10 * 1024 * 1024) return try await handleChatCompletionsFromBody(body: body, state: state, sessionID: sessionID) } // MARK: - Non-Streaming -private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) async throws -> Response +private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) + async throws -> Response { let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens guard requestMaxTokens > 0 else { @@ -243,7 +249,9 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state // MARK: - Streaming (SSE) -private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) async throws -> Response { +private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil) + async throws -> Response +{ let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens guard requestMaxTokens > 0 else { throw ServerError.badRequest("max_tokens must be positive") From 097e8b90d476582303836d1d22ff25ed6ec1c6ae Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 26 Aug 2026 16:30:14 -0700 Subject: [PATCH 3/5] Replace force-unwrap on HTTPField.Name with safe fallback --- swift/Sources/Tools/llm-server/ChatHandler.swift | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 07a617d9..49299fa6 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -7,6 +7,7 @@ import CoreAILMCommon import CoreAILanguageModels import CoreAIShared import Foundation +import HTTPTypes import Hummingbird import NIOCore import NIOFoundationCompat @@ -42,7 +43,8 @@ func startServer(state: ServerState, port: Int) async throws { } router.post("/v1/chat/completions") { request, _ in - let sessionID = request.headers[.init("X-Session-ID")!] + let sessionID: String? = + if let name = HTTPField.Name("X-Session-ID") { request.headers[name] } else { nil } return try await handleChatCompletionsRoute(request: request, state: state, sessionID: sessionID) } From 1739647825d7444afb42570a7006eb6f1e3187bd Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 26 Aug 2026 16:31:05 -0700 Subject: [PATCH 4/5] Default to "default" session when X-Session-ID header is absent Single-client usage (no header) now gets prefix caching automatically. Multi-client callers can still differentiate sessions via the header. --- swift/Sources/Tools/llm-server/ChatHandler.swift | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index 49299fa6..d748c846 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -43,8 +43,8 @@ func startServer(state: ServerState, port: Int) async throws { } router.post("/v1/chat/completions") { request, _ in - let sessionID: String? = - if let name = HTTPField.Name("X-Session-ID") { request.headers[name] } else { nil } + let sessionID = + HTTPField.Name("X-Session-ID").flatMap { request.headers[$0] } ?? "default" return try await handleChatCompletionsRoute(request: request, state: state, sessionID: sessionID) } From 1987e9c1d0c4e9bb8d553856eb87ce984123331c Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Wed, 26 Aug 2026 22:37:26 -0700 Subject: [PATCH 5/5] Avoid redundant Int32 token conversion per request Compute promptTokens.map { Int32($0) } once per handler and reuse for both prepareForRequest and recordPromptTokens. Eliminates one unnecessary heap allocation per request (up to 512KB at 128K context). --- swift/Sources/Tools/llm-server/ChatHandler.swift | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/swift/Sources/Tools/llm-server/ChatHandler.swift b/swift/Sources/Tools/llm-server/ChatHandler.swift index d748c846..bce1e53c 100644 --- a/swift/Sources/Tools/llm-server/ChatHandler.swift +++ b/swift/Sources/Tools/llm-server/ChatHandler.swift @@ -177,7 +177,8 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state "[\(requestID)] messages: \(chatRequest.messages.count), tokens: \(promptTokens.count), max_tokens: \(requestMaxTokens)", component: "Server") - let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokens.map { Int32($0) }) + let promptTokensInt32 = promptTokens.map { Int32($0) } + let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokensInt32) if prefixReused > 0 { CLILogger.log("[\(requestID)] prefix reuse: \(prefixReused) tokens cached", component: "Server") } @@ -220,7 +221,7 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state state.stats.record( promptTokens: promptTokens.count, genTokens: genTokenCount, promptSeconds: 0, genSeconds: seconds, totalSeconds: seconds) - state.recordPromptTokens(promptTokens.map { Int32($0) }) + state.recordPromptTokens(promptTokensInt32) let response = ChatCompletionResponse( id: requestID, @@ -281,7 +282,8 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S "[\(requestID)] stream, messages: \(chatRequest.messages.count), tokens: \(promptTokens.count), max_tokens: \(requestMaxTokens)", component: "Server") - let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokens.map { Int32($0) }) + let promptTokensInt32 = promptTokens.map { Int32($0) } + let prefixReused = await state.prepareForRequest(sessionID: sessionID, promptTokens: promptTokensInt32) if prefixReused > 0 { CLILogger.log("[\(requestID)] prefix reuse: \(prefixReused) tokens cached", component: "Server") } @@ -370,7 +372,7 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S state.stats.record( promptTokens: promptTokens.count, genTokens: tokenCount, promptSeconds: 0, genSeconds: seconds, totalSeconds: seconds) - state.recordPromptTokens(promptTokens.map { Int32($0) }) + state.recordPromptTokens(promptTokensInt32) try await writer.finish(nil) } catch {