engine/shapes.js

import { FFN, DECODE_CAP, D_MODEL, VOCAB, SRC_CAP } from './constants.js';

// WebGPU default maxStorageBufferBindingSize — the conservative fallback when
// a caller-built ctx has no granted-limits info.
const DEFAULT_MAX_BINDING = 134217728;

// SRC_CAP (the active model's max source length) is a multiple of 32, so the
// bucket never exceeds the cap (applyModelConfig invariant).
export function bucketFor(nTokens) {
  if (nTokens > SRC_CAP) {
    throw new Error(`input too long: ${nTokens} tokens (max ${SRC_CAP})`);
  }
  return Math.max(32, alignUp(nTokens, 32));
}

// Per-row decode budget: short rows keep a small ring, long rows get the full
// model cap. The old worker formula clamped at 192 "to stay under the engine
// cap" — but split.js's own estimate (~1.08 tok/source-char) puts a 200-char
// row at ~215 output tokens, so long TYPED rows lost their tails while the
// same text as a file (packed passes maxNewTokens=DECODE_CAP) decoded fully.
// EOS early-stop means rows that finish sooner never pay the larger cap, and
// maxBatchForLimits below already budgets KV at DECODE_CAP per row.
export function maxNewTokensFor(chars) {
  return Math.max(24, Math.min(DECODE_CAP, Math.ceil(chars * 3.2) + 12));
}

// Largest batch B that fits the device's storage-binding limit at bucketed
// source length S, across every per-batch-row allocation the engine makes.
// The per-row ceiling is the max of:
//   - encoder ffnTmp activation  [B·S, FFN=1792]   → S·1792·dtypeBytes
//   - decode KV cache per buffer [B, 224, 448]     → 224·448·dtypeBytes
//     (dominates encoder ffnTmp when S < 56, i.e. the S=32 bucket)
//   - f32 logits                 [B, VOCAB=24000]  → 24000·4 (dtype-independent)
// ctx.limits comes from initDevice(); without it the WebGPU default (128 MiB)
// is assumed.
export function maxBatchForLimits(ctx, S, dtypeBytes) {
  if (!Number.isFinite(S) || S <= 0) throw new Error(`maxBatchForLimits: bad S=${S}`);
  if (!Number.isFinite(dtypeBytes) || dtypeBytes <= 0) {
    throw new Error(`maxBatchForLimits: bad dtypeBytes=${dtypeBytes}`);
  }
  const limit = ctx?.limits?.maxStorageBufferBindingSize ?? DEFAULT_MAX_BINDING;
  const perRow = Math.max(S * FFN * dtypeBytes, DECODE_CAP * D_MODEL * dtypeBytes, VOCAB * 4);
  return Math.max(1, Math.floor(limit / perRow));
}

export function alignUp(n, a) {
  return Math.ceil(n / a) * a;
}