webbrain-one's picture
FP16-weight package of d1-browser-decision-fp32 (FP32 arithmetic)
d83c32b verified
Raw History Blame Contribute Delete
13.8 kB
import { encode, temperature, imageInputs, positionMatrix, pyJSON } from './d1-preprocess.js';
export const D1_PROVIDER = 'webgpu_d1';
const record = value => value && typeof value === 'object' && !Array.isArray(value);
const checkAbort = signal => { if (signal?.aborted) throw signal.reason || new Error('D1 request cancelled.'); };
const dispose = values => { for (const value of Object.values(values || {})) value.dispose?.(); };
// Media blocks are data, never fetched from a page-chosen network URL.
export function splitD1State(state, images = []) {
const media = [...images];
let text = state == null ? '' : state;
if (Array.isArray(state)) {
text = state.filter(part => {
if (part?.type !== 'image_url') return true;
media.push(part.image_url?.url); return false;
});
if (!text.length) text = '';
}
if (media.length > 4 || media.some(url => typeof url !== 'string' || !/^data:image\/(png|jpeg|webp);base64,[A-Za-z0-9+/=]+$/.test(url))) {
throw new Error('D1 accepts at most four inline PNG/JPEG/WebP screenshots; remote image URLs are not allowed.');
}
if (new TextEncoder().encode(JSON.stringify([text, media])).length > 8 * 1024 * 1024) throw new Error('D1 input exceeds 8 MiB.');
return { state: text, images: media };
}
export function d1Answers(questions, probabilities, width) {
const answers = Object.create(null);
const entries = Array.isArray(questions) ? questions : Object.entries(questions);
entries.forEach(([id, question], row) => {
const names = question.type === 'noul' ? ['no', 'yes'] : question.type === 'score' ? question.criteria.map((_, i) => String(i)) : Object.keys(question.criteria);
const values = Array.from(probabilities.slice(row * width, row * width + names.length), Number);
if (values.length !== names.length || values.some(value => !Number.isFinite(value) || value < 0 || value > 1) || Math.abs(values.reduce((a, b) => a + b, 0) - 1) > .001) throw new Error('Invalid D1 probability output.');
const selected = values.indexOf(Math.max(...values));
if (question.type === 'noul') answers[id] = { type: 'noul', noul: values[1] };
else {
const common = { type: question.type, confidence: values[selected], probabilities: Object.fromEntries(names.map((name, i) => [name, values[i]])) };
answers[id] = question.type === 'choice' ? { ...common, choice: names[selected] } : { ...common, score: values.reduce((sum, p, i) => sum + p * i, 0), legend: Object.fromEntries(names.map((name, i) => [name, typeof question.criteria[i] === 'string' ? question.criteria[i] : pyJSON(question.criteria[i])])) };
}
});
return answers;
}
// The encoder's six full-attention layers are dense: each one materializes a
// heads × length × length FP32 score tensor per row (plus elementwise copies
// of it), so memory grows with the square of the row length, not with the
// token count. Batches are bounded by that tensor, and text-only rows are
// capped on WebGPU; the source truncates the state, never the question, when
// a row is over the cap.
export const D1_ATTENTION_HEADS = 16;
export const D1_ATTENTION_BUDGET_BYTES = 256 * 1024 * 1024;
export const D1_WEBGPU_MAX_TEXT_TOKENS = 2048;
export const d1AttentionBytes = (rows, length) => rows * D1_ATTENTION_HEADS * length * length * 4;
export function d1RowBatches(rows, mediaLength, { maxTokens = 65536, budgetBytes = D1_ATTENTION_BUDGET_BYTES } = {}) {
const order = rows.map((_, i) => i).sort((a, b) => rows[b].input_ids.length - rows[a].input_ids.length), batches = [];
while (order.length) {
const longest = rows[order[0]].input_ids.length + mediaLength;
let size = 1;
while (size < order.length && (size + 1) * longest <= maxTokens && d1AttentionBytes(size + 1, longest) <= budgetBytes) size++;
batches.push(order.splice(0, size));
}
return batches;
}
const HIDDEN = 768, PATCHES = 1024, WIDTH = 1024, FLOAT_BYTES = 4, VISION_HEADS = 12;
export function createD1Runtime({ ort, tokenizer, config, ratios, sessions, device, model, dtype = 'float32' }) {
if (!sessions?.decision || !sessions?.vision || !sessions?.projector) throw new Error('D1 requires all three sessions.');
const tensor = (type, data, shape) => new ort.Tensor(type, data, shape);
// One score tensor must fit a single storage binding on this adapter.
const rowLimit = Number(device?.limits?.maxStorageBufferBindingSize) || Infinity;
const budgetBytes = Math.min(D1_ATTENTION_BUDGET_BYTES, rowLimit);
const usage = globalThis.GPUBufferUsage ? GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST : 0;
const gpuTensor = (bytes, dims) => {
const buffer = device.createBuffer({ size: Math.max(16, Math.ceil(bytes / 16) * 16), usage });
return { buffer, tensor: ort.Tensor.fromGpuBuffer(buffer, { dataType: 'float32', dims, dispose: () => buffer.destroy() }) };
};
const copy = (pairs) => {
const encoder = device.createCommandEncoder();
for (const [source, sourceOffset, target, targetOffset, bytes] of pairs) encoder.copyBufferToBuffer(source, sourceOffset, target, targetOffset, bytes);
device.queue.submit([encoder.finish()]);
};
// Vision and projector outputs stay on the GPU when the sessions return
// gpu-buffer tensors: each crop is copied into the projector input and each
// projected chunk into every row of the decision input, without a readback.
async function mediaPrefix(images, signal) {
const chunks = [], layouts = [];
try {
for (const url of images) {
checkAbort(signal);
const image = await imageInputs(url, ratios), crops = image.shapes.length;
// The vision tower's attention is dense too (12 heads × 1024² per crop),
// so a tiled screenshot runs in crop groups that fit the same budget.
const group = Math.max(1, Math.floor(budgetBytes / (VISION_HEADS * PATCHES * PATCHES * FLOAT_BYTES)));
for (let first = 0; first < crops; first += group) {
const count = Math.min(group, crops - first);
const positions = new Float32Array(count * PATCHES * 256);
for (let i = 0; i < count; i++) positions.set(positionMatrix(...image.shapes[first + i]), i * PATCHES * 256);
const feeds = {
pixel_values: tensor('float32', image.pixels.slice(first * PATCHES * HIDDEN, (first + count) * PATCHES * HIDDEN), [count, PATCHES, HIDDEN]),
pixel_mask: tensor('bool', Uint8Array.from(image.mask.slice(first * PATCHES, (first + count) * PATCHES)), [count, PATCHES]),
position_matrix: tensor('float32', positions, [count, PATCHES, 256]),
};
let output;
try {
output = await sessions.vision.run(feeds); checkAbort(signal);
const gpu = output.hidden.location === 'gpu-buffer' && !!device;
const hidden = gpu ? null : await output.hidden.getData();
for (let crop = 0; crop < count; crop++) {
const [h, w] = image.shapes[first + crop], bytes = h * w * HIDDEN * FLOAT_BYTES;
let input, buffer, projected;
if (gpu) ({ tensor: input, buffer } = gpuTensor(bytes, [1, h, w, HIDDEN]));
else input = tensor('float32', hidden.slice(crop * PATCHES * HIDDEN, crop * PATCHES * HIDDEN + h * w * HIDDEN), [1, h, w, HIDDEN]);
try {
if (gpu) copy([[output.hidden.gpuBuffer, crop * PATCHES * HIDDEN * FLOAT_BYTES, buffer, 0, bytes]]);
projected = await sessions.projector.run({ hidden: input }); checkAbort(signal);
if (gpu && projected.prefix.location === 'gpu-buffer') { chunks.push({ tensor: projected.prefix, rows: projected.prefix.dims[1] }); projected = null; }
else { const data = new Float32Array(await projected.prefix.getData()); chunks.push({ data, rows: data.length / WIDTH }); }
} finally { input.dispose?.(); dispose(projected); }
}
} finally { dispose(feeds); dispose(output); }
}
layouts.push({ ...image.plan, shapes: image.shapes });
}
} catch (error) { for (const chunk of chunks) chunk.tensor?.dispose?.(); throw error; }
return { chunks, layouts, mediaLength: chunks.reduce((n, chunk) => n + chunk.rows, 0) };
}
function mediaFeed(chunks, mediaLength, batch) {
if (!mediaLength || !chunks.some(chunk => chunk.tensor)) {
const prefix = new Float32Array(mediaLength * WIDTH);
let offset = 0; for (const chunk of chunks) { prefix.set(chunk.data, offset); offset += chunk.data.length; }
const media = new Float32Array(batch * prefix.length);
for (let row = 0; row < batch; row++) media.set(prefix, row * prefix.length);
return tensor('float32', media, [batch, mediaLength, WIDTH]);
}
const rowBytes = mediaLength * WIDTH * FLOAT_BYTES, { tensor: media, buffer } = gpuTensor(batch * rowBytes, [batch, mediaLength, WIDTH]), pairs = [];
for (let row = 0; row < batch; row++) {
let offset = row * rowBytes;
for (const chunk of chunks) {
const bytes = chunk.rows * WIDTH * FLOAT_BYTES;
if (chunk.tensor) pairs.push([chunk.tensor.gpuBuffer, 0, buffer, offset, bytes]);
else device.queue.writeBuffer(buffer, offset, chunk.data);
offset += bytes;
}
}
copy(pairs);
return media;
}
return {
async evaluate({ state = '', images = [], questions, signal } = {}) {
checkAbort(signal);
if (!record(questions) || !Object.keys(questions).length || Object.keys(questions).length > 32) throw new Error('D1 requires 1 to 32 named questions.');
const input = splitD1State(state, images);
const { chunks, layouts, mediaLength } = await mediaPrefix(input.images, signal);
try {
// The longest row whose score tensor still fits one storage binding.
const maxRow = Math.floor(Math.sqrt(rowLimit / (D1_ATTENTION_HEADS * FLOAT_BYTES)));
const textCap = input.images.length ? config.image_text_length : Math.min(config.max_length, D1_WEBGPU_MAX_TEXT_TOKENS);
const limit = Math.min(textCap, config.max_length - mediaLength, maxRow - mediaLength);
if (!Number.isInteger(mediaLength) || config.max_length - mediaLength < 64) throw new Error('D1 media exceed the source context budget.');
if (limit < 64) throw new Error(`D1 input is too large for this GPU: ${mediaLength} screenshot positions. Send fewer or smaller screenshots.`);
const entries = Object.entries(questions), rows = entries.map(([, q]) => encode(tokenizer, input.state, q, limit, input.images.length > 0));
const longest = Math.max(...rows.map(row => row.input_ids.length)) + mediaLength;
if (d1AttentionBytes(1, longest) > rowLimit) throw new Error(`D1 input is too large for this GPU: ${longest} tokens including screenshots. Send fewer or smaller screenshots.`);
const answers = Object.create(null), batches = [];
for (const indices of d1RowBatches(rows, mediaLength, { budgetBytes })) {
const selectedRows = indices.map(i => rows[i]), selectedEntries = indices.map(i => entries[i]);
const batch = selectedRows.length, length = Math.max(...selectedRows.map(row => row.input_ids.length)), width = Math.max(...selectedRows.map(row => row.markers.length));
const ids = new BigInt64Array(batch * length), mask = new Uint8Array(batch * length), markers = new BigInt64Array(batch * width), markerMask = new Uint8Array(batch * width);
ids.fill(BigInt(tokenizer.pad_token_id ?? 0));
selectedRows.forEach((row, i) => {
ids.set(row.input_ids.map(BigInt), i * length); mask.fill(1, i * length, i * length + row.input_ids.length);
markers.set(row.markers.map(BigInt), i * width); markerMask.fill(1, i * width, i * width + row.markers.length);
});
const feeds = { input_ids: tensor('int64', ids, [batch, length]), text_mask: tensor('bool', mask, [batch, length]), marker_pos: tensor('int64', markers, [batch, width]), marker_mask: tensor('bool', markerMask, [batch, width]), qtype: tensor('int64', BigInt64Array.from(selectedEntries, ([, q]) => BigInt({ choice: 0, score: 1, noul: 2 }[q.type])), [batch]), temperature: tensor('float32', Float32Array.from(selectedEntries, ([, q]) => temperature(q, config, input.images.length > 0)), [batch]), media_prefix: mediaFeed(chunks, mediaLength, batch) };
let output;
try {
checkAbort(signal); output = await sessions.decision.run(feeds);
const probabilities = await output.probabilities.getData(); await device?.queue.onSubmittedWorkDone(); checkAbort(signal);
Object.assign(answers, d1Answers(selectedEntries, probabilities, width));
batches.push({ question_ids: selectedEntries.map(([id]) => id), batch, text_length: length, options: width, attention_bytes: d1AttentionBytes(batch, length + mediaLength) });
} finally { dispose(feeds); dispose(output); }
}
return { model, provider: D1_PROVIDER, answers: Object.fromEntries(entries.map(([id]) => [id, answers[id]])), usage: { input_tokens: rows.reduce((sum, row) => sum + row.input_ids.length + mediaLength, 0), output_tokens: 0 }, diagnostics: { dtype, generated_tokens: 0, batch: rows.length, text_length: Math.max(...rows.map(row => row.input_ids.length)), media_length: mediaLength, media_on_gpu: chunks.some(chunk => chunk.tensor), layouts, batches, attention_budget_bytes: budgetBytes, cpu_fallback_policy: 'CPU/WASM control/shape operators allowed; actual placement requires a separate profile.' } };
} finally { for (const chunk of chunks) chunk.tensor?.dispose?.(); }
},
async dispose() { for (const session of Object.values(sessions)) await session.release(); device?.destroy(); },
};
}