src/bench/state-drift-reference.test.ts

import { describe, expect, it } from 'vitest';
import fixtureJson from '../../public/fixtures/state-drift-27b-cpu-reference.json';
import type {
  EngineCapabilities,
  LoadModelResult,
  ManifestModelV2,
} from '../engine';
import { validateStateDriftReferenceFixture } from './state-drift-reference';

const fixture: unknown = fixtureJson;

const model = {
  id: '27b',
  source: {
    repo: 'WaveCut/Bonsai-web-GGUF',
    revision: 'd85382fa09fe868c0242d81488dfc2edd8d3729b',
    file: 'Bonsai-27B-Q1_0.gguf',
    bytes: 3_803_452_480,
    sha256: '17ef842e47450caeb8eaa3ebfbbab5d2f2278b62b79be107985fb69a2f819aa0',
  },
} as ManifestModelV2;

const capabilities = {
  runtime: {
    llamaCppRevision: '00fa7cb284cbf133fc426733bd64238a3588a33e',
  },
} as EngineCapabilities;

const loadResult = {
  modelId: '27b',
  backend: 'webgpu',
  gate: { requestedBackend: 'webgpu' },
  context: {
    size: 2_048,
    batchSize: 32,
    microBatchSize: 16,
    vocabularySize: 248_320,
  },
  tuning: {
    scope: 'benchmark',
    requested: { nBatch: 32, nUbatch: 16 },
    observed: { nBatch: 32, nUbatch: 16 },
    applied: true,
  },
} as LoadModelResult;

describe('state-drift CPU reference fixture', () => {
  it('validates the pinned fixture against the loaded model, runtime, context, and batching', () => {
    const result = validateStateDriftReferenceFixture(fixture, {
      model,
      capabilities,
      loadResult,
    });

    expect(result.promptTokenIds).toHaveLength(38);
    expect(result.referenceTokenIds).toHaveLength(1_024);
    expect(result.referenceTokenIdsSha256).toBe(
      'c503b2db0dcf10daecfda4f19a5f466d1699b31688076d6c32f7c29daefcef9b',
    );
    expect(result.checkpointPrefixes).toHaveLength(6);
  });

  it('rejects provenance or token-hash drift before scoring', () => {
    const wrongModel = structuredClone(fixture) as Record<string, unknown>;
    const provenance = wrongModel.provenance as { model: { sha256: string } };
    provenance.model.sha256 = '0'.repeat(64);
    expect(() => validateStateDriftReferenceFixture(wrongModel, {
      model,
      capabilities,
      loadResult,
    })).toThrow('model source sha256');

    const wrongTokens = structuredClone(fixture) as { referenceTokenIds: number[] };
    wrongTokens.referenceTokenIds[29] = wrongTokens.referenceTokenIds[29]! + 1;
    expect(() => validateStateDriftReferenceFixture(wrongTokens, {
      model,
      capabilities,
      loadResult,
    })).toThrow('referenceTokenIdsSha256');
  });
});