| import { MODEL, DSPARK, DEFAULTS } from "./core.js"; |
| import { WeightCache } from "./cache.js"; |
| import { downloadForSession } from "./session-download.js"; |
| import { startRuntime } from "./startup.js"; |
| import { splitStreamContent } from "./stream.js"; |
| export class Engine { |
| constructor(onStatus) { |
| this.onStatus = onStatus; |
| this.logs = []; |
| this.ready = false; |
| this.loading = false; |
| this.adapter = null; |
| this.gpuLayers = 0; |
| this.cache = new WeightCache(); |
| this.sessionFiles = new Map(); |
| } |
| async probe() { |
| if (!navigator.gpu) |
| return { |
| ok: false, |
| message: "WebGPU is unavailable. Use a browser with WebGPU enabled.", |
| }; |
| this.adapter = await navigator.gpu.requestAdapter({ |
| powerPreference: "high-performance", |
| }); |
| if (!this.adapter) |
| return { |
| ok: false, |
| message: |
| "No WebGPU adapter found. Enable hardware acceleration in your browser, then reopen this page.", |
| }; |
| const info = this.adapter.info; |
| return { |
| ok: true, |
| name: |
| info?.description || info?.device || info?.vendor || "WebGPU graphics", |
| maxBuffer: this.adapter.limits.maxStorageBufferBindingSize, |
| features: [...this.adapter.features], |
| mobile: /Android|iPhone|iPad/i.test(navigator.userAgent), |
| }; |
| } |
| get url() { |
| if ( |
| import.meta.env.DEV && |
| new URLSearchParams(location.search).get("qa-weights") === "local" |
| ) |
| return location.origin + "/__qa/model/" + MODEL.file; |
| return `https://huggingface.co/${MODEL.weights}/resolve/${MODEL.revision}/${MODEL.file}`; |
| } |
| get urls() { |
| return Array.from({ length: 6 }, (_, i) => |
| this.url.replace("-00001-of-", `-${String(i + 1).padStart(5, "0")}-of-`), |
| ); |
| } |
| get draftURL() { |
| return import.meta.env.DEV && |
| new URLSearchParams(location.search).get("qa-weights") === "local" |
| ? location.origin + "/__qa/dspark" |
| : `https://huggingface.co/${DSPARK.repo}/resolve/${DSPARK.revision}/${DSPARK.file}`; |
| } |
| assets(settings) { |
| const files = this.urls.map((url, index) => ({ |
| url, |
| size: MODEL.shardSizes[index], |
| type: "model", |
| })); |
| if (settings.dspark !== "off") |
| files.push({ url: this.draftURL, size: DSPARK.bytes, type: "draft" }); |
| return files; |
| } |
| async cacheState(settings = this.settings || DEFAULTS) { |
| const entries = await this.cache.list(), |
| files = this.assets(settings), |
| complete = files.map((f) => |
| entries.find( |
| (e) => |
| e.metadata?.originalURL === f.url && |
| e.size === f.size && |
| e.metadata.originalSize === f.size, |
| ), |
| ); |
| return { |
| entries, |
| complete, |
| cached: complete.every(Boolean), |
| bytes: complete.reduce((n, f) => n + (f?.size || 0), 0), |
| }; |
| } |
| async load(settings) { |
| if (this.loading) throw new Error("Model is already loading."); |
| this.loading = true; |
| this.loadController = new AbortController(); |
| const signal = this.loadController.signal, |
| start = performance.now(); |
| this.logs = []; |
| this.gpuLayers = 0; |
| this.downloadMetrics = { requests: 0, retries: 0 }; |
| try { |
| if (this.ready) await this.unload(); |
| const gpu = await this.probe(); |
| if (!gpu.ok) throw new Error(gpu.message); |
| const files = this.assets(settings), |
| total = files.reduce((n, f) => n + f.size, 0), |
| cached = settings.saveModel |
| ? await this.cacheState(settings) |
| : { complete: [], cached: false }; |
| const blobs = await Promise.all( |
| files.map( |
| (file, i) => |
| this.sessionFiles.get(file.url) || |
| (cached.complete[i] |
| ? this.cache.open(cached.complete[i].name) |
| : null), |
| ), |
| ); |
| const restored = blobs.reduce((sum, blob, index) => { |
| if (blob?.size !== files[index].size) blobs[index] = null; |
| return sum + (blobs[index]?.size || 0); |
| }, 0); |
| this.onStatus({ |
| phase: restored === total ? "compile" : "download", |
| progress: restored / total, |
| cached: cached.cached, |
| label: |
| restored === total |
| ? "Restoring this model" |
| : "Downloading Q8 model" + |
| (settings.dspark !== "off" ? " + DSpark" : ""), |
| }); |
| if (settings.saveModel) navigator.storage?.persist?.().catch(() => {}); |
| const missing = files |
| .map((file, index) => ({ ...file, index })) |
| .filter((f) => !blobs[f.index]); |
| if (missing.length) |
| await downloadForSession( |
| missing.map((f) => f.url), |
| { |
| signal, |
| total, |
| fileSizes: missing.map((f) => f.size), |
| parallel: gpu.mobile ? 4 : 8, |
| onProgress: ({ loaded, requests, retries }) => { |
| this.downloadMetrics = { requests, retries }; |
| this.onStatus({ |
| phase: "download", |
| progress: (loaded + restored) / total, |
| loaded: loaded + restored, |
| total, |
| elapsed: (performance.now() - start) / 1000, |
| label: |
| settings.dspark !== "off" |
| ? "Downloading Q8 model + DSpark" |
| : "Downloading Q8 model", |
| persistent: settings.saveModel, |
| }); |
| }, |
| onFileComplete: async (blob, index) => { |
| const file = missing[index]; |
| if ( |
| file.type === "model" && |
| new TextDecoder().decode( |
| await blob.slice(0, 4).arrayBuffer(), |
| ) !== "GGUF" |
| ) |
| throw new Error( |
| "The weight download was not a valid GGUF file. Please retry.", |
| ); |
| this.sessionFiles.set(file.url, blob); |
| blobs[file.index] = blob; |
| if (settings.saveModel) { |
| signal.throwIfAborted(); |
| await this.cache.write( |
| await this.cache.getNameFromURL(file.url), |
| blob.stream(), |
| { originalURL: file.url, originalSize: blob.size, etag: "" }, |
| ); |
| } |
| }, |
| }, |
| ); |
| signal.throwIfAborted(); |
| this.downloadSeconds = (performance.now() - start) / 1000; |
| this.onStatus({ |
| phase: "compile", |
| progress: 1, |
| label: "Preparing dedicated WebGPU kernels", |
| }); |
| const report = (event) => { |
| if (event.diagnostic) { |
| this.logs.push({ |
| level: event.level || "info", |
| line: event.diagnostic, |
| }); |
| if (this.logs.length > 100) this.logs.shift(); |
| } |
| this.onStatus(event); |
| }; |
| const { MiniCPMRuntime } = await import("./minicpm-runtime.js"); |
| this.dedicated = new MiniCPMRuntime(report); |
| const runtime = this.dedicated; |
| await startRuntime( |
| async () => { |
| await runtime.load(blobs.slice(0, 6), settings, signal); |
| if (settings.dspark !== "off") { |
| report({ |
| diagnostic: |
| "Preparing DSpark draft model and target verification", |
| level: "info", |
| }); |
| const { DSparkRuntime } = await import("./dspark-runtime.js"); |
| const draft = new DSparkRuntime(report); |
| runtime.draft = draft; |
| await draft.initialize(runtime, blobs[6], signal); |
| } |
| }, |
| { signal, stop: () => runtime.exit(), timeoutMs: 120000 }, |
| ); |
| this.info = { |
| n_ctx: settings.context, |
| n_layer: 42, |
| n_embd: 2048, |
| n_vocab: 130560, |
| }; |
| this.gpuLayers = 43; |
| this.ready = true; |
| this.settings = { ...settings }; |
| this.sessionFiles.clear(); |
| this.loadSeconds = (performance.now() - start) / 1000; |
| this.gpuSeconds = this.loadSeconds - this.downloadSeconds; |
| this.runtimeName = "MiniCPM5 dedicated WebGPU"; |
| this.runtimeCapabilities = { |
| subgroups: runtime.subgroups, |
| speculative: !!runtime.draft, |
| }; |
| report({ |
| diagnostic: `Dedicated WebGPU ready: ${runtime.subgroups ? "subgroup" : "portable"} kernels${runtime.draft ? ", DSpark enabled" : ""}. All 43 layers on the GPU.`, |
| level: "info", |
| }); |
| this.onStatus({ |
| phase: "ready", |
| label: "Model ready · dedicated WebGPU", |
| seconds: this.loadSeconds, |
| cached: cached.cached, |
| info: this.info, |
| }); |
| } catch (error) { |
| signal.aborted || this.loadController.abort(); |
| await this.dedicated?.exit().catch(() => {}); |
| this.dedicated = null; |
| this.ready = false; |
| throw error; |
| } finally { |
| this.loading = false; |
| } |
| } |
| async chat(messages, settings, tools, signal, onChunk) { |
| if (!this.ready) throw new Error("Load the model first."); |
| const prompt = this.dedicated.renderPrompt(messages, settings, tools), |
| splitter = { thinking: settings.thinking, pending: "" }; |
| let last; |
| await this.dedicated.generate(prompt, settings, signal, async (token) => { |
| if (token.progress) { |
| await onChunk({ choices: [], prompt_progress: token.progress }); |
| return; |
| } |
| last = token; |
| await onChunk({ |
| choices: [ |
| { |
| delta: splitStreamContent(splitter, token.delta), |
| finish_reason: token.finish, |
| }, |
| ], |
| timings: { |
| predicted_n: token.index, |
| predicted_per_second: |
| token.index > 1 |
| ? (token.index - 1) / |
| Math.max(0.001, (token.elapsed - token.prefillMs) / 1000) |
| : 0, |
| prompt_n: token.promptTokens, |
| prompt_ms: token.prefillMs, |
| }, |
| }); |
| }); |
| if (last) |
| await onChunk({ |
| choices: [ |
| { |
| delta: splitStreamContent(splitter, "", true), |
| finish_reason: last.finish || "stop", |
| }, |
| ], |
| usage: { completion_tokens: last.index }, |
| speculation: this.dedicated.specStats, |
| }); |
| } |
| async unload() { |
| try { |
| await this.dedicated?.exit(); |
| } finally { |
| this.ready = false; |
| this.dedicated = null; |
| } |
| } |
| async clearCache() { |
| await this.cache.deleteMany( |
| (f) => |
| f.metadata?.originalURL?.includes(MODEL.weights) || |
| f.metadata?.originalURL?.includes(DSPARK.repo) || |
| f.name.includes("MiniCPM5-2B-Q8_0") || |
| f.metadata?.originalURL?.endsWith("/__qa/dspark"), |
| ); |
| } |
| } |
|
|
