src/bench/state-drift-reference.ts

import { sha256 } from '@noble/hashes/sha2.js';
import { bytesToHex } from '@noble/hashes/utils.js';
import type {
  EngineCapabilities,
  LoadModelResult,
  ManifestModelV2,
} from '../engine';
import { BENCHMARK_WORKLOADS, TOKEN_ID_CHECKPOINTS } from './report';

export const STATE_DRIFT_REFERENCE_PATH = 'fixtures/state-drift-27b-cpu-reference.json';
const STATE_DRIFT_SOURCE_EVIDENCE = 'results/space-model/27b-native-cpu-state-drift-1024-b32-u16.json';
const STATE_DRIFT_RENDERED_PROMPT_SHA256 = '7c46d0de443e6f1235ddaa7c0a55c9da710eb73958d5671247b6f0e3b7d189c7';
const STATE_DRIFT_PROMPT_TOKEN_IDS_SHA256 = 'e7af7d66c96f24cb148575bd45cef079920669eb0d193fdc830ce61dfec85451';
const STATE_DRIFT_REFERENCE_TOKEN_IDS_SHA256 = 'c503b2db0dcf10daecfda4f19a5f466d1699b31688076d6c32f7c29daefcef9b';

export interface StateDriftReferenceFixture {
  schemaVersion: 1;
  kind: 'bonsai-state-drift-reference';
  workload: {
    id: 'state-drift-1k-v1';
    messages: Array<{ role: 'user'; content: string }>;
  };
  provenance: {
    sourceEvidence: string;
    engineRevision: string;
    nativeBinary: { bytes: number; sha256: string };
    model: { file: string; bytes: number; sha256: string };
    renderedPromptSha256: string;
    execution: {
      backend: 'cpu';
      contextSize: number;
      batchSize: number;
      microBatchSize: number;
    };
  };
  tokenEncoding: 'uint32-le';
  promptTokenIds: number[];
  promptTokenIdsSha256: string;
  referenceTokenIds: number[];
  referenceTokenIdsSha256: string;
  checkpointPrefixes: Array<{ tokens: number; sha256: string }>;
}

export interface StateDriftReferenceContext {
  model: ManifestModelV2;
  capabilities: EngineCapabilities;
  loadResult: LoadModelResult;
}

function fail(message: string): never {
  throw new Error(`Invalid state-drift CPU reference: ${message}`);
}

function record(value: unknown, label: string): Record<string, unknown> {
  if (typeof value !== 'object' || value === null || Array.isArray(value)) {
    fail(`${label} must be an object.`);
  }
  return value as Record<string, unknown>;
}

function exact<T>(value: unknown, expected: T, label: string): T {
  if (!Object.is(value, expected)) {
    fail(`${label} must be ${JSON.stringify(expected)}; received ${JSON.stringify(value)}.`);
  }
  return expected;
}

function nonEmptyString(value: unknown, label: string): string {
  if (typeof value !== 'string' || value.length === 0) fail(`${label} must be a non-empty string.`);
  return value;
}

function positiveInteger(value: unknown, label: string): number {
  if (!Number.isSafeInteger(value) || (value as number) <= 0) {
    fail(`${label} must be a positive safe integer.`);
  }
  return value as number;
}

function sha256Digest(value: unknown, label: string): string {
  const digest = nonEmptyString(value, label).toLowerCase();
  if (!/^[0-9a-f]{64}$/.test(digest)) fail(`${label} must be a 64-character SHA256 digest.`);
  return digest;
}

function gitRevision(value: unknown, label: string): string {
  const revision = nonEmptyString(value, label).toLowerCase();
  if (!/^[0-9a-f]{40}$/.test(revision)) fail(`${label} must be a pinned 40-character revision.`);
  return revision;
}

function tokenIds(
  value: unknown,
  expectedLength: number,
  vocabularySize: number,
  label: string,
): number[] {
  if (!Array.isArray(value) || value.length !== expectedLength) {
    fail(`${label} must contain exactly ${expectedLength} token ids.`);
  }
  return value.map((tokenId, index) => {
    if (!Number.isSafeInteger(tokenId) || tokenId < 0 || tokenId >= vocabularySize) {
      fail(`${label}[${index}] is outside the loaded vocabulary.`);
    }
    return tokenId;
  });
}

function tokenIdsSha256(values: readonly number[], count = values.length): string {
  const bytes = new Uint8Array(count * Uint32Array.BYTES_PER_ELEMENT);
  const view = new DataView(bytes.buffer);
  for (let index = 0; index < count; index += 1) {
    view.setUint32(index * Uint32Array.BYTES_PER_ELEMENT, values[index]!, true);
  }
  return bytesToHex(sha256(bytes));
}

export function validateStateDriftReferenceFixture(
  value: unknown,
  context: StateDriftReferenceContext,
): StateDriftReferenceFixture {
  const root = record(value, 'root');
  exact(root.schemaVersion, 1, 'schemaVersion');
  exact(root.kind, 'bonsai-state-drift-reference', 'kind');

  const workload = record(root.workload, 'workload');
  exact(workload.id, 'state-drift-1k-v1', 'workload.id');
  const expectedMessages = BENCHMARK_WORKLOADS['state-drift-1k-v1'].messages;
  if (JSON.stringify(workload.messages) !== JSON.stringify(expectedMessages)) {
    fail('workload.messages do not match the locked state-drift prompt.');
  }

  const { model, capabilities, loadResult } = context;
  if (
    model.id !== '27b'
    || loadResult.modelId !== '27b'
    || loadResult.backend !== 'webgpu'
    || loadResult.gate.requestedBackend !== 'webgpu'
    || loadResult.tuning.scope !== 'benchmark'
    || !loadResult.tuning.applied
  ) {
    fail('the loaded runtime is not the explicit 27B WebGPU benchmark path.');
  }

  const provenance = record(root.provenance, 'provenance');
  const sourceEvidence = nonEmptyString(provenance.sourceEvidence, 'provenance.sourceEvidence');
  exact(sourceEvidence, STATE_DRIFT_SOURCE_EVIDENCE, 'provenance.sourceEvidence');
  const engineRevision = gitRevision(provenance.engineRevision, 'provenance.engineRevision');
  exact(engineRevision, capabilities.runtime.llamaCppRevision, 'runtime llama.cpp revision');

  const nativeBinary = record(provenance.nativeBinary, 'provenance.nativeBinary');
  const nativeBinaryBytes = positiveInteger(nativeBinary.bytes, 'provenance.nativeBinary.bytes');
  const nativeBinarySha256 = sha256Digest(
    nativeBinary.sha256,
    'provenance.nativeBinary.sha256',
  );

  const fixtureModel = record(provenance.model, 'provenance.model');
  const modelFile = nonEmptyString(fixtureModel.file, 'provenance.model.file');
  const modelBytes = positiveInteger(fixtureModel.bytes, 'provenance.model.bytes');
  const modelSha256 = sha256Digest(fixtureModel.sha256, 'provenance.model.sha256');
  exact(modelFile, model.source.file, 'model source file');
  exact(modelBytes, model.source.bytes, 'model source bytes');
  exact(modelSha256, model.source.sha256, 'model source sha256');

  const renderedPromptSha256 = sha256Digest(
    provenance.renderedPromptSha256,
    'provenance.renderedPromptSha256',
  );
  exact(
    renderedPromptSha256,
    STATE_DRIFT_RENDERED_PROMPT_SHA256,
    'provenance.renderedPromptSha256',
  );
  const execution = record(provenance.execution, 'provenance.execution');
  exact(execution.backend, 'cpu', 'provenance.execution.backend');
  const contextSize = positiveInteger(execution.contextSize, 'provenance.execution.contextSize');
  const batchSize = positiveInteger(execution.batchSize, 'provenance.execution.batchSize');
  const microBatchSize = positiveInteger(
    execution.microBatchSize,
    'provenance.execution.microBatchSize',
  );
  exact(contextSize, loadResult.context.size, 'reference/loaded context size');
  exact(batchSize, loadResult.context.batchSize, 'reference/loaded batch size');
  exact(microBatchSize, loadResult.context.microBatchSize, 'reference/loaded micro-batch size');
  exact(loadResult.tuning.requested.nBatch, batchSize, 'requested reference batch size');
  exact(loadResult.tuning.requested.nUbatch, microBatchSize, 'requested reference micro-batch size');
  exact(loadResult.tuning.observed.nBatch, batchSize, 'observed reference batch size');
  exact(loadResult.tuning.observed.nUbatch, microBatchSize, 'observed reference micro-batch size');

  exact(root.tokenEncoding, 'uint32-le', 'tokenEncoding');
  const promptTokenIds = tokenIds(
    root.promptTokenIds,
    38,
    loadResult.context.vocabularySize,
    'promptTokenIds',
  );
  const promptTokenIdsSha256 = sha256Digest(root.promptTokenIdsSha256, 'promptTokenIdsSha256');
  exact(promptTokenIdsSha256, tokenIdsSha256(promptTokenIds), 'promptTokenIdsSha256');
  exact(
    promptTokenIdsSha256,
    STATE_DRIFT_PROMPT_TOKEN_IDS_SHA256,
    'pinned promptTokenIdsSha256',
  );
  const referenceTokenIds = tokenIds(
    root.referenceTokenIds,
    1_024,
    loadResult.context.vocabularySize,
    'referenceTokenIds',
  );
  const referenceTokenIdsSha256 = sha256Digest(
    root.referenceTokenIdsSha256,
    'referenceTokenIdsSha256',
  );
  exact(referenceTokenIdsSha256, tokenIdsSha256(referenceTokenIds), 'referenceTokenIdsSha256');
  exact(
    referenceTokenIdsSha256,
    STATE_DRIFT_REFERENCE_TOKEN_IDS_SHA256,
    'pinned referenceTokenIdsSha256',
  );

  if (!Array.isArray(root.checkpointPrefixes)
    || root.checkpointPrefixes.length !== TOKEN_ID_CHECKPOINTS.length) {
    fail(`checkpointPrefixes must contain exactly ${TOKEN_ID_CHECKPOINTS.length} entries.`);
  }
  const checkpointPrefixes = root.checkpointPrefixes.map((rawCheckpoint, index) => {
    const checkpoint = record(rawCheckpoint, `checkpointPrefixes[${index}]`);
    const tokens = TOKEN_ID_CHECKPOINTS[index]!;
    exact(checkpoint.tokens, tokens, `checkpointPrefixes[${index}].tokens`);
    const embeddedSha256 = sha256Digest(
      checkpoint.sha256,
      `checkpointPrefixes[${index}].sha256`,
    );
    exact(
      embeddedSha256,
      tokenIdsSha256(referenceTokenIds, tokens),
      `checkpointPrefixes[${index}].sha256`,
    );
    return { tokens, sha256: embeddedSha256 };
  });

  return {
    schemaVersion: 1,
    kind: 'bonsai-state-drift-reference',
    workload: {
      id: 'state-drift-1k-v1',
      messages: expectedMessages.map((message) => ({ ...message })) as Array<{
        role: 'user';
        content: string;
      }>,
    },
    provenance: {
      sourceEvidence,
      engineRevision,
      nativeBinary: { bytes: nativeBinaryBytes, sha256: nativeBinarySha256 },
      model: { file: modelFile, bytes: modelBytes, sha256: modelSha256 },
      renderedPromptSha256,
      execution: {
        backend: 'cpu',
        contextSize,
        batchSize,
        microBatchSize,
      },
    },
    tokenEncoding: 'uint32-le',
    promptTokenIds,
    promptTokenIdsSha256,
    referenceTokenIds,
    referenceTokenIdsSha256,
    checkpointPrefixes,
  };
}