Bonsai-Chat-WebGPU / src /engine /runtime.ts
Valeriy Selitskiy
Use 4096-token release defaults
36a170e
Raw
History Blame Contribute Delete
62.8 kB
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,
};
}
}