Spaces:
Running
Running
| import { | |
| Wllama, | |
| type ChatCompletionChunk, | |
| type ChatCompletionMessage, | |
| type ChatCompletionResponse, | |
| type ChatCompletionTool, | |
| type ChatCompletionUsage, | |
| type ResultTimings, | |
| } from '../../vendor/wllama-bonsai/esm/index.js'; | |
| import runtimeSource from '../../vendor/wllama-bonsai/SOURCE.json'; | |
| import { DEFAULT_MAX_TOKENS } from '../lib/runtime-defaults'; | |
| import { | |
| assertBackendPolicy, | |
| estimateStorage, | |
| evaluateModelContextPolicy, | |
| evaluateModelGate, | |
| inspectWebGpuAdapter, | |
| persistStorage, | |
| } from './device-gate'; | |
| import { verifyBlobSha256 } from './blob-integrity'; | |
| import { | |
| EngineRuntimeError, | |
| throwIfAborted, | |
| type WebGpuDeviceLostDetails, | |
| } from './errors'; | |
| import { | |
| findManifestModel, | |
| loadModelManifestV2, | |
| orderedShardUrls, | |
| parseModelManifestV2, | |
| type ManifestModelV2, | |
| type ModelManifestV2, | |
| } from './manifest'; | |
| import { NativeLogCollector, type BackendReport } from './native-log'; | |
| import { ToolCallAccumulator } from './tool-call-accumulator'; | |
| import type { | |
| BenchmarkFlashMode, | |
| BenchmarkKvCacheType, | |
| BenchmarkWasmFlavor, | |
| EngineCapabilities, | |
| EngineEvent, | |
| EngineSampledTokenTraceEntry, | |
| EngineShardProgress, | |
| GenerateParams, | |
| GenerateResult, | |
| LoadModelParams, | |
| LoadModelResult, | |
| RuntimeBackend, | |
| ScoreSequenceParams, | |
| ScoreSequenceResult, | |
| ShardDownloadFailureDetails, | |
| StorageEstimate, | |
| } from './protocol'; | |
| const WLLAMA_REVISION = '912c18b75d4358c1405a64646b8dbe43a205943b'; | |
| const LLAMA_CPP_REVISION = '00fa7cb284cbf133fc426733bd64238a3588a33e'; | |
| const WLLAMA_WASM_PATH = '/wasm/wllama.wasm'; | |
| const WLLAMA_COMPAT_WASM_PATH = '/wasm/wllama-compat.wasm'; | |
| const WLLAMA_COMPAT_WORKER_PATH = '/wasm/wllama-compat.js'; | |
| const MAX_GLUE_INT = 2_147_483_647; | |
| const DEVICE_LOST_EXIT_GRACE_MS = 100; | |
| const TOKEN_TRACE_TOP_LOGPROBS = 5; | |
| const MIN_LIVE_RATE_WINDOW_MS = 50; | |
| const COMPAT_STABLE_BATCH_SIZE = 32; | |
| const COMPAT_STABLE_MICRO_BATCH_SIZE = 16; | |
| const STATE_DRIFT_REFERENCE_TOKENS = 1_024; | |
| const STATE_DRIFT_PROMPT_TOKENS = 38; | |
| const TEACHER_FORCE_LOGIT_BIAS = 1_000; | |
| export interface ResolvedLoadTuning { | |
| scope: 'release-defaults' | 'benchmark'; | |
| nBatch: number | null; | |
| nUbatch: number | null; | |
| flashMode: BenchmarkFlashMode; | |
| kvCacheType: BenchmarkKvCacheType; | |
| wasmFlavor: BenchmarkWasmFlavor; | |
| } | |
| export function resolveRuntimeBatchShape( | |
| tuning: Pick<ResolvedLoadTuning, 'nBatch' | 'nUbatch'>, | |
| wasmFlavor: 'jspi' | 'compat', | |
| ): { nBatch: number | undefined; nUbatch: number | undefined } { | |
| if (wasmFlavor !== 'compat') { | |
| return { | |
| nBatch: tuning.nBatch ?? undefined, | |
| nUbatch: tuning.nUbatch ?? undefined, | |
| }; | |
| } | |
| const nBatch = tuning.nBatch ?? COMPAT_STABLE_BATCH_SIZE; | |
| return { | |
| nBatch, | |
| nUbatch: tuning.nUbatch ?? Math.min(COMPAT_STABLE_MICRO_BATCH_SIZE, nBatch), | |
| }; | |
| } | |
| function validateBatchValue(value: unknown, label: 'n_batch' | 'n_ubatch'): number | null { | |
| if (value === undefined) return null; | |
| if (!Number.isSafeInteger(value) || (value as number) <= 0 || (value as number) > MAX_GLUE_INT) { | |
| throw new EngineRuntimeError( | |
| label === 'n_batch' ? 'INVALID_BATCH_SIZE' : 'INVALID_UBATCH_SIZE', | |
| `${label} must be a positive signed 32-bit integer (maximum ${MAX_GLUE_INT}).`, | |
| ); | |
| } | |
| return value as number; | |
| } | |
| export function resolveLoadTuning( | |
| input: unknown, | |
| backend: LoadModelParams['backend'], | |
| ): ResolvedLoadTuning { | |
| if (input === undefined) { | |
| return { | |
| scope: 'release-defaults', | |
| nBatch: null, | |
| nUbatch: null, | |
| flashMode: 'off', | |
| kvCacheType: 'f16', | |
| wasmFlavor: 'auto', | |
| }; | |
| } | |
| if (typeof input !== 'object' || input === null || Array.isArray(input)) { | |
| throw new EngineRuntimeError('INVALID_BENCHMARK_TUNING', 'Benchmark tuning must be an object.'); | |
| } | |
| const tuning = input as Record<string, unknown>; | |
| const flashMode = tuning.flashMode; | |
| const kvCacheType = tuning.kvCacheType; | |
| const wasmFlavor = tuning.wasmFlavor; | |
| if (flashMode !== 'off' && flashMode !== 'auto') { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Benchmark flash mode must be off or auto.', | |
| ); | |
| } | |
| if (kvCacheType !== 'f16' && kvCacheType !== 'q8_0' && kvCacheType !== 'q4_0') { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Benchmark KV cache type must be f16, q8_0, or q4_0.', | |
| ); | |
| } | |
| if (wasmFlavor !== 'auto' && wasmFlavor !== 'jspi' && wasmFlavor !== 'compat') { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Benchmark WASM flavor must be auto, jspi, or compat.', | |
| ); | |
| } | |
| if (kvCacheType !== 'f16' && flashMode !== 'auto') { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Quantized V cache requires Flash Attention in auto mode.', | |
| ); | |
| } | |
| if ((flashMode === 'auto' || kvCacheType !== 'f16') && backend !== 'webgpu') { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Flash Attention and quantized KV experiments require the explicit WebGPU backend.', | |
| ); | |
| } | |
| const nBatch = validateBatchValue(tuning.nBatch, 'n_batch'); | |
| const nUbatch = validateBatchValue(tuning.nUbatch, 'n_ubatch'); | |
| if (nBatch !== null && nUbatch !== null && nUbatch > nBatch) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_BENCHMARK_TUNING', | |
| 'Benchmark n_ubatch must not exceed n_batch.', | |
| ); | |
| } | |
| return { | |
| scope: 'benchmark', | |
| nBatch, | |
| nUbatch, | |
| flashMode, | |
| kvCacheType, | |
| wasmFlavor, | |
| }; | |
| } | |
| function runtimeArtifact(path: string): { bytes: number; sha256: string } { | |
| const artifact = runtimeSource.files.find((file) => file.path === path); | |
| if (!artifact) throw new Error(`Runtime provenance is missing ${path}`); | |
| return artifact; | |
| } | |
| function needsCompatWasm(): boolean { | |
| if (!(WebAssembly as typeof WebAssembly & { Suspending?: unknown }).Suspending) return true; | |
| try { | |
| new WebAssembly.Memory( | |
| { address: 'i64', initial: 1n } as unknown as WebAssembly.MemoryDescriptor, | |
| ); | |
| return false; | |
| } catch { | |
| return true; | |
| } | |
| } | |
| export function selectWasmFlavor( | |
| requested: BenchmarkWasmFlavor, | |
| compatRequired = needsCompatWasm(), | |
| ): 'jspi' | 'compat' { | |
| if (requested === 'compat') return 'compat'; | |
| if (requested === 'jspi') { | |
| if (compatRequired) { | |
| throw new EngineRuntimeError( | |
| 'BENCHMARK_WASM_FLAVOR_UNAVAILABLE', | |
| 'JSPI was requested, but this browser requires the compatibility runtime.', | |
| ); | |
| } | |
| return 'jspi'; | |
| } | |
| return compatRequired ? 'compat' : 'jspi'; | |
| } | |
| type EventSink = (event: EngineEvent) => void; | |
| interface LoadedModelState { | |
| manifest: ModelManifestV2; | |
| model: ManifestModelV2; | |
| backend: RuntimeBackend; | |
| tuningScope: ResolvedLoadTuning['scope']; | |
| contextSize: number; | |
| batchSize: number; | |
| microBatchSize: number; | |
| vocabularySize: number; | |
| } | |
| interface CachedShardState { | |
| blobs: Array<Blob | null>; | |
| cachedBytes: number; | |
| shards?: EngineShardProgress[]; | |
| } | |
| function initialShardProgress( | |
| model: ManifestModelV2, | |
| blobs: readonly (Blob | null)[] = [], | |
| ): EngineShardProgress[] { | |
| return model.files.map((file, index) => ({ | |
| index, | |
| path: file.path, | |
| loadedBytes: blobs[index] ? file.bytes : 0, | |
| verifiedBytes: blobs[index] ? file.bytes : 0, | |
| totalBytes: file.bytes, | |
| state: blobs[index] ? 'cached' : 'queued', | |
| })); | |
| } | |
| function copyShardProgress(shards: readonly EngineShardProgress[]): EngineShardProgress[] { | |
| return shards.map((shard) => ({ ...shard })); | |
| } | |
| function totalShardProgress(shards: readonly EngineShardProgress[]): number { | |
| return shards.reduce((sum, shard) => sum + shard.loadedBytes, 0); | |
| } | |
| function totalVerifiedShardProgress(shards: readonly EngineShardProgress[]): number { | |
| return shards.reduce((sum, shard) => sum + shard.verifiedBytes, 0); | |
| } | |
| function errorText(error: unknown): string { | |
| return error instanceof Error ? error.message : String(error); | |
| } | |
| function isAbortFailure(error: unknown, signal: AbortSignal): boolean { | |
| return signal.aborted || (error instanceof DOMException && error.name === 'AbortError') | |
| || (error instanceof Error && error.name === 'AbortError'); | |
| } | |
| function classifyShardFailure( | |
| error: unknown, | |
| signal: AbortSignal, | |
| ): ShardDownloadFailureDetails['failure'] { | |
| if (error instanceof EngineRuntimeError | |
| && (error.code === 'SHARD_SIZE_MISMATCH' || error.code === 'SHARD_HASH_MISMATCH')) { | |
| return 'verification'; | |
| } | |
| return isAbortFailure(error, signal) ? 'abort' : 'network'; | |
| } | |
| function shardFailure( | |
| error: unknown, | |
| failure: ShardDownloadFailureDetails['failure'], | |
| shardIndex: number, | |
| shardCount: number, | |
| shardPath: string, | |
| partialDeleted: boolean, | |
| cleanupError?: unknown, | |
| ): EngineRuntimeError { | |
| const verification = failure === 'verification'; | |
| const aborted = failure === 'abort'; | |
| const causeCode = error instanceof EngineRuntimeError | |
| ? error.code | |
| : aborted ? 'ABORTED' : error instanceof Error ? error.name : 'UNKNOWN'; | |
| const details: ShardDownloadFailureDetails = { | |
| kind: 'shard-download-failure', | |
| failure, | |
| causeCode, | |
| shardIndex, | |
| shardCount, | |
| shardPath, | |
| retryFromByteZero: true, | |
| partialDeleted, | |
| }; | |
| const code = verification | |
| ? causeCode | |
| : aborted ? 'SHARD_DOWNLOAD_ABORTED' : 'SHARD_DOWNLOAD_FAILED'; | |
| const cleanupSuffix = cleanupError === undefined | |
| ? '' | |
| : ` Partial cache cleanup also failed: ${errorText(cleanupError)}`; | |
| const causeMessage = errorText(error).replace(/\.+$/, ''); | |
| return new EngineRuntimeError( | |
| code, | |
| `Shard ${shardIndex + 1}/${shardCount} (${shardPath}) failed: ${causeMessage}.${cleanupSuffix}`, | |
| details, | |
| ); | |
| } | |
| function absoluteAssetUrl(path: string): string { | |
| return new URL(path, globalThis.location.origin).href; | |
| } | |
| function installWorkerDocumentShim(): void { | |
| if (!('document' in globalThis)) { | |
| Object.defineProperty(globalThis, 'document', { | |
| configurable: true, | |
| value: { baseURI: globalThis.location.href }, | |
| }); | |
| } | |
| } | |
| function emitProgress( | |
| sink: EventSink, | |
| requestId: string, | |
| values: Omit< | |
| Extract<EngineEvent, { event: 'progress' }>, | |
| 'type' | 'requestId' | 'event' | 'nativeStage' | |
| > & { nativeStage?: string | null }, | |
| ): void { | |
| sink({ | |
| type: 'event', | |
| requestId, | |
| event: 'progress', | |
| ...values, | |
| nativeStage: values.nativeStage ?? null, | |
| }); | |
| } | |
| interface PromptProgressSample { | |
| total: number; | |
| cache: number; | |
| processed: number; | |
| time_ms: number; | |
| } | |
| type LiveCompletionChunk = ChatCompletionChunk & { | |
| prompt_progress?: PromptProgressSample; | |
| }; | |
| function emitGenerationProgress( | |
| sink: EventSink, | |
| requestId: string, | |
| values: Omit<Extract<EngineEvent, { event: 'generation' }>, 'type' | 'requestId' | 'event'>, | |
| ): void { | |
| sink({ | |
| type: 'event', | |
| requestId, | |
| event: 'generation', | |
| ...values, | |
| }); | |
| } | |
| function tokenCount(value: number | null | undefined): number { | |
| return typeof value === 'number' && Number.isFinite(value) && value >= 0 | |
| ? Math.floor(value) | |
| : 0; | |
| } | |
| function mapUsage( | |
| usage: ChatCompletionUsage | null | undefined, | |
| timings: ResultTimings | undefined, | |
| ): GenerateResult['usage'] { | |
| if (!usage && !timings) { | |
| return null; | |
| } | |
| const promptTokens = Math.max( | |
| tokenCount(usage?.prompt_tokens), | |
| tokenCount(timings?.prompt_n), | |
| ); | |
| const completionTokens = Math.max( | |
| tokenCount(usage?.completion_tokens), | |
| tokenCount(timings?.predicted_n), | |
| ); | |
| return { | |
| promptTokens, | |
| completionTokens, | |
| totalTokens: Math.max(tokenCount(usage?.total_tokens), promptTokens + completionTokens), | |
| }; | |
| } | |
| function mapTimings(timings: ResultTimings | undefined): GenerateResult['timings'] { | |
| if (!timings) { | |
| return null; | |
| } | |
| return { | |
| promptTokensPerSecond: timings.prompt_per_second, | |
| predictedTokensPerSecond: timings.predicted_per_second, | |
| }; | |
| } | |
| function traceTokenId(value: unknown, index: number, field: string): number { | |
| if (!Number.isSafeInteger(value) || (value as number) < 0 || (value as number) > 0xffff_ffff) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_TOKEN_ID_TRACE', | |
| 'The native runtime returned a missing or invalid token id in the sampled-token trace.', | |
| { index, field, id: value }, | |
| ); | |
| } | |
| return value as number; | |
| } | |
| function traceLogprob(value: unknown, index: number, field: string): number { | |
| if (typeof value !== 'number' || !Number.isFinite(value)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_TOKEN_LOGPROB_TRACE', | |
| 'The native runtime returned a missing or invalid logprob in the sampled-token trace.', | |
| { index, field, logprob: value }, | |
| ); | |
| } | |
| return value; | |
| } | |
| function appendSampledTokenTrace( | |
| target: EngineSampledTokenTraceEntry[], | |
| completion: ChatCompletionChunk | ChatCompletionResponse, | |
| ): void { | |
| const entries = completion.choices[0]?.logprobs?.content; | |
| if (!entries) return; | |
| for (const entry of entries) { | |
| const index = target.length; | |
| const selected = { | |
| id: traceTokenId(entry.id, index, 'selected.id'), | |
| logprob: traceLogprob(entry.logprob, index, 'selected.logprob'), | |
| }; | |
| if (!Array.isArray(entry.top_logprobs) || entry.top_logprobs.length !== TOKEN_TRACE_TOP_LOGPROBS) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_TOKEN_LOGPROB_TRACE', | |
| `The native runtime must return exactly ${TOKEN_TRACE_TOP_LOGPROBS} top logprob candidates per sampled token.`, | |
| { | |
| index, | |
| expected: TOKEN_TRACE_TOP_LOGPROBS, | |
| observed: Array.isArray(entry.top_logprobs) ? entry.top_logprobs.length : null, | |
| }, | |
| ); | |
| } | |
| const seenIds = new Set<number>(); | |
| const topCandidates = entry.top_logprobs.map((candidate, candidateIndex) => { | |
| const id = traceTokenId(candidate.id, index, `topCandidates[${candidateIndex}].id`); | |
| if (seenIds.has(id)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_TOKEN_LOGPROB_TRACE', | |
| 'The native runtime returned duplicate ids in a top-logprob candidate list.', | |
| { index, candidateIndex, id }, | |
| ); | |
| } | |
| seenIds.add(id); | |
| return { | |
| id, | |
| logprob: traceLogprob( | |
| candidate.logprob, | |
| index, | |
| `topCandidates[${candidateIndex}].logprob`, | |
| ), | |
| }; | |
| }).sort((left, right) => right.logprob - left.logprob || left.id - right.id); | |
| const selectedCandidate = topCandidates.find((candidate) => candidate.id === selected.id); | |
| if (!selectedCandidate || selectedCandidate.logprob !== selected.logprob) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_TOKEN_LOGPROB_TRACE', | |
| 'The sampled token must appear exactly once in its top-logprob candidates with the same logprob.', | |
| { index, selected, selectedCandidate: selectedCandidate ?? null }, | |
| ); | |
| } | |
| target.push({ selected, topCandidates }); | |
| } | |
| } | |
| function tokenTraceAccounting( | |
| tokenTrace: EngineSampledTokenTraceEntry[] | null, | |
| usage: GenerateResult['usage'], | |
| ): GenerateResult['tokenTraceAccounting'] { | |
| if (tokenTrace === null) return null; | |
| const usageCompletionTokens = usage?.completionTokens ?? null; | |
| return { | |
| usageCompletionTokens, | |
| tracedTokens: tokenTrace.length, | |
| delta: usageCompletionTokens === null ? null : tokenTrace.length - usageCompletionTokens, | |
| }; | |
| } | |
| function scoreRecord(value: unknown, index: number, field: string): Record<string, unknown> { | |
| if (typeof value !== 'object' || value === null || Array.isArray(value)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} has an invalid ${field}.`, | |
| { index, field }, | |
| ); | |
| } | |
| return value as Record<string, unknown>; | |
| } | |
| function scoreTokenId( | |
| value: unknown, | |
| index: number, | |
| field: string, | |
| vocabularySize: number, | |
| ): number { | |
| if (!Number.isSafeInteger(value) || (value as number) < 0 || (value as number) >= vocabularySize) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} has an invalid ${field}.`, | |
| { index, field, id: value, vocabularySize }, | |
| ); | |
| } | |
| return value as number; | |
| } | |
| function scoreLogprob(value: unknown, index: number, field: string): number { | |
| if (typeof value !== 'number' || !Number.isFinite(value)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} has an invalid ${field}.`, | |
| { index, field, logprob: value }, | |
| ); | |
| } | |
| return value; | |
| } | |
| function parseTeacherForcedResponse( | |
| response: unknown, | |
| index: number, | |
| referenceTokenId: number, | |
| vocabularySize: number, | |
| ): ScoreSequenceResult['entries'][number] { | |
| const root = scoreRecord(response, index, 'root'); | |
| if (!Array.isArray(root.choices) || root.choices.length !== 1) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} must contain exactly one completion choice.`, | |
| { index, observed: Array.isArray(root.choices) ? root.choices.length : null }, | |
| ); | |
| } | |
| const choice = scoreRecord(root.choices[0], index, 'choices[0]'); | |
| const logprobs = scoreRecord(choice.logprobs, index, 'choices[0].logprobs'); | |
| if (!Array.isArray(logprobs.content) || logprobs.content.length !== 1) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} must contain one selected-token logprob entry.`, | |
| { index, observed: Array.isArray(logprobs.content) ? logprobs.content.length : null }, | |
| ); | |
| } | |
| const selectedEntry = scoreRecord(logprobs.content[0], index, 'logprobs.content[0]'); | |
| const selectedReference = { | |
| id: scoreTokenId(selectedEntry.id, index, 'selectedReference.id', vocabularySize), | |
| logprob: scoreLogprob(selectedEntry.logprob, index, 'selectedReference.logprob'), | |
| }; | |
| if (selectedReference.id !== referenceTokenId) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher forcing selected token ${selectedReference.id} instead of reference token ${referenceTokenId} at position ${index + 1}.`, | |
| { index, selectedTokenId: selectedReference.id, referenceTokenId }, | |
| ); | |
| } | |
| const rawTopLogprobs = selectedEntry.top_logprobs; | |
| if (!Array.isArray(rawTopLogprobs) | |
| || rawTopLogprobs.length !== TOKEN_TRACE_TOP_LOGPROBS) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} must contain exactly ${TOKEN_TRACE_TOP_LOGPROBS} natural top candidates.`, | |
| { | |
| index, | |
| observed: Array.isArray(rawTopLogprobs) | |
| ? rawTopLogprobs.length | |
| : null, | |
| }, | |
| ); | |
| } | |
| const seenIds = new Set<number>(); | |
| const topCandidates = rawTopLogprobs.map((rawCandidate, candidateIndex) => { | |
| const candidate = scoreRecord( | |
| rawCandidate, | |
| index, | |
| `topCandidates[${candidateIndex}]`, | |
| ); | |
| const parsed = { | |
| id: scoreTokenId( | |
| candidate.id, | |
| index, | |
| `topCandidates[${candidateIndex}].id`, | |
| vocabularySize, | |
| ), | |
| logprob: scoreLogprob( | |
| candidate.logprob, | |
| index, | |
| `topCandidates[${candidateIndex}].logprob`, | |
| ), | |
| }; | |
| if (seenIds.has(parsed.id)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} contains duplicate natural candidate id ${parsed.id}.`, | |
| { index, candidateIndex, id: parsed.id }, | |
| ); | |
| } | |
| seenIds.add(parsed.id); | |
| return parsed; | |
| }).sort((left, right) => right.logprob - left.logprob || left.id - right.id); | |
| const referenceRankInTopCandidatesZeroBased = topCandidates.findIndex( | |
| (candidate) => candidate.id === referenceTokenId, | |
| ); | |
| if ( | |
| referenceRankInTopCandidatesZeroBased !== -1 | |
| && topCandidates[referenceRankInTopCandidatesZeroBased]?.logprob !== selectedReference.logprob | |
| ) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| `Teacher-forced response ${index} reports inconsistent reference logprobs.`, | |
| { index, selectedReference, referenceRankInTopCandidatesZeroBased }, | |
| ); | |
| } | |
| const naturalTop1 = topCandidates[0]!; | |
| const runnerUp = topCandidates[1]!; | |
| return { | |
| index, | |
| selectedReference, | |
| naturalTop1: { ...naturalTop1 }, | |
| topCandidates, | |
| referenceRankInTopCandidatesZeroBased: referenceRankInTopCandidatesZeroBased === -1 | |
| ? null | |
| : referenceRankInTopCandidatesZeroBased, | |
| top1Top2Margin: naturalTop1.logprob - runnerUp.logprob, | |
| }; | |
| } | |
| function asReasoningDelta(chunk: ChatCompletionChunk): string | undefined { | |
| const choice = chunk.choices[0] as unknown as Record<string, unknown> | undefined; | |
| const delta = choice?.delta; | |
| if (typeof delta !== 'object' || delta === null) { | |
| return undefined; | |
| } | |
| const reasoning = (delta as Record<string, unknown>).reasoning_content; | |
| return typeof reasoning === 'string' && reasoning.length > 0 ? reasoning : undefined; | |
| } | |
| function asReasoningContent(message: ChatCompletionResponse['choices'][number]['message']): string { | |
| const reasoning = (message as unknown as Record<string, unknown>).reasoning_content; | |
| return typeof reasoning === 'string' ? reasoning : ''; | |
| } | |
| export class BrowserEngineRuntime { | |
| private wllama: Wllama | null = null; | |
| private wllamaFlavor: 'jspi' | 'compat' | null = null; | |
| private compatWorkerCode: string | null = null; | |
| private loaded: LoadedModelState | null = null; | |
| private readonly nativeLog = new NativeLogCollector(); | |
| private async raceWebGpuDeviceLoss<Result>( | |
| stage: WebGpuDeviceLostDetails['stage'], | |
| backend: RuntimeBackend, | |
| model: ManifestModelV2, | |
| operation: Promise<Result>, | |
| ): Promise<Result> { | |
| if (backend !== 'webgpu') return operation; | |
| const watch = this.nativeLog.watchDeviceLost(); | |
| try { | |
| return await Promise.race([ | |
| operation, | |
| watch.promise.then(async () => { | |
| await this.assertWebGpuAlive(stage, backend, model); | |
| throw new Error('WebGPU device-loss handler returned unexpectedly.'); | |
| }), | |
| ]); | |
| } finally { | |
| watch.dispose(); | |
| } | |
| } | |
| private async assertWebGpuAlive( | |
| stage: WebGpuDeviceLostDetails['stage'], | |
| backend: RuntimeBackend, | |
| model: ManifestModelV2, | |
| ): Promise<void> { | |
| const signal = this.nativeLog.deviceLostSignal(); | |
| if (backend !== 'webgpu' || !signal) { | |
| return; | |
| } | |
| const invalidatedWllama = this.wllama; | |
| this.loaded = null; | |
| this.wllama = null; | |
| this.wllamaFlavor = null; | |
| if (invalidatedWllama) { | |
| let timeout: ReturnType<typeof setTimeout> | undefined; | |
| try { | |
| await Promise.race([ | |
| invalidatedWllama.exit().catch(() => undefined), | |
| new Promise<void>((resolve) => { | |
| timeout = setTimeout(resolve, DEVICE_LOST_EXIT_GRACE_MS); | |
| }), | |
| ]); | |
| } catch { | |
| // The worker is replaced after this typed failure; exit is best-effort only. | |
| } finally { | |
| if (timeout !== undefined) clearTimeout(timeout); | |
| } | |
| } | |
| const fallbackSteer = model.cpuFallback | |
| ? ' If this repeats, select CPU-WASM before retrying.' | |
| : ' This model tier requires WebGPU; retry after the browser obtains a fresh GPU device.'; | |
| const details: WebGpuDeviceLostDetails = { | |
| recoverable: true, | |
| nextAction: 'reload-model', | |
| stage, | |
| modelId: model.id, | |
| cpuFallbackAvailable: model.cpuFallback, | |
| signal, | |
| }; | |
| throw new EngineRuntimeError( | |
| 'WEBGPU_DEVICE_LOST', | |
| `The WebGPU device was lost during ${stage}. The loaded model was invalidated. Run the action again to reload it.${fallbackSteer}`, | |
| details, | |
| ); | |
| } | |
| private async ensureWllama(flavor = selectWasmFlavor('auto')): Promise<Wllama> { | |
| if (this.wllama && this.wllamaFlavor === flavor) { | |
| return this.wllama; | |
| } | |
| if (this.wllama) { | |
| await this.wllama.exit().catch(() => undefined); | |
| this.wllama = null; | |
| this.loaded = null; | |
| this.wllamaFlavor = null; | |
| } | |
| installWorkerDocumentShim(); | |
| if (this.compatWorkerCode === null) { | |
| const response = await fetch(absoluteAssetUrl(WLLAMA_COMPAT_WORKER_PATH)); | |
| if (!response.ok) { | |
| throw new EngineRuntimeError( | |
| 'COMPAT_ASSET_LOAD_FAILED', | |
| `Failed to load local wllama compat worker: HTTP ${response.status}`, | |
| ); | |
| } | |
| this.compatWorkerCode = await response.text(); | |
| } | |
| const wllama = new Wllama( | |
| { default: absoluteAssetUrl(WLLAMA_WASM_PATH) }, | |
| { | |
| parallelDownloads: 3, | |
| forceCompat: flavor === 'compat', | |
| logger: { | |
| debug: (...values: unknown[]) => this.nativeLog.append(values), | |
| log: (...values: unknown[]) => this.nativeLog.append(values), | |
| warn: (...values: unknown[]) => this.nativeLog.append(values), | |
| error: (...values: unknown[]) => this.nativeLog.append(values), | |
| }, | |
| }, | |
| ); | |
| wllama.setCompat({ | |
| wasm: absoluteAssetUrl(WLLAMA_COMPAT_WASM_PATH), | |
| worker: { code: this.compatWorkerCode }, | |
| }, 'firefox_safari'); | |
| this.wllama = wllama; | |
| this.wllamaFlavor = flavor; | |
| return wllama; | |
| } | |
| async capabilities(): Promise<EngineCapabilities> { | |
| const [webgpu, storage] = await Promise.all([ | |
| inspectWebGpuAdapter(), | |
| estimateStorage(), | |
| ]); | |
| const wasmFlavor = needsCompatWasm() ? 'compat' : 'jspi'; | |
| const wasmArtifact = runtimeArtifact( | |
| wasmFlavor === 'compat' ? 'public/wasm/wllama-compat.wasm' : 'public/wasm/wllama.wasm', | |
| ); | |
| return { | |
| crossOriginIsolated: globalThis.crossOriginIsolated, | |
| sharedArrayBuffer: typeof SharedArrayBuffer !== 'undefined', | |
| hardwareConcurrency: navigator.hardwareConcurrency || 1, | |
| webgpu, | |
| storage, | |
| browser: { | |
| userAgent: navigator.userAgent, | |
| platform: navigator.platform, | |
| language: navigator.language, | |
| }, | |
| runtime: { | |
| implementation: 'bonsai-wllama', | |
| wllamaRevision: WLLAMA_REVISION, | |
| llamaCppRevision: LLAMA_CPP_REVISION, | |
| patchSetSha256: runtimeSource.patchSet.sha256, | |
| moduleSha256: runtimeArtifact('vendor/wllama-bonsai/esm/index.js').sha256, | |
| wasmFlavor, | |
| wasmSha256: wasmArtifact.sha256, | |
| compatWorkerSha256: wasmFlavor === 'compat' | |
| ? runtimeArtifact('public/wasm/wllama-compat.js').sha256 | |
| : null, | |
| tokenEmbeddingOnWebGPU: true, | |
| tensorPlacementOverrides: false, | |
| }, | |
| }; | |
| } | |
| private async inspectCachedShards( | |
| requestId: string, | |
| wllama: Wllama, | |
| model: ManifestModelV2, | |
| urls: readonly string[], | |
| signal: AbortSignal, | |
| sink: EventSink, | |
| ): Promise<CachedShardState> { | |
| const blobs: Array<Blob | null> = []; | |
| const shards = initialShardProgress(model); | |
| let cachedBytes = 0; | |
| const emitCacheProgress = (shardIndex: number, phase: 'cache' | 'verify'): void => { | |
| const shard = shards[shardIndex] ?? null; | |
| emitProgress(sink, requestId, { | |
| phase, | |
| loadedBytes: totalShardProgress(shards), | |
| totalBytes: model.downloadBytes, | |
| shardIndex, | |
| shardCount: model.files.length, | |
| shardPath: shard?.path ?? null, | |
| stageProgress: phase === 'verify' | |
| ? totalVerifiedShardProgress(shards) / model.downloadBytes | |
| : totalShardProgress(shards) / model.downloadBytes, | |
| residentBytes: null, | |
| shards: copyShardProgress(shards), | |
| }); | |
| }; | |
| for (let index = 0; index < urls.length; index += 1) { | |
| const url = urls[index]; | |
| const file = model.files[index]; | |
| if (!url || !file) { | |
| throw new EngineRuntimeError('INVALID_SHARD_ORDER', 'Manifest shard URL and file counts differ.'); | |
| } | |
| const key = await wllama.cacheManager.getNameFromURL(url); | |
| const [blob, metadata] = await Promise.all([ | |
| wllama.cacheManager.open(url), | |
| wllama.cacheManager.getMetadata(key), | |
| ]); | |
| if (blob?.size === file.bytes && metadata?.sha256 === file.sha256) { | |
| const shard = shards[index]; | |
| if (!shard) throw new EngineRuntimeError('INVALID_SHARD_ORDER', 'Manifest shard progress is incomplete.'); | |
| shard.state = 'verifying'; | |
| shard.loadedBytes = file.bytes; | |
| shard.verifiedBytes = 0; | |
| emitCacheProgress(index, 'verify'); | |
| const integrity = await verifyBlobSha256(blob, file.sha256, signal, (loadedBytes) => { | |
| shard.verifiedBytes = Math.min(loadedBytes, shard.totalBytes); | |
| emitCacheProgress(index, 'verify'); | |
| }); | |
| if (integrity.matches) { | |
| blobs.push(blob); | |
| cachedBytes += file.bytes; | |
| shard.loadedBytes = file.bytes; | |
| shard.verifiedBytes = file.bytes; | |
| shard.state = 'cached'; | |
| emitCacheProgress(index, 'cache'); | |
| continue; | |
| } | |
| } | |
| if (blob) { | |
| await wllama.cacheManager.delete(url); | |
| } | |
| blobs.push(null); | |
| const shard = shards[index]; | |
| if (shard) { | |
| shard.loadedBytes = 0; | |
| shard.verifiedBytes = 0; | |
| shard.state = 'queued'; | |
| emitCacheProgress(index, 'cache'); | |
| } | |
| } | |
| return { blobs, cachedBytes, shards }; | |
| } | |
| private async downloadShards( | |
| requestId: string, | |
| wllama: Wllama, | |
| model: ManifestModelV2, | |
| urls: readonly string[], | |
| cached: CachedShardState, | |
| signal: AbortSignal, | |
| sink: EventSink, | |
| ): Promise<Blob[]> { | |
| const downloadController = new AbortController(); | |
| const abortDownloads = () => downloadController.abort(signal.reason); | |
| if (signal.aborted) abortDownloads(); | |
| else signal.addEventListener('abort', abortDownloads, { once: true }); | |
| const downloadSignal = downloadController.signal; | |
| const loadedByShard = model.files.map((file, index) => cached.blobs[index] ? file.bytes : 0); | |
| const shards = cached.shards ?? initialShardProgress(model, cached.blobs); | |
| const emitDownloadProgress = (shardIndex: number, phase: 'download' | 'verify' = 'download'): void => { | |
| const file = model.files[shardIndex] ?? null; | |
| emitProgress(sink, requestId, { | |
| phase, | |
| loadedBytes: loadedByShard.reduce((sum, value) => sum + value, 0), | |
| totalBytes: model.downloadBytes, | |
| shardIndex, | |
| shardCount: model.files.length, | |
| shardPath: file?.path ?? null, | |
| stageProgress: phase === 'verify' | |
| ? totalVerifiedShardProgress(shards) / model.downloadBytes | |
| : totalShardProgress(shards) / model.downloadBytes, | |
| residentBytes: null, | |
| shards: copyShardProgress(shards), | |
| }); | |
| }; | |
| emitDownloadProgress(0); | |
| const queue = model.files | |
| .map((_, index) => index) | |
| .filter((index) => cached.blobs[index] === null); | |
| let stopDequeuing = false; | |
| let firstFailureReserved = false; | |
| let firstFailure: EngineRuntimeError | null = null; | |
| const worker = async (): Promise<void> => { | |
| while (!stopDequeuing && queue.length > 0) { | |
| throwIfAborted(downloadSignal); | |
| const index = queue.shift(); | |
| if (index === undefined) { | |
| return; | |
| } | |
| const file = model.files[index]; | |
| const url = urls[index]; | |
| if (!file || !url) { | |
| throw new EngineRuntimeError('INVALID_SHARD_ORDER', 'Manifest shard URL and file counts differ.'); | |
| } | |
| const shard = shards[index]; | |
| if (!shard) throw new EngineRuntimeError('INVALID_SHARD_ORDER', 'Manifest shard progress is incomplete.'); | |
| try { | |
| shard.state = 'downloading'; | |
| shard.loadedBytes = 0; | |
| shard.verifiedBytes = 0; | |
| emitDownloadProgress(index); | |
| await wllama.cacheManager.download(url, { | |
| signal: downloadSignal, | |
| metadataAdditional: { | |
| originalURL: url, | |
| originalSize: file.bytes, | |
| sha256: file.sha256, | |
| }, | |
| progressCallback: ({ loaded }) => { | |
| loadedByShard[index] = Math.min(loaded, file.bytes); | |
| shard.loadedBytes = loadedByShard[index] ?? 0; | |
| emitDownloadProgress(index); | |
| }, | |
| }); | |
| throwIfAborted(downloadSignal); | |
| const blob = await wllama.cacheManager.open(url); | |
| if (!blob || blob.size !== file.bytes) { | |
| throw new EngineRuntimeError( | |
| 'SHARD_SIZE_MISMATCH', | |
| `Downloaded shard ${file.path} has ${blob?.size ?? 0} bytes; expected ${file.bytes}.`, | |
| ); | |
| } | |
| shard.state = 'verifying'; | |
| shard.loadedBytes = file.bytes; | |
| shard.verifiedBytes = 0; | |
| loadedByShard[index] = file.bytes; | |
| emitDownloadProgress(index, 'verify'); | |
| const integrity = await verifyBlobSha256(blob, file.sha256, downloadSignal, (loadedBytes) => { | |
| shard.verifiedBytes = Math.min(loadedBytes, file.bytes); | |
| emitDownloadProgress(index, 'verify'); | |
| }); | |
| if (!integrity.matches) { | |
| throw new EngineRuntimeError( | |
| 'SHARD_HASH_MISMATCH', | |
| `Downloaded shard ${file.path} has SHA-256 ${integrity.actualSha256}; expected ${file.sha256}.`, | |
| ); | |
| } | |
| const cacheKey = await wllama.cacheManager.getNameFromURL(url); | |
| await wllama.cacheManager.writeMetadata(cacheKey, { | |
| etag: 'manifest-v2', | |
| originalURL: url, | |
| originalSize: file.bytes, | |
| sha256: file.sha256, | |
| }); | |
| cached.blobs[index] = blob; | |
| loadedByShard[index] = file.bytes; | |
| shard.loadedBytes = file.bytes; | |
| shard.verifiedBytes = file.bytes; | |
| shard.state = 'complete'; | |
| emitDownloadProgress(index); | |
| } catch (error) { | |
| const failure = classifyShardFailure(error, downloadSignal); | |
| stopDequeuing = true; | |
| const isPrimaryFailure = !firstFailureReserved; | |
| if (isPrimaryFailure) firstFailureReserved = true; | |
| cached.blobs[index] = null; | |
| loadedByShard[index] = 0; | |
| shard.loadedBytes = 0; | |
| shard.verifiedBytes = 0; | |
| shard.state = 'error'; | |
| let partialDeleted = false; | |
| let cleanupError: unknown; | |
| try { | |
| await wllama.cacheManager.delete(url); | |
| partialDeleted = true; | |
| } catch (deleteError) { | |
| cleanupError = deleteError; | |
| } | |
| emitDownloadProgress(index, failure === 'verification' ? 'verify' : 'download'); | |
| const normalizedFailure = shardFailure( | |
| error, | |
| failure, | |
| index, | |
| model.files.length, | |
| file.path, | |
| partialDeleted, | |
| cleanupError, | |
| ); | |
| if (isPrimaryFailure) firstFailure = normalizedFailure; | |
| return; | |
| } | |
| } | |
| }; | |
| const workers = Array.from({ length: Math.min(3, Math.max(1, queue.length)) }, () => worker()); | |
| let workerResults: PromiseSettledResult<void>[]; | |
| try { | |
| workerResults = await Promise.allSettled(workers); | |
| } finally { | |
| signal.removeEventListener('abort', abortDownloads); | |
| } | |
| if (firstFailure !== null) throw firstFailure; | |
| const rejectedWorker = workerResults.find( | |
| (result): result is PromiseRejectedResult => result.status === 'rejected', | |
| ); | |
| if (rejectedWorker) throw rejectedWorker.reason; | |
| throwIfAborted(signal); | |
| const blobs = cached.blobs; | |
| if (blobs.some((blob) => blob === null)) { | |
| throw new EngineRuntimeError('SHARD_LOAD_INCOMPLETE', 'One or more model shards are unavailable.'); | |
| } | |
| return blobs as Blob[]; | |
| } | |
| async loadModel( | |
| requestId: string, | |
| params: LoadModelParams, | |
| signal: AbortSignal, | |
| sink: EventSink, | |
| ): Promise<LoadModelResult> { | |
| if (params.tensorPlacement) { | |
| throw new EngineRuntimeError( | |
| 'TENSOR_PLACEMENT_UNSUPPORTED', | |
| 'Arbitrary tensor placement is not exposed. The pinned Bonsai build applies its validated token-embedding WebGPU override automatically.', | |
| params.tensorPlacement, | |
| ); | |
| } | |
| const tuning = resolveLoadTuning(params.benchmarkTuning, params.backend); | |
| const wasmFlavor = selectWasmFlavor(tuning.wasmFlavor); | |
| await this.unload(); | |
| this.nativeLog.clear(); | |
| emitProgress(sink, requestId, { | |
| phase: 'manifest', | |
| loadedBytes: 0, | |
| totalBytes: 1, | |
| shardIndex: null, | |
| shardCount: 0, | |
| shardPath: null, | |
| stageProgress: 0, | |
| residentBytes: null, | |
| shards: [], | |
| }); | |
| const manifest = params.manifestUrl !== undefined | |
| ? await loadModelManifestV2(params.manifestUrl, signal) | |
| : parseModelManifestV2(params.manifest); | |
| throwIfAborted(signal); | |
| emitProgress(sink, requestId, { | |
| phase: 'manifest', | |
| loadedBytes: 1, | |
| totalBytes: 1, | |
| shardIndex: null, | |
| shardCount: 0, | |
| shardPath: null, | |
| stageProgress: 1, | |
| residentBytes: null, | |
| shards: [], | |
| }); | |
| const model = findManifestModel(manifest, params.modelId); | |
| const contextSize = params.contextSize ?? model.defaultContext; | |
| const contextPolicy = evaluateModelContextPolicy(model, contextSize, { | |
| tuningScope: tuning.scope, | |
| requestedBackend: params.backend, | |
| flashMode: tuning.flashMode, | |
| kvCacheType: tuning.kvCacheType, | |
| }); | |
| if (!contextPolicy.allowed) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_CONTEXT_SIZE', | |
| `Context size must be between 1 and ${contextPolicy.limit} for this model and runtime policy.`, | |
| ); | |
| } | |
| const urls = orderedShardUrls(manifest, model); | |
| const wllama = await this.ensureWllama(wasmFlavor); | |
| const cached = await this.inspectCachedShards(requestId, wllama, model, urls, signal, sink); | |
| const [adapter, storage] = await Promise.all([ | |
| inspectWebGpuAdapter(), | |
| estimateStorage(), | |
| ]); | |
| const gate = evaluateModelGate(model, params.backend, adapter, storage, cached.cachedBytes); | |
| if (!gate.allowed || gate.selectedBackend === null) { | |
| throw new EngineRuntimeError('MODEL_GATE_REJECTED', gate.reasons.join(' '), gate); | |
| } | |
| if (gate.selectedBackend === 'webgpu' && !wllama.isSupportWebGPU()) { | |
| throw new EngineRuntimeError( | |
| 'WEBGPU_RUNTIME_UNAVAILABLE', | |
| 'The adapter gate passed, but the pinned Bonsai wllama build reports WebGPU unavailable in this browser runtime.', | |
| ); | |
| } | |
| const blobs = await this.downloadShards(requestId, wllama, model, urls, cached, signal, sink); | |
| throwIfAborted(signal); | |
| emitProgress(sink, requestId, { | |
| phase: 'load', | |
| loadedBytes: model.downloadBytes, | |
| totalBytes: model.downloadBytes, | |
| shardIndex: null, | |
| shardCount: model.files.length, | |
| shardPath: null, | |
| stageProgress: 0, | |
| residentBytes: null, | |
| shards: copyShardProgress(cached.shards ?? initialShardProgress(model, blobs)), | |
| }); | |
| const defaultThreads = Math.max(1, Math.floor((navigator.hardwareConcurrency || 1) / 2)); | |
| const threads = params.threads ?? (gate.selectedBackend === 'wasm' ? defaultThreads : 1); | |
| if (!Number.isSafeInteger(threads) || threads <= 0) { | |
| throw new EngineRuntimeError('INVALID_THREAD_COUNT', 'Thread count must be a positive integer.'); | |
| } | |
| const loadShards = copyShardProgress(cached.shards ?? initialShardProgress(model, blobs)); | |
| const runtimeBatch = resolveRuntimeBatchShape(tuning, wasmFlavor); | |
| const stopLoadProgress = this.nativeLog.onModelLoadProgress((progress) => { | |
| const liveReport = this.nativeLog.report(); | |
| emitProgress(sink, requestId, { | |
| phase: 'load', | |
| loadedBytes: model.downloadBytes, | |
| totalBytes: model.downloadBytes, | |
| shardIndex: null, | |
| shardCount: model.files.length, | |
| shardPath: null, | |
| stageProgress: progress.value, | |
| nativeStage: progress.current, | |
| residentBytes: liveReport.allocatedBufferBytes ?? null, | |
| shards: copyShardProgress(loadShards), | |
| }); | |
| }); | |
| try { | |
| await this.raceWebGpuDeviceLoss('load', gate.selectedBackend, model, wllama.loadModel(blobs, { | |
| n_gpu_layers: gate.selectedBackend === 'webgpu' ? 99999 : 0, | |
| offload_token_embedding: gate.selectedBackend === 'webgpu', | |
| n_ctx: contextSize, | |
| n_threads: threads, | |
| n_batch: runtimeBatch.nBatch, | |
| n_ubatch: runtimeBatch.nUbatch, | |
| seed: 42, | |
| flash_attn: tuning.flashMode === 'auto', | |
| warmup: false, | |
| cache_type_k: tuning.kvCacheType, | |
| cache_type_v: tuning.kvCacheType, | |
| })); | |
| await this.assertWebGpuAlive('load', gate.selectedBackend, model); | |
| throwIfAborted(signal); | |
| const report = this.nativeLog.report(); | |
| try { | |
| assertBackendPolicy(model, gate.selectedBackend, report); | |
| } catch (error) { | |
| throw new EngineRuntimeError( | |
| 'BACKEND_TRIPWIRE', | |
| error instanceof Error ? error.message : String(error), | |
| report, | |
| ); | |
| } | |
| const context = wllama.getLoadedContextInfo(); | |
| const template = wllama.getChatTemplate() ?? ''; | |
| const wasmArtifact = runtimeArtifact( | |
| wasmFlavor === 'compat' ? 'public/wasm/wllama-compat.wasm' : 'public/wasm/wllama.wasm', | |
| ); | |
| const compatWorkerSha256 = wasmFlavor === 'compat' | |
| ? runtimeArtifact('public/wasm/wllama-compat.js').sha256 | |
| : null; | |
| const tuningApplied = (tuning.nBatch === null || context.n_batch === tuning.nBatch) | |
| && (tuning.nUbatch === null || context.n_ubatch === tuning.nUbatch) | |
| && report.flashAttention === (tuning.flashMode === 'auto') | |
| && report.cacheTypeK === tuning.kvCacheType | |
| && report.cacheTypeV === tuning.kvCacheType | |
| && (gate.selectedBackend !== 'webgpu' | |
| || (report.webgpuKvBufferBytes !== null && report.webgpuKvBufferBytes > 0)) | |
| && (tuning.wasmFlavor === 'auto' || tuning.wasmFlavor === wasmFlavor); | |
| const tuningResult: LoadModelResult['tuning'] = { | |
| scope: tuning.scope, | |
| requested: { | |
| nBatch: tuning.nBatch, | |
| nUbatch: tuning.nUbatch, | |
| flashMode: tuning.flashMode, | |
| cacheTypeK: tuning.kvCacheType, | |
| cacheTypeV: tuning.kvCacheType, | |
| wasmFlavor: tuning.wasmFlavor, | |
| }, | |
| observed: { | |
| nBatch: context.n_batch, | |
| nUbatch: context.n_ubatch, | |
| flashAttention: report.flashAttention, | |
| cacheTypeK: report.cacheTypeK, | |
| cacheTypeV: report.cacheTypeV, | |
| kvBufferBytes: report.webgpuKvBufferBytes, | |
| wasmFlavor, | |
| wasmSha256: wasmArtifact.sha256, | |
| compatWorkerSha256, | |
| }, | |
| applied: tuningApplied, | |
| }; | |
| if (!tuningApplied) { | |
| throw new EngineRuntimeError( | |
| tuning.scope === 'benchmark' | |
| ? 'BENCHMARK_TUNING_NOT_APPLIED' | |
| : 'RUNTIME_POLICY_NOT_APPLIED', | |
| tuning.scope === 'benchmark' | |
| ? 'The native runtime did not apply the requested benchmark tuning exactly.' | |
| : 'The native runtime did not preserve the fixed release tuning policy.', | |
| tuningResult, | |
| ); | |
| } | |
| this.loaded = { | |
| manifest, | |
| model, | |
| backend: gate.selectedBackend, | |
| tuningScope: tuning.scope, | |
| contextSize: context.n_ctx, | |
| batchSize: context.n_batch, | |
| microBatchSize: context.n_ubatch, | |
| vocabularySize: context.n_vocab, | |
| }; | |
| emitProgress(sink, requestId, { | |
| phase: 'load', | |
| loadedBytes: model.downloadBytes, | |
| totalBytes: model.downloadBytes, | |
| shardIndex: null, | |
| shardCount: model.files.length, | |
| shardPath: null, | |
| stageProgress: 1, | |
| residentBytes: report.allocatedBufferBytes ?? null, | |
| shards: copyShardProgress(loadShards), | |
| }); | |
| return { | |
| modelId: model.id, | |
| backend: gate.selectedBackend, | |
| gate, | |
| shardUrls: urls, | |
| context: { | |
| size: context.n_ctx, | |
| trainingSize: context.n_ctx_train, | |
| layerCount: context.n_layer, | |
| vocabularySize: context.n_vocab, | |
| batchSize: context.n_batch, | |
| microBatchSize: context.n_ubatch, | |
| }, | |
| tuning: tuningResult, | |
| chatTemplate: { | |
| bytes: new TextEncoder().encode(template).byteLength, | |
| hasThinkMarker: template.includes('<think>'), | |
| hasToolCallMarker: template.includes('<tool_call>'), | |
| hasToolResponseMarker: template.includes('<tool_response>'), | |
| }, | |
| backendReport: report, | |
| }; | |
| } catch (error) { | |
| if (error instanceof EngineRuntimeError && error.code === 'WEBGPU_DEVICE_LOST') throw error; | |
| await this.assertWebGpuAlive('load', gate.selectedBackend, model); | |
| await wllama.exit().catch(() => undefined); | |
| this.wllama = null; | |
| this.wllamaFlavor = null; | |
| this.loaded = null; | |
| throw error; | |
| } finally { | |
| stopLoadProgress(); | |
| } | |
| } | |
| async generate( | |
| requestId: string, | |
| params: GenerateParams, | |
| signal: AbortSignal, | |
| sink: EventSink, | |
| ): Promise<GenerateResult> { | |
| const loaded = this.loaded; | |
| const wllama = this.wllama; | |
| if (!loaded || !wllama?.isModelLoaded()) { | |
| throw new EngineRuntimeError('MODEL_NOT_LOADED', 'Load a model before generating.'); | |
| } | |
| await this.assertWebGpuAlive('generate', loaded.backend, loaded.model); | |
| if (params.messages.length === 0) { | |
| throw new EngineRuntimeError('EMPTY_MESSAGES', 'Generation requires at least one chat message.'); | |
| } | |
| throwIfAborted(signal); | |
| let text = ''; | |
| let reasoningText = ''; | |
| let finishReason: GenerateResult['finishReason'] = null; | |
| let usage: GenerateResult['usage'] = null; | |
| let timings: GenerateResult['timings'] = null; | |
| let livePromptTokensPerSecond = 0; | |
| let liveDecodeTokensPerSecond = 0; | |
| let livePromptProcessed = 0; | |
| let livePromptTotal = 0; | |
| let livePromptCached = 0; | |
| let fallbackCompletionTokens = 0; | |
| let fallbackDecodeStartedAt: number | null = null; | |
| const decodeRateSamples: Array<{ tokens: number; elapsedMs: number }> = []; | |
| const usePerTokenTimings = this.wllamaFlavor !== 'compat'; | |
| const sampleDecodeRate = (tokens: number, elapsedMs: number): void => { | |
| const lastSample = decodeRateSamples.at(-1); | |
| if ( | |
| tokens <= 0 | |
| || !Number.isFinite(elapsedMs) | |
| || (lastSample && (tokens <= lastSample.tokens || elapsedMs < lastSample.elapsedMs)) | |
| ) { | |
| return; | |
| } | |
| decodeRateSamples.push({ tokens, elapsedMs }); | |
| while (decodeRateSamples.length > 2) { | |
| const first = decodeRateSamples[0]; | |
| if (!first) break; | |
| const tokenSpan = tokens - first.tokens; | |
| const timeSpan = elapsedMs - first.elapsedMs; | |
| if (tokenSpan <= 12 && timeSpan <= 2_000) break; | |
| decodeRateSamples.shift(); | |
| } | |
| const first = decodeRateSamples[0]; | |
| if (!first) return; | |
| const tokenSpan = tokens - first.tokens; | |
| const timeSpan = elapsedMs - first.elapsedMs; | |
| if (tokenSpan <= 0 || timeSpan <= 0) return; | |
| const rate = tokenSpan * 1_000 / timeSpan; | |
| if (Number.isFinite(rate) && rate >= 0) { | |
| liveDecodeTokensPerSecond = rate; | |
| } | |
| }; | |
| const tokenTrace: EngineSampledTokenTraceEntry[] | null = params.returnTokenIds === true ? [] : null; | |
| const streamedToolCalls = new ToolCallAccumulator(); | |
| emitGenerationProgress(sink, requestId, { | |
| phase: 'prefill', | |
| promptProcessed: 0, | |
| promptTotal: 0, | |
| promptCached: 0, | |
| completionTokens: 0, | |
| elapsedMs: 0, | |
| promptTokensPerSecond: 0, | |
| decodeTokensPerSecond: 0, | |
| }); | |
| try { | |
| const completionOptions = { | |
| messages: params.messages as ChatCompletionMessage[], | |
| max_tokens: params.maxTokens ?? DEFAULT_MAX_TOKENS, | |
| temperature: params.temperature ?? 0, | |
| top_p: params.topP, | |
| top_k: params.topK ?? 1, | |
| min_p: params.minP, | |
| seed: params.seed ?? 42, | |
| tools: params.tools as ChatCompletionTool[] | undefined, | |
| tool_choice: params.toolChoice, | |
| cache_prompt: params.cachePrompt ?? true, | |
| ...(params.returnTokenIds === true | |
| ? { logprobs: true, top_logprobs: TOKEN_TRACE_TOP_LOGPROBS } | |
| : {}), | |
| abortSignal: signal, | |
| }; | |
| const nonStreamingToolTrace = tokenTrace !== null && (params.tools?.length ?? 0) > 0; | |
| if (nonStreamingToolTrace) { | |
| const response = await this.raceWebGpuDeviceLoss( | |
| 'generate', | |
| loaded.backend, | |
| loaded.model, | |
| wllama.createChatCompletion({ ...completionOptions, stream: false }), | |
| ); | |
| const choice = response.choices[0]; | |
| if (!choice) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_COMPLETION_RESPONSE', | |
| 'The native runtime returned a chat completion without a choice.', | |
| ); | |
| } | |
| text = typeof choice.message.content === 'string' ? choice.message.content : ''; | |
| reasoningText = asReasoningContent(choice.message); | |
| finishReason = choice.finish_reason; | |
| appendSampledTokenTrace(tokenTrace, response); | |
| streamedToolCalls.append(choice.message.tool_calls?.map((call, index) => ({ index, ...call }))); | |
| const responseTimings = (response as unknown as { timings?: ResultTimings }).timings; | |
| usage = mapUsage(response.usage, responseTimings); | |
| timings = mapTimings(responseTimings); | |
| if (text || reasoningText) { | |
| sink({ | |
| type: 'event', | |
| requestId, | |
| event: 'token', | |
| text, | |
| ...(reasoningText ? { reasoningDelta: reasoningText } : {}), | |
| }); | |
| } | |
| } else { | |
| await this.raceWebGpuDeviceLoss('generate', loaded.backend, loaded.model, wllama.createChatCompletion({ | |
| ...completionOptions, | |
| stream: true, | |
| ...(usePerTokenTimings ? { timings_per_token: true, return_progress: true } : {}), | |
| onData: (chunk: ChatCompletionChunk) => { | |
| const liveChunk = chunk as LiveCompletionChunk; | |
| const delta = chunk.choices[0]?.delta.content; | |
| const textDelta = typeof delta === 'string' ? delta : ''; | |
| const reasoningDelta = asReasoningDelta(chunk); | |
| const progress = liveChunk.prompt_progress; | |
| if (progress) { | |
| const processed = Math.max(0, progress.processed); | |
| const cached = Math.max(0, Math.min(processed, progress.cache)); | |
| const elapsedMs = Math.max(0, progress.time_ms); | |
| const processedNow = Math.max(0, processed - cached); | |
| livePromptProcessed = processed; | |
| livePromptTotal = Math.max(processed, progress.total); | |
| livePromptCached = cached; | |
| const promptRate = elapsedMs >= MIN_LIVE_RATE_WINDOW_MS | |
| ? processedNow * 1000 / elapsedMs | |
| : 0; | |
| if (Number.isFinite(promptRate) && promptRate >= 0) { | |
| livePromptTokensPerSecond = promptRate; | |
| } | |
| emitGenerationProgress(sink, requestId, { | |
| phase: 'prefill', | |
| promptProcessed: processed, | |
| promptTotal: livePromptTotal, | |
| promptCached: cached, | |
| completionTokens: 0, | |
| elapsedMs, | |
| promptTokensPerSecond: livePromptTokensPerSecond, | |
| decodeTokensPerSecond: 0, | |
| }); | |
| } else if (usePerTokenTimings && chunk.timings) { | |
| const predictedTokens = Math.max(0, chunk.timings.predicted_n); | |
| const predictedMs = Math.max(0, chunk.timings.predicted_ms); | |
| sampleDecodeRate(predictedTokens, predictedMs); | |
| const finalPromptRate = Math.max(0, chunk.timings.prompt_per_second); | |
| if (Number.isFinite(finalPromptRate)) { | |
| livePromptTokensPerSecond = finalPromptRate; | |
| } | |
| emitGenerationProgress(sink, requestId, { | |
| phase: 'decode', | |
| promptProcessed: Math.max(0, chunk.timings.prompt_n), | |
| promptTotal: Math.max(0, chunk.timings.prompt_n), | |
| promptCached: Math.max(0, chunk.timings.cache_n), | |
| completionTokens: predictedTokens, | |
| elapsedMs: predictedMs, | |
| promptTokensPerSecond: livePromptTokensPerSecond, | |
| decodeTokensPerSecond: liveDecodeTokensPerSecond, | |
| }); | |
| } else if (!usePerTokenTimings && (textDelta || reasoningDelta)) { | |
| fallbackCompletionTokens += 1; | |
| const now = performance.now(); | |
| fallbackDecodeStartedAt ??= now; | |
| const elapsedMs = now - fallbackDecodeStartedAt; | |
| sampleDecodeRate(fallbackCompletionTokens, elapsedMs); | |
| emitGenerationProgress(sink, requestId, { | |
| phase: 'decode', | |
| promptProcessed: livePromptProcessed, | |
| promptTotal: livePromptTotal, | |
| promptCached: livePromptCached, | |
| completionTokens: fallbackCompletionTokens, | |
| elapsedMs, | |
| promptTokensPerSecond: livePromptTokensPerSecond, | |
| decodeTokensPerSecond: liveDecodeTokensPerSecond, | |
| }); | |
| } | |
| if (tokenTrace !== null) appendSampledTokenTrace(tokenTrace, chunk); | |
| const toolCallProgress = streamedToolCalls.append(chunk.choices[0]?.delta.tool_calls); | |
| for (const progress of toolCallProgress) { | |
| sink({ | |
| type: 'event', | |
| requestId, | |
| event: 'tool-call', | |
| ...progress, | |
| }); | |
| } | |
| if (textDelta || reasoningDelta) { | |
| text += textDelta; | |
| reasoningText += reasoningDelta ?? ''; | |
| sink({ | |
| type: 'event', | |
| requestId, | |
| event: 'token', | |
| text: textDelta, | |
| ...(reasoningDelta ? { reasoningDelta } : {}), | |
| }); | |
| } | |
| finishReason = chunk.choices[0]?.finish_reason ?? finishReason; | |
| usage = mapUsage(chunk.usage, chunk.timings) ?? usage; | |
| timings = mapTimings(chunk.timings) ?? timings; | |
| }, | |
| })); | |
| } | |
| } catch (error) { | |
| if (error instanceof EngineRuntimeError) throw error; | |
| await this.assertWebGpuAlive('generate', loaded.backend, loaded.model); | |
| throw new EngineRuntimeError( | |
| 'NATIVE_COMPLETION_FAILED', | |
| error instanceof Error ? error.message : String(error), | |
| { nativeLog: this.nativeLog.recent() }, | |
| ); | |
| } | |
| await this.assertWebGpuAlive('generate', loaded.backend, loaded.model); | |
| throwIfAborted(signal); | |
| const tokenIds = tokenTrace?.map((entry) => entry.selected.id) ?? null; | |
| const accounting = tokenTraceAccounting(tokenTrace, usage); | |
| const report = this.nativeLog.report(); | |
| try { | |
| assertBackendPolicy(loaded.model, loaded.backend, report); | |
| } catch (error) { | |
| await this.unload(); | |
| throw new EngineRuntimeError( | |
| 'BACKEND_TRIPWIRE', | |
| error instanceof Error ? error.message : String(error), | |
| report, | |
| ); | |
| } | |
| let toolCalls: GenerateResult['toolCalls']; | |
| let toolCallError: string | undefined; | |
| try { | |
| const completedToolCalls = streamedToolCalls.finish( | |
| finishReason === 'tool_calls' || streamedToolCalls.hasCalls(), | |
| ); | |
| toolCalls = finishReason === 'tool_calls' ? completedToolCalls : []; | |
| } catch (error) { | |
| toolCalls = []; | |
| toolCallError = error instanceof Error ? error.message : String(error); | |
| } | |
| return { | |
| text, | |
| reasoningText, | |
| ...(toolCallError ? { toolCallError } : {}), | |
| tokenIds, | |
| tokenTrace, | |
| tokenTraceAccounting: accounting, | |
| finishReason, | |
| toolCalls, | |
| usage, | |
| timings, | |
| }; | |
| } | |
| async scoreSequence( | |
| params: ScoreSequenceParams, | |
| signal: AbortSignal, | |
| ): Promise<ScoreSequenceResult> { | |
| const loaded = this.loaded; | |
| const wllama = this.wllama; | |
| if (!loaded || !wllama?.isModelLoaded()) { | |
| throw new EngineRuntimeError('MODEL_NOT_LOADED', 'Load a model before scoring a sequence.'); | |
| } | |
| if ( | |
| loaded.model.id !== '27b' | |
| || loaded.backend !== 'webgpu' | |
| || loaded.tuningScope !== 'benchmark' | |
| ) { | |
| throw new EngineRuntimeError( | |
| 'SCORE_SEQUENCE_UNAVAILABLE', | |
| 'Teacher-forced scoring is diagnostic-only and requires the loaded 27B WebGPU benchmark path.', | |
| { | |
| modelId: loaded.model.id, | |
| backend: loaded.backend, | |
| tuningScope: loaded.tuningScope, | |
| }, | |
| ); | |
| } | |
| if (params.topK !== TOKEN_TRACE_TOP_LOGPROBS) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE', | |
| `Teacher-forced scoring requires topK=${TOKEN_TRACE_TOP_LOGPROBS}.`, | |
| ); | |
| } | |
| if (!Array.isArray(params.promptTokenIds) | |
| || params.promptTokenIds.length !== STATE_DRIFT_PROMPT_TOKENS) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE', | |
| `Teacher-forced scoring requires exactly ${STATE_DRIFT_PROMPT_TOKENS} rendered prompt token ids.`, | |
| ); | |
| } | |
| if (!Array.isArray(params.referenceTokenIds) | |
| || params.referenceTokenIds.length !== STATE_DRIFT_REFERENCE_TOKENS) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE', | |
| `Teacher-forced scoring requires exactly ${STATE_DRIFT_REFERENCE_TOKENS} reference token ids.`, | |
| ); | |
| } | |
| if (params.promptTokenIds.length + params.referenceTokenIds.length > loaded.contextSize) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE', | |
| 'The rendered prompt and fixed reference sequence exceed the loaded context.', | |
| { | |
| promptTokens: params.promptTokenIds.length, | |
| referenceTokens: params.referenceTokenIds.length, | |
| contextSize: loaded.contextSize, | |
| }, | |
| ); | |
| } | |
| const validateInputTokenIds = (values: readonly number[], field: string): void => { | |
| for (const [index, tokenId] of values.entries()) { | |
| if (!Number.isSafeInteger(tokenId) || tokenId < 0 || tokenId >= loaded.vocabularySize) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE', | |
| `${field}[${index}] is outside the loaded vocabulary.`, | |
| { field, index, tokenId, vocabularySize: loaded.vocabularySize }, | |
| ); | |
| } | |
| } | |
| }; | |
| validateInputTokenIds(params.promptTokenIds, 'promptTokenIds'); | |
| validateInputTokenIds(params.referenceTokenIds, 'referenceTokenIds'); | |
| await this.assertWebGpuAlive('score-sequence', loaded.backend, loaded.model); | |
| throwIfAborted(signal); | |
| const score = async (): Promise<ScoreSequenceResult> => { | |
| const entries: ScoreSequenceResult['entries'] = []; | |
| const createRawCompletion = wllama.createCompletion.bind(wllama) as unknown as ( | |
| options: Record<string, unknown>, | |
| ) => Promise<unknown>; | |
| for (let index = 0; index < params.referenceTokenIds.length; index += 1) { | |
| throwIfAborted(signal); | |
| const referenceTokenId = params.referenceTokenIds[index]!; | |
| const prompt = [ | |
| ...params.promptTokenIds, | |
| ...params.referenceTokenIds.slice(0, index), | |
| ]; | |
| const response = await createRawCompletion({ | |
| // The pinned llama.cpp endpoint accepts a raw token-id prompt even though | |
| // upstream wllama's OAI declaration still narrows this field to strings. | |
| prompt, | |
| stream: false, | |
| max_tokens: 1, | |
| temperature: 0, | |
| top_k: 1, | |
| logprobs: TOKEN_TRACE_TOP_LOGPROBS, | |
| logit_bias: { [String(referenceTokenId)]: TEACHER_FORCE_LOGIT_BIAS }, | |
| cache_prompt: index !== 0, | |
| post_sampling_probs: false, | |
| abortSignal: signal, | |
| }); | |
| throwIfAborted(signal); | |
| entries.push(parseTeacherForcedResponse( | |
| response, | |
| index, | |
| referenceTokenId, | |
| loaded.vocabularySize, | |
| )); | |
| } | |
| const meanNll = -entries.reduce( | |
| (sum, entry) => sum + entry.selectedReference.logprob, | |
| 0, | |
| ) / entries.length; | |
| const perplexity = Math.exp(meanNll); | |
| if (!Number.isFinite(meanNll) || !Number.isFinite(perplexity)) { | |
| throw new EngineRuntimeError( | |
| 'INVALID_SCORE_SEQUENCE_RESPONSE', | |
| 'Teacher-forced scoring produced a non-finite mean NLL or perplexity.', | |
| { meanNll, perplexity }, | |
| ); | |
| } | |
| return { | |
| method: { | |
| promptMode: 'raw-token-id-prefix', | |
| maxTokensPerStep: 1, | |
| temperature: 0, | |
| topK: 1, | |
| reportedTopLogprobs: TOKEN_TRACE_TOP_LOGPROBS, | |
| logitBias: TEACHER_FORCE_LOGIT_BIAS, | |
| cachePromptFirst: false, | |
| cachePromptSubsequent: true, | |
| }, | |
| entries, | |
| summary: { | |
| tokenCount: STATE_DRIFT_REFERENCE_TOKENS, | |
| meanNll, | |
| perplexity, | |
| }, | |
| }; | |
| }; | |
| try { | |
| const result = await this.raceWebGpuDeviceLoss( | |
| 'score-sequence', | |
| loaded.backend, | |
| loaded.model, | |
| score(), | |
| ); | |
| await this.assertWebGpuAlive('score-sequence', loaded.backend, loaded.model); | |
| throwIfAborted(signal); | |
| const report = this.nativeLog.report(); | |
| try { | |
| assertBackendPolicy(loaded.model, loaded.backend, report); | |
| } catch (error) { | |
| await this.unload(); | |
| throw new EngineRuntimeError( | |
| 'BACKEND_TRIPWIRE', | |
| error instanceof Error ? error.message : String(error), | |
| report, | |
| ); | |
| } | |
| return result; | |
| } catch (error) { | |
| if (error instanceof EngineRuntimeError && error.code === 'WEBGPU_DEVICE_LOST') throw error; | |
| await this.assertWebGpuAlive('score-sequence', loaded.backend, loaded.model); | |
| throw error; | |
| } | |
| } | |
| async backendReport(): Promise<BackendReport> { | |
| const report = this.nativeLog.report(); | |
| if (this.loaded) { | |
| await this.assertWebGpuAlive('backend-report', this.loaded.backend, this.loaded.model); | |
| } | |
| return report; | |
| } | |
| async unload(): Promise<{ unloaded: true }> { | |
| await this.wllama?.exit(); | |
| this.loaded = null; | |
| return { unloaded: true }; | |
| } | |
| storageEstimate(): Promise<StorageEstimate> { | |
| return estimateStorage(); | |
| } | |
| storagePersist(): Promise<StorageEstimate & { granted: boolean }> { | |
| return persistStorage(); | |
| } | |
| async storageClear(): Promise<StorageEstimate & { cleared: true }> { | |
| await this.unload(); | |
| const wllama = await this.ensureWllama(); | |
| await wllama.cacheManager.clear(); | |
| return { | |
| ...await estimateStorage(), | |
| cleared: true, | |
| }; | |
| } | |
| } | |