| import { MiniCPMRuntime, layoutTypes, cacheBytes } from "./minicpm-runtime.js"; |
| import { buildDSparkKernels } from "./dspark-kernels.js"; |
| import { |
| H, |
| F, |
| KV, |
| VOCAB, |
| MAX_CHUNKS, |
| PREFILL_BATCH, |
| } from "./minicpm-kernels.js"; |
| export const DSPARK_LAYERS = [1, 10, 20, 30, 39]; |
| const extraTypes = { |
| projectBatch: ["r", "r", "w"], |
| projectBlock: ["r", "r", "w"], |
| contextBatch: ["r", "r", "w"], |
| contextBlock: ["r", "r", "w"], |
| capture: ["r", "r", "w", "u"], |
| noiseNormRope: ["r", "r", "r", "w", "w", "w"], |
| contextNormRope: ["r", "r", "r", "w", "w", "w"], |
| noise: ["r", "w"], |
| markovEmbed: ["r", "r", "w"], |
| markov: ["r", "r", "w"], |
| addBias: ["r", "r", "w"], |
| argmax: ["r", "w"], |
| argmaxFinal: ["r", "w"], |
| }; |
| export class DSparkRuntime extends MiniCPMRuntime { |
| async initialize(parent, blob, signal) { |
| this.parent = parent; |
| this.device = parent.device; |
| this.adapter = parent.adapter; |
| this.subgroups = parent.subgroups; |
| this.maxContext = parent.maxContext; |
| this.stepLayout = parent.stepLayout; |
| this.settings = parent.settings; |
| const all = buildDSparkKernels({ |
| subgroups: this.subgroups, |
| kv: this.settings.cacheK, |
| }); |
| const needed = [ |
| "norm", |
| "qkvBlock", |
| "outputBlock", |
| "upGateBlock", |
| "downBlock", |
| "silu", |
| "projectBatch", |
| "projectBlock", |
| "contextBatch", |
| "contextBlock", |
| "markov", |
| "store", |
| "attention", |
| "combine", |
| ...Object.keys(extraTypes), |
| ]; |
| await Promise.all( |
| [...new Set(needed)].map(async (name) => { |
| const module = this.device.createShaderModule({ |
| code: all[name], |
| label: "dspark " + name, |
| }); |
| const info = await module.getCompilationInfo(); |
| const errors = info.messages.filter((m) => m.type === "error"); |
| if (errors.length) |
| throw new Error( |
| "DSpark " + |
| name + |
| ": " + |
| errors.map((e) => e.message + " at " + e.lineNum).join("; "), |
| ); |
| const types = |
| extraTypes[name] || |
| layoutTypes[name] || |
| layoutTypes[name.replace(/Block$/, "")]; |
| const layout = this.device.createBindGroupLayout({ |
| entries: types.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({ |
| layout: this.device.createPipelineLayout({ |
| bindGroupLayouts: [this.stepLayout, layout], |
| }), |
| compute: { module, entryPoint: "main" }, |
| label: "DSpark " + name, |
| }); |
| }), |
| ); |
| const headSize = Number( |
| new DataView(await blob.slice(0, 8).arrayBuffer()).getBigUint64(0, true), |
| ); |
| if (headSize > 1024 * 1024) |
| throw new Error("Invalid DSpark weights header."); |
| const header = JSON.parse(await blob.slice(8, 8 + headSize).text()); |
| const offset = 8 + headSize; |
| const upload = async (names, label, fp32 = false) => { |
| const list = []; |
| for (const name of names) { |
| signal?.throwIfAborted(); |
| const item = header[name]; |
| if (!item || item.dtype !== "BF16") |
| throw new Error("Missing BF16 DSpark tensor: " + name); |
| const bytes = await blob |
| .slice(offset + item.data_offsets[0], offset + item.data_offsets[1]) |
| .arrayBuffer(); |
| if (fp32) { |
| const source = new Uint16Array(bytes), |
| out = new Uint32Array(source.length); |
| for (let i = 0; i < source.length; i++) out[i] = source[i] << 16; |
| list.push(out.buffer); |
| } else list.push(bytes); |
| } |
| const b = this.buffer( |
| list.reduce((n, a) => n + Math.ceil(a.byteLength / 4) * 4, 0), |
| label, |
| ); |
| let at = 0; |
| for (const bytes of list) { |
| let data = bytes; |
| if (data.byteLength % 4) { |
| const padded = new Uint8Array(Math.ceil(data.byteLength / 4) * 4); |
| padded.set(new Uint8Array(data)); |
| data = padded.buffer; |
| } |
| this.device.queue.writeBuffer(b, at, data); |
| at += data.byteLength; |
| } |
| return b; |
| }; |
| const w = (this.work = {}); |
| for (const [name, count] of Object.entries({ |
| features: PREFILL_BATCH * H * 5, |
| projected: PREFILL_BATCH * H, |
| context: PREFILL_BATCH * H, |
| contextKV: PREFILL_BATCH * KV * 2, |
| x: PREFILL_BATCH * H, |
| norm: PREFILL_BATCH * H, |
| temp: PREFILL_BATCH * H, |
| zeros: PREFILL_BATCH * H, |
| qkv: PREFILL_BATCH * (H + 2 * KV), |
| 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, |
| tokens: this.maxContext + 9, |
| baseLogits: 8 * VOCAB, |
| bias: VOCAB, |
| logits: VOCAB, |
| markov: 256, |
| argmax: 256, |
| })) |
| w[name] = this.buffer(count * 4, "DSpark " + name); |
| w.sampler = this.uniform(new Float32Array(4), "draft sampler"); |
| w.steps = this.buffer( |
| (this.maxContext + 9) * 256, |
| "DSpark steps", |
| GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, |
| ); |
| w.sampleSteps = this.buffer( |
| (this.maxContext + 9) * 256, |
| "DSpark sample steps", |
| GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, |
| ); |
| this.stepGroup = this.device.createBindGroup({ |
| layout: this.stepLayout, |
| entries: [{ binding: 0, resource: { buffer: w.steps, size: 16 } }], |
| }); |
| this.sampleStepGroup = this.device.createBindGroup({ |
| layout: this.stepLayout, |
| entries: [{ binding: 0, resource: { buffer: w.sampleSteps, size: 16 } }], |
| }); |
| const [fc, hiddenNorm, finalNorm, w1, w2, confidence, confidenceBias] = |
| await Promise.all([ |
| upload(["fc.weight"], "DSpark projection"), |
| upload(["hidden_norm.weight"], "DSpark hidden norm", true), |
| upload(["norm.weight"], "DSpark output norm", true), |
| upload(["markov_head.markov_w1.weight"], "DSpark Markov embeddings"), |
| upload(["markov_head.markov_w2.weight"], "DSpark Markov projection"), |
| upload(["confidence_head.proj.weight"], "DSpark confidence", true), |
| upload(["confidence_head.proj.bias"], "DSpark confidence bias", true), |
| ]); |
| this.confidence = { weight: confidence, bias: confidenceBias }; |
| this.groups = { |
| projectBatch: this.group("projectBatch", [w.features, fc, w.projected]), |
| projectBlock: this.group("projectBlock", [w.features, fc, w.projected]), |
| hiddenNorm: this.group("norm", [ |
| w.projected, |
| w.zeros, |
| hiddenNorm, |
| w.context, |
| ]), |
| noise: this.group("noise", [parent.work.tokens, w.tokens]), |
| embed: parent.group("embed", [...parent.embedding, w.tokens, w.x]), |
| finalNorm: this.group("norm", [w.x, w.temp, finalNorm, w.norm]), |
| head: parent.group("headBlock", [w.norm, ...parent.head, w.baseLogits]), |
| markovEmbed: this.group("markovEmbed", [w1, w.tokens, w.markov]), |
| markov: this.group("markov", [w.markov, w2, w.bias]), |
| argmax: this.group("argmax", [w.logits, w.argmax]), |
| argmaxFinal: this.group("argmaxFinal", [w.argmax, w.tokens]), |
| }; |
| this.biasGroups = Array.from({ length: 7 }, (_, i) => |
| this.group("addBias", [ |
| { buffer: w.baseLogits, offset: i * VOCAB * 4, size: VOCAB * 4 }, |
| w.bias, |
| w.logits, |
| ]), |
| ); |
| this.captureGroups = DSPARK_LAYERS.map((_, i) => |
| this.group("capture", [ |
| parent.work.x, |
| parent.work.temp, |
| w.features, |
| this.uniform(new Uint32Array([i, 0, 0, 0]), "capture layer"), |
| ]), |
| ); |
| for (let i = 0; i < 5; i++) { |
| const prefix = "layers." + i + "."; |
| const get = (names, label, float = false) => |
| upload( |
| names.map((n) => prefix + n + ".weight"), |
| label + " " + i, |
| float, |
| ); |
| const [norm, post, qkv, contextKV, o, upGate, down, qNorm, kNorm] = |
| await Promise.all([ |
| get(["input_layernorm"], "DSpark norm", true), |
| get(["post_attention_layernorm"], "DSpark post norm", true), |
| get( |
| ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], |
| "DSpark qkv", |
| ), |
| get(["self_attn.k_proj", "self_attn.v_proj"], "DSpark context KV"), |
| get(["self_attn.o_proj"], "DSpark attention out"), |
| get(["mlp.gate_proj", "mlp.up_proj"], "DSpark gate/up"), |
| get(["mlp.down_proj"], "DSpark down"), |
| get(["self_attn.q_norm"], "DSpark Q norm", true), |
| get(["self_attn.k_norm"], "DSpark K norm", true), |
| ]); |
| const k = this.buffer( |
| cacheBytes(this.maxContext, this.settings.cacheK), |
| "DSpark key cache", |
| ), |
| v = this.buffer( |
| cacheBytes(this.maxContext, this.settings.cacheK), |
| "DSpark value cache", |
| ); |
| 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("qkvBlock", [w.norm, qkv, w.qkv]), |
| contextBatch: this.group("contextBatch", [ |
| w.context, |
| contextKV, |
| w.contextKV, |
| ]), |
| contextBlock: this.group("contextBlock", [ |
| w.context, |
| contextKV, |
| w.contextKV, |
| ]), |
| noiseRope: this.group("noiseNormRope", [ |
| w.qkv, |
| qNorm, |
| kNorm, |
| w.q, |
| w.k, |
| w.v, |
| ]), |
| contextRope: this.group("contextNormRope", [ |
| w.contextKV, |
| qNorm, |
| kNorm, |
| 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("outputBlock", [w.attention, o, w.temp]), |
| upGate: this.group("upGateBlock", [w.norm, upGate, w.upGate]), |
| silu: this.group("silu", [w.upGate, w.hidden]), |
| down: this.group("downBlock", [w.hidden, down, w.temp]), |
| }; |
| this.layers.push({ groups, k, v }); |
| await this.device.queue.onSubmittedWorkDone(); |
| parent.check(); |
| } |
| this.sortGroups = {}; |
| for (const keep of [64, 256]) { |
| const pg = parent.sortGroups[keep]; |
| this.sortGroups[keep] = { |
| ...pg, |
| first: parent.group("sortFirst" + keep, [ |
| w.logits, |
| parent.work.sortA, |
| w.tokens, |
| parent.work.penalty, |
| ]), |
| sample: parent.group("sample", [pg.source, w.tokens, w.sampler]), |
| }; |
| } |
| this.ready = true; |
| } |
| capture(pass, layer, count) { |
| const index = DSPARK_LAYERS.indexOf(layer); |
| if (index >= 0) |
| this.dispatch(pass, "capture", this.captureGroups[index], H / 256, count); |
| } |
| updateContext(encoder, position, count) { |
| if (!count) return; |
| this.contextRun = position + count; |
| this.prepareStep(position, count); |
| const pass = encoder.beginComputePass(); |
| pass.setBindGroup(0, this.stepGroup, [position * 256]); |
| const run = (name, g, x = 1, y = 1, z = 1) => |
| this.dispatch(pass, name, g, x, y, z), |
| block = count <= 8; |
| run( |
| block ? "projectBlock" : "projectBatch", |
| block ? this.groups.projectBlock : this.groups.projectBatch, |
| H / (block ? 8 : 32), |
| ); |
| run("norm", this.groups.hiddenNorm, count); |
| for (const { groups: g } of this.layers) { |
| run( |
| block ? "contextBlock" : "contextBatch", |
| block ? g.contextBlock : g.contextBatch, |
| (KV * 2) / (block ? 8 : 32), |
| ); |
| run("contextNormRope", g.contextRope, 2, count); |
| run("store", g.store, KV / 32, count, 2); |
| } |
| pass.end(); |
| } |
| propose(encoder, position) { |
| const from = Math.max(0, position - 64); |
| encoder.copyBufferToBuffer( |
| this.parent.work.tokens, |
| from * 4, |
| this.work.tokens, |
| from * 4, |
| (position - from + 1) * 4, |
| ); |
| const count = 7, |
| plan = this.prepareStep(position, count), |
| pass = encoder.beginComputePass(); |
| pass.setBindGroup(0, this.stepGroup, [position * 256]); |
| const run = (name, g, x = 1, y = 1, z = 1) => |
| this.dispatch(pass, name, g, x, y, z); |
| run("noise", this.groups.noise); |
| this.parent.dispatch(pass, "embed", this.groups.embed, H / 256, count); |
| for (const { groups: g } of this.layers) { |
| run("norm", g.input, count); |
| run("qkvBlock", g.qkv, (H + KV * 2) / 8); |
| run("noiseNormRope", g.noiseRope, 18, count); |
| run("store", g.store, KV / 32, count, 2); |
| run("attention", g.attention, plan.chunks, 16, count); |
| run("combine", g.combine, 16, count); |
| run("outputBlock", g.output, H / 8); |
| run("norm", g.post, count); |
| run("upGateBlock", g.upGate, (F * 2) / 8); |
| run("silu", g.silu, F / 256, count); |
| run("downBlock", g.down, H / 8); |
| } |
| run("norm", this.groups.finalNorm, count); |
| this.parent.dispatch(pass, "headBlock", this.groups.head, VOCAB / 8); |
| for (let i = 0; i < 7; i++) { |
| this.device.queue.writeBuffer( |
| this.work.sampleSteps, |
| (position + i) * 256, |
| new Uint32Array([position + i, 1, 0, 1]), |
| ); |
| pass.setBindGroup(0, this.sampleStepGroup, [(position + i) * 256]); |
| run("markovEmbed", this.groups.markovEmbed); |
| run("markov", this.groups.markov, VOCAB / 8); |
| run("addBias", this.biasGroups[i], VOCAB / 256); |
| if (this.parent.samplingSettings.temperature > 0) { |
| const keep = this.parent.sampleK <= 64 ? 64 : 256, |
| g = this.sortGroups[keep]; |
| this.parent.encodeSample(pass, g.first, g); |
| this.parent.dispatch(pass, "saveCandidates", g.save[8 + i], keep / 64); |
| } else { |
| run("argmax", this.groups.argmax, 128); |
| run("argmaxFinal", this.groups.argmaxFinal); |
| } |
| } |
| pass.end(); |
| encoder.copyBufferToBuffer( |
| this.work.tokens, |
| (position + 1) * 4, |
| this.parent.work.tokens, |
| (position + 1) * 4, |
| 7 * 4, |
| ); |
| } |
| } |
|
|
