| import { |
| sampleVocabulary, |
| candidateDistribution, |
| drawDistribution, |
| verifySampledDraft, |
| randomStream, |
| } from "./sampling.js"; |
| import { TokenStream } from "./token-stream.js"; |
| import { Tokenizer } from "@huggingface/tokenizers"; |
| import { renderMiniCPMPrompt } from "./prompt.js"; |
| import { inspectGGUF } from "./gguf.js"; |
| import { |
| buildMiniCPMKernels, |
| shardVocabularyKernel, |
| H, |
| F, |
| VOCAB, |
| LAYERS, |
| KV, |
| MAX_CHUNKS, |
| PREFILL_BATCH, |
| } from "./minicpm-kernels.js"; |
| const STRIDE = 256, |
| DECODE_BATCH = 4, |
| EOS = new Set([1, 130073]); |
| export const layoutTypes = { |
| saveCandidates: ["r", "w"], |
| embed: ["r", "r", "w"], |
| norm: ["w", "r", "r", "w"], |
| qkv: ["r", "r", "w"], |
| output: ["r", "r", "w"], |
| down: ["r", "r", "w"], |
| head: ["r", "r", "w"], |
| upGate: ["r", "r", "w"], |
| qkvBatch: ["r", "r", "w"], |
| outputBatch: ["r", "r", "w"], |
| downBatch: ["r", "r", "w"], |
| upGateBatch: ["r", "r", "w"], |
| silu: ["r", "w"], |
| rope: ["r", "w", "w", "w"], |
| store: ["r", "r", "w", "w"], |
| attention: ["r", "r", "r", "w", "w"], |
| combine: ["r", "r", "w"], |
| sortFirst64: ["r", "w", "r", "u"], |
| sortFirst256: ["r", "w", "r", "u"], |
| sortMerge64: ["r", "w", "u"], |
| sortMerge256: ["r", "w", "u"], |
| sample: ["r", "w", "u"], |
| }; |
| export function cacheBytes(context, format) { |
| return ( |
| context * |
| KV * |
| ({ f32: 4, f16: 2, q8_0: 36 / 32, q4_0: 20 / 32 }[format] || 2) |
| ); |
| } |
| export class MiniCPMRuntime { |
| constructor(onStatus) { |
| this.onStatus = onStatus; |
| this.buffers = new Set(); |
| this.resident = []; |
| this.residentRun = 0; |
| this.pipelines = {}; |
| this.layouts = {}; |
| this.layers = []; |
| } |
| buffer(size, label = "", usage) { |
| const b = this.device.createBuffer({ |
| size: Math.ceil(size / 4) * 4, |
| label, |
| usage: |
| usage ?? |
| GPUBufferUsage.STORAGE | |
| GPUBufferUsage.COPY_SRC | |
| GPUBufferUsage.COPY_DST, |
| }); |
| this.buffers.add(b); |
| return b; |
| } |
| uniform(data, label) { |
| const b = this.buffer( |
| Math.max(16, data.byteLength), |
| label, |
| GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, |
| ); |
| this.device.queue.writeBuffer(b, 0, data); |
| return b; |
| } |
| group(name, buffers) { |
| return this.device.createBindGroup({ |
| layout: this.layouts[name], |
| entries: buffers.map((b, binding) => ({ |
| binding, |
| resource: b.buffer ? b : { buffer: b }, |
| })), |
| }); |
| } |
| async load(blobs, settings, signal) { |
| this.settings = settings; |
| this.maxContext = settings.context; |
| this.adapter = await navigator.gpu.requestAdapter({ |
| powerPreference: "high-performance", |
| }); |
| if (!this.adapter) throw new Error("No WebGPU adapter available."); |
| this.subgroups = |
| this.adapter.features.has("subgroups") && |
| (this.adapter.info?.subgroupMinSize || 0) >= 32 && |
| !( |
| import.meta.env.DEV && |
| new URLSearchParams(location.search).has("qa-portable") |
| ); |
| const limits = this.adapter.limits; |
| this.sharded = |
| limits.maxStorageBufferBindingSize < 284098560 || |
| limits.maxBufferSize < 284098560 || |
| (import.meta.env.DEV && |
| new URLSearchParams(location.search).has("qa-sharded")); |
| this.smallWorkgroups = |
| limits.maxComputeWorkgroupStorageSize < 25600 || |
| (import.meta.env.DEV && |
| new URLSearchParams(location.search).has("qa-small-workgroups")); |
| this.device = await this.adapter.requestDevice({ |
| requiredFeatures: this.subgroups ? ["subgroups"] : [], |
| requiredLimits: { |
| maxStorageBufferBindingSize: Math.min( |
| limits.maxStorageBufferBindingSize, |
| 536870912, |
| ), |
| maxBufferSize: Math.min(limits.maxBufferSize, 536870912), |
| maxComputeWorkgroupStorageSize: Math.min( |
| limits.maxComputeWorkgroupStorageSize, |
| 32768, |
| ), |
| }, |
| }); |
| this.device.lost.then((info) => { |
| this.failure = new Error("GPU device lost: " + info.message); |
| }); |
| this.device.addEventListener("uncapturederror", (event) => { |
| this.failure = new Error(event.error.message); |
| console.error(this.failure); |
| }); |
| this.stepLayout = this.device.createBindGroupLayout({ |
| entries: [ |
| { |
| binding: 0, |
| visibility: GPUShaderStage.COMPUTE, |
| buffer: { |
| type: "uniform", |
| hasDynamicOffset: true, |
| minBindingSize: 16, |
| }, |
| }, |
| ], |
| }); |
| const kernels = buildMiniCPMKernels({ |
| subgroups: this.subgroups, |
| kv: settings.cacheK, |
| smallWorkgroups: this.smallWorkgroups, |
| }); |
| if (this.sharded) |
| for (const name of ["embed", "head", "headBlock"]) |
| kernels[name] = shardVocabularyKernel(kernels[name], name === "embed"); |
| const tokenizerPromise = this.loadTokenizer(); |
| await Promise.all( |
| Object.entries(kernels).map(async ([name, code]) => { |
| const module = this.device.createShaderModule({ code, label: name }); |
| const info = await module.getCompilationInfo(); |
| const errors = info.messages.filter((m) => m.type === "error"); |
| if (errors.length) |
| throw new Error( |
| name + |
| ": " + |
| errors.map((e) => e.message + " at " + e.lineNum).join("; "), |
| ); |
| const layout = this.device.createBindGroupLayout({ |
| entries: this.layoutFor(name).map((type, binding) => ({ |
| binding, |
| visibility: GPUShaderStage.COMPUTE, |
| buffer: { |
| type: |
| type === "u" |
| ? "uniform" |
| : type === "w" |
| ? "storage" |
| : "read-only-storage", |
| }, |
| })), |
| }); |
| this.layouts[name] = layout; |
| this.pipelines[name] = await this.device.createComputePipelineAsync({ |
| label: name, |
| layout: this.device.createPipelineLayout({ |
| bindGroupLayouts: [this.stepLayout, layout], |
| }), |
| compute: { module, entryPoint: "main" }, |
| }); |
| }), |
| ); |
| const tensors = new Map(); |
| for (const blob of blobs) { |
| signal?.throwIfAborted(); |
| const header = await inspectGGUF(blob); |
| for (const tensor of header.tensors) { |
| if (tensors.has(tensor.name)) |
| throw new Error("Duplicate weight " + tensor.name); |
| tensors.set(tensor.name, { ...tensor, blob }); |
| } |
| } |
| if (tensors.size !== 381) |
| throw new Error("MiniCPM5 requires all 381 original weight tensors."); |
| const upload = async (names, label) => { |
| const list = names.map((name) => { |
| const t = tensors.get(name); |
| if (!t) throw new Error("Missing tensor: " + name); |
| return t; |
| }); |
| const b = this.buffer( |
| list.reduce((n, t) => n + t.bytes, 0), |
| label, |
| ); |
| let offset = 0; |
| for (const t of list) { |
| for (let start = 0; start < t.bytes; start += 16 * 1024 * 1024) { |
| signal?.throwIfAborted(); |
| const bytes = await t.blob |
| .slice( |
| t.offset + start, |
| t.offset + Math.min(t.bytes, start + 16 * 1024 * 1024), |
| ) |
| .arrayBuffer(); |
| this.device.queue.writeBuffer(b, offset + start, bytes); |
| } |
| offset += t.bytes; |
| } |
| return b; |
| }; |
| const w = (this.work = {}); |
| for (const [name, size] of Object.entries({ |
| x: PREFILL_BATCH * H, |
| norm: PREFILL_BATCH * H, |
| temp: PREFILL_BATCH * H, |
| zeros: PREFILL_BATCH * H, |
| qkv: PREFILL_BATCH * (H + KV * 2), |
| q: PREFILL_BATCH * H, |
| k: PREFILL_BATCH * KV, |
| v: PREFILL_BATCH * KV, |
| attention: PREFILL_BATCH * H, |
| upGate: PREFILL_BATCH * F * 2, |
| hidden: PREFILL_BATCH * F, |
| partial: PREFILL_BATCH * 16 * MAX_CHUNKS * 128, |
| md: PREFILL_BATCH * 16 * MAX_CHUNKS * 2, |
| logits: VOCAB, |
| tokens: this.maxContext + DECODE_BATCH + 1, |
| sortA: 128 * 256 * 2, |
| sortB: 128 * 256 * 2, |
| })) |
| w[name] = this.buffer(size * 4, name); |
| w.candidates = this.buffer(15 * 256 * 8, "speculative candidates"); |
| w.candidatesRead = this.buffer( |
| 15 * 256 * 8, |
| "candidate readback", |
| GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, |
| ); |
| w.verifyLogits = this.buffer(8 * VOCAB * 4, "verification logits"); |
| w.verifyRead = this.buffer( |
| 64, |
| "verification readback", |
| GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, |
| ); |
| w.sampleSteps = this.buffer( |
| (this.maxContext + 9) * STRIDE, |
| "sampler steps", |
| GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, |
| ); |
| w.steps = this.buffer( |
| (this.maxContext + DECODE_BATCH + 2) * STRIDE, |
| "steps", |
| GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, |
| ); |
| w.read = this.buffer( |
| 4 * DECODE_BATCH, |
| "sampled tokens", |
| GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, |
| ); |
| w.logitsRead = this.buffer( |
| VOCAB * 4, |
| "logits readback", |
| GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, |
| ); |
| w.penalty = this.uniform(new Float32Array([1, 0, 0, 0]), "repeat penalty"); |
| w.sampler = this.uniform(new Float32Array(4), "sampler"); |
| this.sampleStepGroup = this.device.createBindGroup({ |
| layout: this.stepLayout, |
| entries: [{ binding: 0, resource: { buffer: w.sampleSteps, size: 16 } }], |
| }); |
| this.stepGroup = this.device.createBindGroup({ |
| layout: this.stepLayout, |
| entries: [{ binding: 0, resource: { buffer: w.steps, size: 16 } }], |
| }); |
| const uploadVocab = async (name, label) => { |
| if (!this.sharded) return [await upload([name], label)]; |
| const t = tensors.get(name); |
| if (t.bytes !== 284098560) |
| throw new Error("Unexpected vocabulary weight size."); |
| const parts = []; |
| for (let part = 0; part < 4; part++) { |
| const size = 71024640, |
| b = this.buffer(size, label + " " + part); |
| for (let at = 0; at < size; at += 16 * 1024 * 1024) { |
| signal?.throwIfAborted(); |
| const start = t.offset + part * size + at; |
| this.device.queue.writeBuffer( |
| b, |
| at, |
| await t.blob |
| .slice(start, start + Math.min(size - at, 16 * 1024 * 1024)) |
| .arrayBuffer(), |
| ); |
| } |
| parts.push(b); |
| } |
| return parts; |
| }; |
| this.embedding = await uploadVocab("token_embd.weight", "embeddings"); |
| this.head = await uploadVocab("output.weight", "output weights"); |
| this.finalNorm = await upload(["output_norm.weight"], "output norm"); |
| for (let i = 0; i < LAYERS; i++) { |
| const prefix = "blk." + i + "."; |
| const get = (names, label) => |
| upload( |
| names.map((n) => prefix + n + ".weight"), |
| label + " " + i, |
| ); |
| const [norm, post, qkv, o, upGate, down] = await Promise.all([ |
| get(["attn_norm"], "norm"), |
| get(["ffn_norm"], "post norm"), |
| get(["attn_q", "attn_k", "attn_v"], "qkv"), |
| get(["attn_output"], "attention output"), |
| get(["ffn_gate", "ffn_up"], "gate/up"), |
| get(["ffn_down"], "down"), |
| ]); |
| const k = this.buffer( |
| cacheBytes(settings.context, settings.cacheK), |
| "key cache " + i, |
| ), |
| v = this.buffer( |
| cacheBytes(settings.context, settings.cacheK), |
| "value cache " + i, |
| ); |
| const groups = { |
| input: this.group("norm", [w.x, i ? w.temp : w.zeros, norm, w.norm]), |
| post: this.group("norm", [w.x, w.temp, post, w.norm]), |
| qkv: this.group("qkv", [w.norm, qkv, w.qkv]), |
| qkvBatch: this.group("qkvBatch", [w.norm, qkv, w.qkv]), |
| rope: this.group("rope", [w.qkv, w.q, w.k, w.v]), |
| store: this.group("store", [w.k, w.v, k, v]), |
| attention: this.group("attention", [w.q, k, v, w.partial, w.md]), |
| combine: this.group("combine", [w.partial, w.md, w.attention]), |
| output: this.group("output", [w.attention, o, w.temp]), |
| outputBatch: this.group("outputBatch", [w.attention, o, w.temp]), |
| upGate: this.group("upGate", [w.norm, upGate, w.hidden]), |
| upGateBatch: this.group("upGateBatch", [w.norm, upGate, w.upGate]), |
| silu: this.group("silu", [w.upGate, w.hidden]), |
| down: this.group("down", [w.hidden, down, w.temp]), |
| downBatch: this.group("downBatch", [w.hidden, down, w.temp]), |
| }; |
| this.layers.push({ index: i, groups, k, v }); |
| await this.device.queue.onSubmittedWorkDone(); |
| this.check(); |
| this.onStatus?.({ |
| diagnostic: `Dedicated WebGPU: layer ${i + 1} / ${LAYERS} uploaded`, |
| level: "info", |
| }); |
| } |
| this.groups = { |
| embed: this.group("embed", [...this.embedding, w.tokens, w.x]), |
| finalNorm: this.group("norm", [w.x, w.temp, this.finalNorm, w.norm]), |
| }; |
| this.headGroups = Array.from({ length: PREFILL_BATCH }, (_, i) => |
| this.group("head", [ |
| { buffer: w.norm, offset: i * H * 4, size: H * 4 }, |
| ...this.head, |
| w.logits, |
| ]), |
| ); |
| this.verifyHead = this.group("headBlock", [ |
| w.norm, |
| ...this.head, |
| w.verifyLogits, |
| ]); |
| this.sortGroups = {}; |
| for (const keep of [64, 256]) { |
| const rounds = []; |
| let count = 128 * keep, |
| source = w.sortA, |
| target = w.sortB; |
| while (count > keep) { |
| const n = Math.ceil(count / 1024); |
| rounds.push({ |
| n, |
| group: this.group("sortMerge" + keep, [ |
| source, |
| target, |
| this.uniform(new Uint32Array([count, 0, 0, 0]), "sort size"), |
| ]), |
| }); |
| count = n * keep; |
| [source, target] = [target, source]; |
| } |
| this.sortGroups[keep] = { |
| source, |
| save: Array.from({ length: 15 }, (_, i) => |
| this.group("saveCandidates", [ |
| source, |
| { buffer: w.candidates, offset: i * 256 * 8, size: 256 * 8 }, |
| ]), |
| ), |
| verify: Array.from({ length: 8 }, (_, i) => |
| this.group("sortFirst" + keep, [ |
| { buffer: w.verifyLogits, offset: i * VOCAB * 4, size: VOCAB * 4 }, |
| w.sortA, |
| w.tokens, |
| w.penalty, |
| ]), |
| ), |
| first: this.group("sortFirst" + keep, [ |
| w.logits, |
| w.sortA, |
| w.tokens, |
| w.penalty, |
| ]), |
| rounds, |
| sample: this.group("sample", [source, w.tokens, w.sampler]), |
| }; |
| } |
| await tokenizerPromise; |
| await this.device.queue.onSubmittedWorkDone(); |
| this.check(); |
| await this.warmup(); |
| this.ready = true; |
| } |
| layoutFor(name) { |
| return this.sharded && ["embed", "head", "headBlock"].includes(name) |
| ? ["r", "r", "r", "r", "r", "w"] |
| : layoutTypes[name] || layoutTypes[name.replace(/Block$/, "")]; |
| } |
| async warmup() { |
| this.sampleK = 40; |
| this.device.queue.writeBuffer(this.work.tokens, 0, new Uint32Array([0])); |
| const values = new ArrayBuffer(16), |
| view = new DataView(values); |
| view.setFloat32(4, 0.95, true); |
| view.setUint32(8, 40, true); |
| this.device.queue.writeBuffer(this.work.sampler, 0, values); |
| for (let round = 0; round < 2; round++) { |
| const start = performance.now(), |
| encoder = this.device.createCommandEncoder(); |
| for (let i = 0; i < 4; i++) this.encode(encoder, i, 1, true); |
| await this.readTokens(encoder, 1, 4); |
| this.baselineMs = (performance.now() - start) / 4; |
| } |
| } |
| async loadTokenizer() { |
| const [json, config] = await Promise.all( |
| ["tokenizer.json", "tokenizer_config.json"].map(async (name) => { |
| const r = await fetch(new URL("../tokenizer/" + name, import.meta.url)); |
| if (!r.ok) throw new Error("Could not load tokenizer."); |
| return r.json(); |
| }), |
| ); |
| this.tokenizer = new Tokenizer(json, config); |
| this.tokenizerConfig = config; |
| } |
| check() { |
| if (this.failure) throw this.failure; |
| } |
| prepareStep(position, count = 1) { |
| const tokens = position + count; |
| let chunks = Math.min(MAX_CHUNKS, Math.max(1, Math.ceil(tokens / 64))); |
| const chunk = Math.ceil(Math.ceil(tokens / chunks) / 4) * 4; |
| chunks = Math.ceil(tokens / chunk); |
| this.device.queue.writeBuffer( |
| this.work.steps, |
| position * STRIDE, |
| new Uint32Array([position, count, chunk, chunks]), |
| ); |
| return { chunk, chunks }; |
| } |
| dispatch(pass, name, group, x = 1, y = 1, z = 1) { |
| pass.setPipeline(this.pipelines[name]); |
| pass.setBindGroup(1, group); |
| pass.dispatchWorkgroups(x, y, z); |
| } |
| encode(encoder, position, count = 1, sample = true, updateContext = true) { |
| const plan = this.prepareStep(position, count), |
| pass = encoder.beginComputePass(); |
| pass.setBindGroup(0, this.stepGroup, [position * STRIDE]); |
| const run = (name, group, x = 1, y = 1, z = 1) => |
| this.dispatch(pass, name, group, x, y, z); |
| const batch = count > 1, |
| block = batch && count <= 8, |
| suffix = block ? "Block" : batch ? "Batch" : "", |
| divisor = batch && !block ? 32 : 8; |
| run("embed", this.groups.embed, H / 256, count); |
| for (const { index, groups: g } of this.layers) { |
| run("norm", g.input, count); |
| run("qkv" + suffix, batch ? g.qkvBatch : g.qkv, (H + KV * 2) / divisor); |
| run("rope", g.rope, (H + KV) / 256, count); |
| run("store", g.store, KV / 32, count, 2); |
| run("attention", g.attention, plan.chunks, 16, count); |
| run("combine", g.combine, 16, count); |
| run("output" + suffix, batch ? g.outputBatch : g.output, H / divisor); |
| run("norm", g.post, count); |
| run( |
| "upGate" + suffix, |
| batch ? g.upGateBatch : g.upGate, |
| batch ? (F * 2) / divisor : F / 8, |
| ); |
| if (batch) run("silu", g.silu, F / 256, count); |
| run("down" + suffix, batch ? g.downBatch : g.down, H / divisor); |
| if (this.draftActive) this.draft.capture(pass, index, count); |
| } |
| if (sample) { |
| run("norm", this.groups.finalNorm, count); |
| run("head", this.headGroups[count - 1], VOCAB / 8); |
| if (!this.fullSampler) { |
| this.prepareSample(pass, position + count - 1); |
| this.encodeSample(pass); |
| } |
| } |
| pass.end(); |
| if (this.draftActive && updateContext) |
| this.draft.updateContext(encoder, position, count); |
| } |
| prepareSample(pass, position) { |
| this.device.queue.writeBuffer( |
| this.work.sampleSteps, |
| position * STRIDE, |
| new Uint32Array([position, 1, 0, 1]), |
| ); |
| pass.setBindGroup(0, this.sampleStepGroup, [position * STRIDE]); |
| } |
| encodeSample( |
| pass, |
| first, |
| g = this.sortGroups[this.sampleK <= 64 ? 64 : 256], |
| sample = true, |
| ) { |
| const keep = this.sampleK <= 64 ? 64 : 256; |
| this.dispatch(pass, "sortFirst" + keep, first || g.first, 128); |
| for (const round of g.rounds) |
| this.dispatch(pass, "sortMerge" + keep, round.group, round.n); |
| if (sample) this.dispatch(pass, "sample", g.sample); |
| } |
| async speculate(position) { |
| const started = performance.now(); |
| const encoder = this.device.createCommandEncoder(); |
| this.draft.propose(encoder, position); |
| this.encode(encoder, position, 8, false, false); |
| const pass = encoder.beginComputePass(); |
| pass.setBindGroup(0, this.stepGroup, [position * STRIDE]); |
| this.dispatch(pass, "norm", this.groups.finalNorm, 8); |
| this.dispatch(pass, "headBlock", this.verifyHead, VOCAB / 8); |
| const groups = this.sortGroups[this.sampleK <= 64 ? 64 : 256]; |
| for (let i = 0; i < 8; i++) { |
| this.prepareSample(pass, position + i); |
| this.encodeSample( |
| pass, |
| groups.verify[i], |
| groups, |
| this.samplingSettings.temperature <= 0, |
| ); |
| if (this.samplingSettings.temperature > 0) |
| this.dispatch( |
| pass, |
| "saveCandidates", |
| groups.save[i], |
| this.sampleK <= 64 ? 1 : 4, |
| ); |
| } |
| pass.end(); |
| encoder.copyBufferToBuffer( |
| this.work.tokens, |
| (position + 1) * 4, |
| this.work.verifyRead, |
| 0, |
| 32, |
| ); |
| encoder.copyBufferToBuffer( |
| this.draft.work.tokens, |
| (position + 1) * 4, |
| this.work.verifyRead, |
| 32, |
| 28, |
| ); |
| if (this.samplingSettings.temperature > 0) |
| encoder.copyBufferToBuffer( |
| this.work.candidates, |
| 0, |
| this.work.candidatesRead, |
| 0, |
| 15 * 256 * 8, |
| ); |
| this.device.queue.submit([encoder.finish()]); |
| await this.work.verifyRead.mapAsync(GPUMapMode.READ); |
| const data = new Uint32Array(this.work.verifyRead.getMappedRange()).slice(); |
| this.work.verifyRead.unmap(); |
| this.check(); |
| let accepted = 0, |
| result; |
| if (this.samplingSettings.temperature > 0) { |
| await this.work.candidatesRead.mapAsync(GPUMapMode.READ); |
| const words = new Uint32Array(this.work.candidatesRead.getMappedRange()); |
| result = []; |
| for (let i = 0; i < 8; i++) { |
| const p = candidateDistribution( |
| words.subarray(i * 512, (i + 1) * 512), |
| this.samplingSettings, |
| ), |
| random = (n) => |
| randomStream(this.samplingSettings.seed, position + i, n); |
| if (i === 7) { |
| result.push(drawDistribution(p, random(2))); |
| break; |
| } |
| const q = candidateDistribution( |
| words.subarray((8 + i) * 512, (9 + i) * 512), |
| this.samplingSettings, |
| ), |
| check = verifySampledDraft(p, q, data[8 + i], random(1), random(2)); |
| result.push(check.token); |
| if (!check.accepted) break; |
| accepted++; |
| } |
|
|
| this.work.candidatesRead.unmap(); |
| this.device.queue.writeBuffer( |
| this.work.tokens, |
| (position + 1) * 4, |
| new Uint32Array(result), |
| ); |
| } else { |
| while (accepted < 7 && data[accepted] === data[8 + accepted]) accepted++; |
| result = Array.from(data.slice(0, accepted + 1)); |
| } |
| const count = accepted + 1; |
| const context = this.device.createCommandEncoder(); |
| this.draft.updateContext(context, position, count); |
| this.device.queue.submit([context.finish()]); |
| this.specStats.steps++; |
| this.specStats.accepted += accepted; |
| this.specStats.milliseconds += performance.now() - started; |
| if ( |
| this.settings.dspark !== "on" && |
| ((this.specStats.steps >= 3 && |
| this.specStats.milliseconds / |
| (this.specStats.steps + this.specStats.accepted) > |
| this.baselineMs * 1.5) || |
| (this.specStats.steps >= 6 && |
| this.specStats.milliseconds / |
| (this.specStats.steps + this.specStats.accepted) > |
| this.baselineMs * 1.1)) |
| ) { |
| this.draftActive = false; |
| this.specStats.fallback = true; |
| } |
| return result; |
| } |
| async readTokens(encoder, from, count) { |
| if (this.fullSampler) { |
| encoder.copyBufferToBuffer( |
| this.work.logits, |
| 0, |
| this.work.logitsRead, |
| 0, |
| VOCAB * 4, |
| ); |
| this.device.queue.submit([encoder.finish()]); |
| await this.work.logitsRead.mapAsync(GPUMapMode.READ); |
| const values = new Float32Array(this.work.logitsRead.getMappedRange()); |
| const token = sampleVocabulary( |
| values, |
| this.samplingSettings, |
| this.resident, |
| from - 1, |
| ); |
| this.work.logitsRead.unmap(); |
| this.check(); |
| this.device.queue.writeBuffer( |
| this.work.tokens, |
| from * 4, |
| new Uint32Array([token]), |
| ); |
| return [token]; |
| } |
| encoder.copyBufferToBuffer( |
| this.work.tokens, |
| from * 4, |
| this.work.read, |
| 0, |
| count * 4, |
| ); |
| this.device.queue.submit([encoder.finish()]); |
| await this.work.read.mapAsync(GPUMapMode.READ, 0, count * 4); |
| const ids = Array.from( |
| new Uint32Array(this.work.read.getMappedRange(0, count * 4)), |
| ); |
| this.work.read.unmap(); |
| this.check(); |
| return ids; |
| } |
| renderPrompt(messages, settings, tools) { |
| return renderMiniCPMPrompt(messages, { |
| thinking: settings.thinking, |
| tools, |
| }); |
| } |
|
|
| |
| |
| async generate(prompt, settings, signal, onToken, onProgress) { |
| this.check(); |
| const ids = this.tokenizer.encode(prompt, { |
| add_special_tokens: false, |
| }).ids; |
| if (ids.length + 1 >= this.maxContext) |
| throw new Error( |
| `The prompt needs ${ids.length} tokens and exceeds the ${this.maxContext}-token context window.`, |
| ); |
| this.fullSampler = settings.topK === 0 && settings.temperature > 0; |
| this.sampleK = settings.topK || 1; |
| this.samplingSettings = { |
| ...settings, |
| seed: |
| settings.seed < 0 |
| ? crypto.getRandomValues(new Uint32Array(1))[0] |
| : settings.seed, |
| }; |
| const config = new ArrayBuffer(16), |
| v = new DataView(config); |
| v.setFloat32(0, settings.temperature, true); |
| v.setFloat32(4, settings.topP, true); |
| v.setUint32(8, this.sampleK, true); |
| v.setUint32(12, this.samplingSettings.seed, true); |
| this.device.queue.writeBuffer(this.work.sampler, 0, config); |
| this.device.queue.writeBuffer( |
| this.work.penalty, |
| 0, |
| new Float32Array([settings.repeatPenalty, 0, 0, 0]), |
| ); |
| if (this.draft) { |
| v.setUint32(12, (this.samplingSettings.seed ^ 0x9e3779b9) >>> 0, true); |
| this.device.queue.writeBuffer(this.draft.work.sampler, 0, config); |
| } |
| this.specStats = { steps: 0, accepted: 0, milliseconds: 0 }; |
| this.draftActive = |
| !this.fullSampler && !!this.draft && settings.dspark !== "off"; |
| this.settings.dspark = settings.dspark || "auto"; |
| let reuse = 0; |
| while ( |
| reuse < Math.min(ids.length - 1, this.residentRun) && |
| ids[reuse] === this.resident[reuse] |
| ) |
| reuse++; |
| if (this.draftActive) reuse = Math.min(reuse, this.draft.contextRun || 0); |
| this.device.queue.writeBuffer(this.work.tokens, 0, new Uint32Array(ids)); |
| this.resident = [...ids]; |
| const start = performance.now(); |
| let ready = []; |
| for (let p = reuse; p < ids.length; ) { |
| signal?.throwIfAborted(); |
| const count = Math.min( |
| settings.performance === "balanced" ? 8 : PREFILL_BATCH, |
| ids.length - p, |
| ), |
| last = p + count === ids.length, |
| encoder = this.device.createCommandEncoder(); |
| this.encode(encoder, p, count, last); |
| if (last) ready = await this.readTokens(encoder, ids.length, 1); |
| else { |
| this.device.queue.submit([encoder.finish()]); |
| await this.device.queue.onSubmittedWorkDone(); |
| this.check(); |
| } |
| p += count; |
| onProgress?.({ total: ids.length, processed: p, cache: reuse }); |
| await onToken({ |
| progress: { total: ids.length, processed: p, cache: reuse }, |
| }); |
| } |
| this.residentRun = ids.length; |
| const prefillMs = performance.now() - start; |
| const decoder = new TokenStream(this.tokenizer); |
| let text = "", |
| position = ids.length; |
| const max = Math.min(settings.maxTokens, this.maxContext - ids.length); |
| for (let emitted = 0; emitted < max; ) { |
| signal?.throwIfAborted(); |
| if ( |
| !ready.length && |
| this.draftActive && |
| position + 8 < this.maxContext && |
| max - emitted >= 8 |
| ) { |
| ready = await this.speculate(position); |
| this.residentRun = position + ready.length; |
| position += ready.length; |
| } |
| if (!ready.length) { |
| const n = Math.min(this.fullSampler ? 1 : DECODE_BATCH, max - emitted), |
| encoder = this.device.createCommandEncoder(); |
| for (let i = 0; i < n; i++) this.encode(encoder, position + i, 1, true); |
| ready = await this.readTokens(encoder, position + 1, n); |
| this.residentRun = position + n; |
| position += n; |
| } |
| const token = ready.shift(); |
| this.resident.push(token); |
| if (EOS.has(token)) break; |
| emitted++; |
| const delta = decoder.push(token); |
| text += delta; |
| const elapsed = performance.now() - start; |
| await onToken({ |
| token, |
| delta, |
| text, |
| index: emitted, |
| prefillMs, |
| elapsed, |
| promptTokens: ids.length, |
| cacheTokens: reuse, |
| finish: emitted >= max ? "length" : null, |
| }); |
| } |
| } |
| async exit() { |
| this.ready = false; |
| if (this.draft) { |
| for (const b of this.draft.buffers) b.destroy(); |
| this.draft.buffers.clear(); |
| } |
| for (const b of this.buffers) b.destroy(); |
| this.buffers.clear(); |
| this.device?.destroy(); |
| } |
| } |
|
|
