@simulatte/doppler 0.1.5 → 0.1.7

package/src/training/distillation/student-fixture.js

@@ -0,0 +1,846 @@
+import { getDevice } from '../../gpu/device.js';
+import { createTrainingConfig } from '../../config/training-defaults.js';
+import {
+  runAttention,
+  castF16ToF32,
+  runGather,
+  runMatmul,
+  runResidualAdd,
+  runRMSNorm,
+  runRoPE,
+  runSiLURowSplit,
+} from '../../gpu/kernels/index.js';
+import { createTensor } from '../../gpu/tensor.js';
+import { acquireBuffer, uploadData, releaseBuffer } from '../../memory/buffer-pool.js';
+import { getBufferDtype, getWeightDtype, isCpuWeightBuffer, isWeightBuffer } from '../../gpu/weight-buffer.js';
+import { OpType } from '../autograd.js';
+import { normalizeOptionalString } from './suite-data.js';
+
+const DISTILL_ADAPTER_TOP_K = 64;
+const DISTILL_STUDENT_GRAPH_PROJECTION = 'projection_head';
+const DISTILL_STUDENT_GRAPH_FULL = 'transformer_full';
+
+function makeTensorFromFloat32(values, shape, label) {
+  const data = values instanceof Float32Array ? values : new Float32Array(values);
+  const buffer = acquireBuffer(data.byteLength, undefined, label || 'train_tensor');
+  uploadData(buffer, data);
+  return createTensor(buffer, 'f32', shape, label || 'train_tensor');
+}
+
+function makeTensorFromF16Bits(values, shape, label) {
+  const data = values instanceof Uint16Array ? values : new Uint16Array(values);
+  const buffer = acquireBuffer(data.byteLength, undefined, label || 'train_tensor_f16');
+  uploadData(buffer, data);
+  return createTensor(buffer, 'f16', shape, label || 'train_tensor_f16');
+}
+
+function makeTensorFromUint32(values, shape, label) {
+  const data = values instanceof Uint32Array ? values : new Uint32Array(values);
+  const buffer = acquireBuffer(data.byteLength, undefined, label || 'train_tokens');
+  uploadData(buffer, data);
+  return createTensor(buffer, 'f32', shape, label || 'train_tokens');
+}
+
+function releaseTensor(tensor) {
+  if (!tensor?.buffer) return;
+  releaseBuffer(tensor.buffer);
+}
+
+function toFloat32Array(values, label = 'values') {
+  if (values instanceof Float32Array) return values;
+  if (ArrayBuffer.isView(values)) {
+    return new Float32Array(values.buffer.slice(values.byteOffset, values.byteOffset + values.byteLength));
+  }
+  if (values instanceof ArrayBuffer) {
+    return new Float32Array(values.slice(0));
+  }
+  if (Array.isArray(values)) {
+    return new Float32Array(values);
+  }
+  throw new Error(`Expected ${label} to be a Float32Array-compatible value.`);
+}
+
+function disposePrefillSnapshot(result) {
+  const cache = result?.cache;
+  if (cache && typeof cache.clear === 'function') {
+    cache.clear();
+  }
+}
+
+function toFiniteNumber(value, fallback) {
+  const parsed = Number(value);
+  return Number.isFinite(parsed) ? parsed : fallback;
+}
+
+function clampDistillTopK(value) {
+  const parsed = Math.floor(toFiniteNumber(value, DISTILL_ADAPTER_TOP_K));
+  return Math.max(2, Math.min(256, parsed));
+}
+
+function normalizeDistillStudentGraphMode(value) {
+  const normalized = normalizeOptionalString(value);
+  if (!normalized) return DISTILL_STUDENT_GRAPH_FULL;
+  const compact = normalized.toLowerCase().replace(/[-\s]/g, '_');
+  if (compact === 'projection_head' || compact === 'projection') {
+    return DISTILL_STUDENT_GRAPH_PROJECTION;
+  }
+  return DISTILL_STUDENT_GRAPH_FULL;
+}
+
+function resolveTensorDtype(value, fallback = 'f32') {
+  const dtype = isWeightBuffer(value)
+    ? value.dtype
+    : (value?.dtype || getWeightDtype(value) || null);
+  const normalized = String(dtype || '').toLowerCase();
+  return normalized === 'f16' ? 'f16' : (normalized === 'f32' ? 'f32' : fallback);
+}
+
+async function ensureTrainableTensor(value, shape, label, ownedTrainables = null) {
+  if (!value) {
+    throw new Error(`Distill full-graph student missing required weight "${label}".`);
+  }
+  const registerOwned = (tensor) => {
+    if (ownedTrainables instanceof Set && tensor?.buffer instanceof GPUBuffer) {
+      ownedTrainables.add(tensor);
+    }
+    return tensor;
+  };
+  if (isWeightBuffer(value)) {
+    if (value.dtype === 'f32') {
+      return value;
+    }
+    if (value.dtype === 'f16') {
+      const sourceShape = Array.isArray(value.shape) && value.shape.length > 0 ? value.shape : [...shape];
+      const source = createTensor(value.buffer, 'f16', sourceShape, `${label}_source_f16`);
+      const promoted = await castF16ToF32(source);
+      return registerOwned(createTensor(promoted.buffer, 'f32', sourceShape, `${label}_trainable_f32`));
+    }
+    throw new Error(`Distill full-graph student weight "${label}" uses unsupported dtype "${value.dtype}".`);
+  }
+  if (value instanceof GPUBuffer) {
+    const sourceShape = [...shape];
+    const rawDtype = String(getBufferDtype(value) || 'f32').toLowerCase();
+    const dtype = rawDtype === 'f16' ? 'f16' : 'f32';
+    const tensor = createTensor(value, dtype, sourceShape, label);
+    if (dtype === 'f16') {
+      const promoted = await castF16ToF32(tensor);
+      return registerOwned(createTensor(promoted.buffer, 'f32', sourceShape, `${label}_trainable_f32`));
+    }
+    return tensor;
+  }
+  if (isCpuWeightBuffer(value)) {
+    const sourceShape = Array.isArray(value.shape) && value.shape.length > 0 ? value.shape : [...shape];
+    const dtype = resolveTensorDtype(value, 'f32');
+    if (dtype === 'f32') {
+      const tensor = makeTensorFromFloat32(value.data, sourceShape, `${label}_cpu_f32`);
+      return registerOwned(tensor);
+    }
+    if (dtype === 'f16') {
+      let raw = null;
+      if (value.data instanceof Uint16Array) {
+        raw = value.data;
+      } else if (ArrayBuffer.isView(value.data)) {
+        raw = new Uint16Array(
+          value.data.buffer,
+          value.data.byteOffset,
+          Math.floor(value.data.byteLength / 2)
+        );
+      } else if (value.data instanceof ArrayBuffer) {
+        raw = new Uint16Array(value.data);
+      }
+      if (!raw) {
+        throw new Error(`Distill full-graph student weight "${label}" has non-typed f16 CPU data.`);
+      }
+      const source = makeTensorFromF16Bits(raw, sourceShape, `${label}_cpu_f16`);
+      const promoted = await castF16ToF32(source);
+      releaseTensor(source);
+      return registerOwned(createTensor(promoted.buffer, 'f32', sourceShape, `${label}_trainable_f32`));
+    }
+    throw new Error(`Distill full-graph student weight "${label}" has unsupported CPU dtype "${dtype}".`);
+  }
+  if (value.buffer instanceof GPUBuffer) {
+    const resolvedShape = Array.isArray(value.shape) && value.shape.length > 0 ? value.shape : [...shape];
+    const tensor = createTensor(
+      value.buffer,
+      resolveTensorDtype(value, 'f32'),
+      resolvedShape,
+      label
+    );
+    if (tensor.dtype === 'f16') {
+      const promoted = await castF16ToF32(tensor);
+      return registerOwned(createTensor(promoted.buffer, 'f32', resolvedShape, `${label}_trainable_f32`));
+    }
+    return tensor;
+  }
+  throw new Error(`Distill full-graph student weight "${label}" is not GPU-resident.`);
+}
+
+async function ensureNormTensor(value, hiddenSize, label, ownedTrainables = null) {
+  return ensureTrainableTensor(value, [hiddenSize], label, ownedTrainables);
+}
+
+function hasTensorPayload(value) {
+  if (!value) return false;
+  if (value instanceof GPUBuffer) return true;
+  if (isWeightBuffer(value) || isCpuWeightBuffer(value)) return true;
+  if (value?.buffer instanceof GPUBuffer) return true;
+  if (ArrayBuffer.isView(value) || Array.isArray(value)) return true;
+  return false;
+}
+
+async function fuseGateUpTensors(gateTensor, upTensor, intermediateSize, hiddenSize, label, ownedTrainables = null) {
+  const device = getDevice();
+  if (!device) {
+    throw new Error('Distill full-graph student requires active GPU device.');
+  }
+  if (gateTensor?.dtype !== 'f32' || upTensor?.dtype !== 'f32') {
+    throw new Error(`Distill fused gate_up expects f32 tensors for "${label}".`);
+  }
+  const expectedRows = intermediateSize;
+  const expectedCols = hiddenSize;
+  const gateRows = Number.isFinite(gateTensor?.shape?.[0]) ? gateTensor.shape[0] : 0;
+  const gateCols = Number.isFinite(gateTensor?.shape?.[1]) ? gateTensor.shape[1] : 0;
+  const upRows = Number.isFinite(upTensor?.shape?.[0]) ? upTensor.shape[0] : 0;
+  const upCols = Number.isFinite(upTensor?.shape?.[1]) ? upTensor.shape[1] : 0;
+  if (gateRows !== expectedRows || gateCols !== expectedCols || upRows !== expectedRows || upCols !== expectedCols) {
+    throw new Error(
+      `Distill gate/up shape mismatch for "${label}": gate=[${gateRows},${gateCols}] up=[${upRows},${upCols}] ` +
+      `expected=[${expectedRows},${expectedCols}]`
+    );
+  }
+  const rowBytes = expectedCols * 4;
+  const blockBytes = expectedRows * rowBytes;
+  const fusedBuffer = acquireBuffer(blockBytes * 2, undefined, `${label}_fused`);
+  const encoder = device.createCommandEncoder();
+  encoder.copyBufferToBuffer(gateTensor.buffer, 0, fusedBuffer, 0, blockBytes);
+  encoder.copyBufferToBuffer(upTensor.buffer, 0, fusedBuffer, blockBytes, blockBytes);
+  device.queue.submit([encoder.finish()]);
+  const fused = createTensor(fusedBuffer, 'f32', [expectedRows * 2, expectedCols], `${label}_fused`);
+  if (ownedTrainables instanceof Set) {
+    ownedTrainables.add(fused);
+  }
+  return fused;
+}
+
+function resolvePhasePrompts(batch, phase) {
+  const distill = batch?.distill || {};
+  const prompts = phase === 'positive'
+    ? distill.tripletPositivePrompts
+    : (phase === 'negative' ? distill.tripletNegativePrompts : distill.prompts);
+  if (!Array.isArray(prompts) || prompts.length === 0) {
+    throw new Error(`Distill student fixture requires distill prompts for phase "${phase}".`);
+  }
+  return prompts;
+}
+
+function createRowSliceTensor(inputTensor, rows, cols, rowIndex, label) {
+  const device = getDevice();
+  if (!device) {
+    throw new Error('Distill full-graph student requires active GPU device.');
+  }
+  const dtype = inputTensor?.dtype === 'f16' ? 'f16' : 'f32';
+  const bytesPerElement = dtype === 'f16' ? 2 : 4;
+  const rowBytes = cols * bytesPerElement;
+  const clampedRow = Math.max(0, Math.min(rows - 1, rowIndex));
+  const outputBuffer = acquireBuffer(rowBytes, undefined, label);
+  const encoder = device.createCommandEncoder();
+  encoder.copyBufferToBuffer(
+    inputTensor.buffer,
+    clampedRow * rowBytes,
+    outputBuffer,
+    0,
+    rowBytes
+  );
+  device.queue.submit([encoder.finish()]);
+  return createTensor(outputBuffer, dtype, [1, cols], label);
+}
+
+function createDistillStudentProjectionModelFixture(overrides = {}, options = {}) {
+  const distillRuntime = options.distillRuntime && typeof options.distillRuntime === 'object'
+    ? options.distillRuntime
+    : null;
+  if (!distillRuntime?.studentPipeline) {
+    throw new Error('Distill student fixture requires distillRuntime.studentPipeline.');
+  }
+  const outputDim = clampDistillTopK(
+    options.outputDim
+    ?? options.inputDim
+    ?? DISTILL_ADAPTER_TOP_K
+  );
+  const inferredEmbeddingDim = Math.floor(
+    Number(distillRuntime.studentPipeline?.modelConfig?.hiddenSize)
+  );
+  const embeddingDim = Number.isInteger(options.embeddingDim) && options.embeddingDim > 0
+    ? options.embeddingDim
+    : (Number.isFinite(inferredEmbeddingDim) && inferredEmbeddingDim > 0
+      ? inferredEmbeddingDim
+      : outputDim);
+  const config = createTrainingConfig({
+    ...overrides,
+    training: {
+      enabled: true,
+      lossScaling: { enabled: false },
+      gradient: { maxNorm: 0 },
+      ...(overrides.training || {}),
+    },
+  });
+
+  const projectionWeights = new Float32Array(embeddingDim * outputDim);
+  const projectionWeight = makeTensorFromFloat32(
+    projectionWeights,
+    [embeddingDim, outputDim],
+    'distill_student_head_weight'
+  );
+  const temporaryInputs = new Set();
+
+  async function projectEmbeddingInput(inputTensor, tape) {
+    const rows = Number.isFinite(inputTensor?.shape?.[0]) ? inputTensor.shape[0] : 1;
+    return tape.record(
+      OpType.MATMUL,
+      (a, b) => runMatmul(a, b, rows, outputDim, embeddingDim, { transposeB: false }),
+      [inputTensor, projectionWeight],
+      { M: rows, N: outputDim, K: embeddingDim, transposeB: false }
+    );
+  }
+
+  async function buildStudentEmbeddingInput(batch, phase = 'anchor') {
+    const prompts = resolvePhasePrompts(batch, phase);
+    const rows = prompts.length;
+    const features = new Float32Array(rows * embeddingDim);
+    for (let row = 0; row < rows; row += 1) {
+      const prompt = String(prompts[row] || '').trim();
+      const studentResult = await distillRuntime.studentPipeline.prefillWithEmbedding(prompt, {
+        useChatTemplate: false,
+        embeddingMode: 'last',
+      });
+      try {
+        const studentEmbedding = toFloat32Array(studentResult?.embedding, 'student embedding');
+        const rowOffset = row * embeddingDim;
+        const copyCount = Math.min(embeddingDim, studentEmbedding.length);
+        features.set(studentEmbedding.subarray(0, copyCount), rowOffset);
+      } finally {
+        disposePrefillSnapshot(studentResult);
+        distillRuntime.studentPipeline.reset();
+      }
+    }
+    const inputTensor = makeTensorFromFloat32(
+      features,
+      [rows, embeddingDim],
+      `distill_student_${phase}_embedding`
+    );
+    temporaryInputs.add(inputTensor);
+    return inputTensor;
+  }
+
+  const model = {
+    async forward(inputTensor, tape) {
+      return projectEmbeddingInput(inputTensor, tape);
+    },
+    async forwardDistill(batch, tape, forwardOptions = {}) {
+      const requestedPhase = String(forwardOptions?.phase || 'anchor').trim();
+      const phase = requestedPhase === 'positive'
+        ? 'positive'
+        : (requestedPhase === 'negative' ? 'negative' : 'anchor');
+      const inputTensor = await buildStudentEmbeddingInput(batch, phase);
+      const logits = await projectEmbeddingInput(inputTensor, tape);
+      return { logits };
+    },
+    cleanupDistillStep() {
+      for (const tensor of temporaryInputs) {
+        releaseTensor(tensor);
+      }
+      temporaryInputs.clear();
+    },
+    loraParams() {
+      return [projectionWeight];
+    },
+    paramGroups() {
+      return {
+        encoder: [],
+        prior: [],
+        decoder: [],
+        base: [projectionWeight],
+        lora: [projectionWeight],
+      };
+    },
+  };
+
+  return {
+    config,
+    model,
+    outputDim,
+    embeddingDim,
+    cleanup() {
+      model.cleanupDistillStep();
+      releaseTensor(projectionWeight);
+    },
+  };
+}
+
+async function createDistillStudentTransformerModelFixture(overrides = {}, options = {}) {
+  const distillRuntime = options.distillRuntime && typeof options.distillRuntime === 'object'
+    ? options.distillRuntime
+    : null;
+  const studentPipeline = distillRuntime?.studentPipeline || null;
+  if (!studentPipeline?.modelConfig || !(studentPipeline.weights instanceof Map)) {
+    throw new Error('Distill full-graph student fixture requires loaded student pipeline weights.');
+  }
+  const modelConfig = studentPipeline.modelConfig;
+  const hiddenSize = Math.max(1, Math.floor(Number(modelConfig.hiddenSize) || 0));
+  const intermediateSize = Math.max(1, Math.floor(Number(modelConfig.intermediateSize) || 0));
+  const numLayers = Math.max(1, Math.floor(Number(modelConfig.numLayers) || 0));
+  const numHeads = Math.max(1, Math.floor(Number(modelConfig.numHeads) || 0));
+  const numKVHeads = Math.max(1, Math.floor(Number(modelConfig.numKVHeads || numHeads) || 0));
+  const headDim = Math.max(1, Math.floor(Number(modelConfig.headDim) || 0));
+  const vocabSize = Math.max(1, Math.floor(Number(modelConfig.vocabSize) || 0));
+  const rmsNormEps = Number.isFinite(modelConfig.rmsNormEps) ? modelConfig.rmsNormEps : 1e-6;
+  const hiddenActivation = String(modelConfig.hiddenActivation || 'silu').toLowerCase();
+  const swigluLimit = Number.isFinite(modelConfig.swigluLimit) ? modelConfig.swigluLimit : 0;
+  const useEmbeddingTranspose = modelConfig.embeddingTranspose === true;
+  const tieWordEmbeddings = modelConfig.useTiedEmbeddings === true;
+
+  const config = createTrainingConfig({
+    ...overrides,
+    training: {
+      enabled: true,
+      lossScaling: { enabled: false },
+      gradient: { maxNorm: 0 },
+      ...(overrides.training || {}),
+    },
+  });
+
+  const ownedTrainables = new Set();
+  const embeddingWeight = await ensureTrainableTensor(
+    studentPipeline.weights.get('embed'),
+    [vocabSize, hiddenSize],
+    'embed',
+    ownedTrainables
+  );
+  const lmHeadWeight = tieWordEmbeddings
+    ? embeddingWeight
+    : await ensureTrainableTensor(
+      studentPipeline.weights.get('lm_head'),
+      [vocabSize, hiddenSize],
+      'lm_head',
+      ownedTrainables
+    );
+  const finalNormWeight = await ensureNormTensor(
+    studentPipeline.weights.get('final_norm'),
+    hiddenSize,
+    'final_norm',
+    ownedTrainables
+  );
+
+  const ropeDim = Math.max(1, Math.floor(headDim / 2));
+  const ropeRows = Math.max(1, Math.floor(Number(modelConfig.maxSeqLen) || 1));
+  const ropeCos = await ensureTrainableTensor(
+    createTensor(studentPipeline.ropeFreqsCos, 'f32', [ropeRows, ropeDim], 'rope_cos'),
+    [ropeRows, ropeDim],
+    'rope_cos',
+    ownedTrainables
+  );
+  const ropeSin = await ensureTrainableTensor(
+    createTensor(studentPipeline.ropeFreqsSin, 'f32', [ropeRows, ropeDim], 'rope_sin'),
+    [ropeRows, ropeDim],
+    'rope_sin',
+    ownedTrainables
+  );
+
+  const layerParams = [];
+  const layers = [];
+  for (let layerIdx = 0; layerIdx < numLayers; layerIdx += 1) {
+    const layerWeights = studentPipeline.weights.get(`layer_${layerIdx}`);
+    if (!layerWeights) {
+      throw new Error(`Distill full-graph student missing layer_${layerIdx} weights.`);
+    }
+    const gateUpWeight = layerWeights.gateUp || layerWeights.ffnGateUp || null;
+    let layerGateUp = null;
+    if (hasTensorPayload(gateUpWeight)) {
+      layerGateUp = await ensureTrainableTensor(
+        gateUpWeight,
+        [intermediateSize * 2, hiddenSize],
+        `layer_${layerIdx}.ffn_gate_up`,
+        ownedTrainables
+      );
+    } else {
+      const gateWeight = layerWeights.gate || layerWeights.ffnGate || null;
+      const upWeight = layerWeights.up || layerWeights.ffnUp || null;
+      if (!hasTensorPayload(gateWeight) || !hasTensorPayload(upWeight)) {
+        throw new Error(
+          `Distill full-graph student missing gate/up projections on layer ${layerIdx}.`
+        );
+      }
+      const gateTensor = await ensureTrainableTensor(
+        gateWeight,
+        [intermediateSize, hiddenSize],
+        `layer_${layerIdx}.ffn_gate`,
+        ownedTrainables
+      );
+      const upTensor = await ensureTrainableTensor(
+        upWeight,
+        [intermediateSize, hiddenSize],
+        `layer_${layerIdx}.ffn_up`,
+        ownedTrainables
+      );
+      layerGateUp = await fuseGateUpTensors(
+        gateTensor,
+        upTensor,
+        intermediateSize,
+        hiddenSize,
+        `layer_${layerIdx}.ffn_gate_up`,
+        ownedTrainables
+      );
+    }
+    const layer = {
+      inputNorm: await ensureNormTensor(
+        layerWeights.inputNorm,
+        hiddenSize,
+        `layer_${layerIdx}.input_norm`,
+        ownedTrainables
+      ),
+      qProj: await ensureTrainableTensor(
+        layerWeights.qProj,
+        [numHeads * headDim, hiddenSize],
+        `layer_${layerIdx}.q_proj`,
+        ownedTrainables
+      ),
+      kProj: await ensureTrainableTensor(
+        layerWeights.kProj,
+        [numKVHeads * headDim, hiddenSize],
+        `layer_${layerIdx}.k_proj`,
+        ownedTrainables
+      ),
+      vProj: await ensureTrainableTensor(
+        layerWeights.vProj,
+        [numKVHeads * headDim, hiddenSize],
+        `layer_${layerIdx}.v_proj`,
+        ownedTrainables
+      ),
+      oProj: await ensureTrainableTensor(
+        layerWeights.oProj,
+        [hiddenSize, hiddenSize],
+        `layer_${layerIdx}.o_proj`,
+        ownedTrainables
+      ),
+      postAttentionNorm: layerWeights.postAttentionNorm
+        ? await ensureNormTensor(
+          layerWeights.postAttentionNorm,
+          hiddenSize,
+          `layer_${layerIdx}.post_attention_norm`,
+          ownedTrainables
+        )
+        : null,
+      gateUp: layerGateUp,
+      down: await ensureTrainableTensor(
+        layerWeights.down || layerWeights.ffnDown,
+        [hiddenSize, intermediateSize],
+        `layer_${layerIdx}.ffn_down`,
+        ownedTrainables
+      ),
+    };
+    layers.push(layer);
+    layerParams.push(layer.inputNorm, layer.qProj, layer.kProj, layer.vProj, layer.oProj, layer.gateUp, layer.down);
+    if (layer.postAttentionNorm) {
+      layerParams.push(layer.postAttentionNorm);
+    }
+  }
+
+  const encoderParams = [embeddingWeight, ...layerParams];
+  const decoderParams = [finalNormWeight, lmHeadWeight];
+  const baseParams = [...encoderParams, ...decoderParams];
+  const temporaryInputs = new Set();
+
+  async function buildPromptTokens(prompt) {
+    const normalized = String(prompt || '').trim();
+    if (!normalized) {
+      throw new Error('Distill full-graph student prompt is empty.');
+    }
+    const tokenIds = studentPipeline.tokenizer.encode(normalized);
+    if (!Array.isArray(tokenIds) || tokenIds.length === 0) {
+      throw new Error('Distill full-graph student tokenizer produced no tokens.');
+    }
+    const tokenTensor = makeTensorFromUint32(
+      tokenIds,
+      [tokenIds.length],
+      'distill_student_prompt_tokens'
+    );
+    temporaryInputs.add(tokenTensor);
+    return { tokenTensor, seqLen: tokenIds.length };
+  }
+
+  async function runTransformerPrompt(prompt, tape) {
+    const { tokenTensor, seqLen } = await buildPromptTokens(prompt);
+    let hidden = await tape.record(
+      OpType.EMBED,
+      (indices, embeddings) => runGather(
+        indices,
+        embeddings,
+        seqLen,
+        hiddenSize,
+        vocabSize,
+        {
+          embeddingDtype: resolveTensorDtype(embeddingWeight, 'f32'),
+          outputDtype: 'f32',
+          transpose: useEmbeddingTranspose,
+        }
+      ),
+      [tokenTensor, embeddingWeight],
+      {
+        numTokens: seqLen,
+        hiddenSize,
+        vocabSize,
+        transpose: useEmbeddingTranspose,
+        indexOffset: 0,
+      }
+    );
+
+    for (let layerIdx = 0; layerIdx < layers.length; layerIdx += 1) {
+      const layer = layers[layerIdx];
+      const normed = await tape.record(
+        OpType.RMSNORM,
+        (x, gamma) => runRMSNorm(x, gamma, rmsNormEps, {
+          batchSize: seqLen,
+          hiddenSize,
+          rmsNormWeightOffset: modelConfig.rmsNormWeightOffset === true,
+        }),
+        [hidden, layer.inputNorm],
+        { numTokens: seqLen, hiddenSize, eps: rmsNormEps }
+      );
+
+      const q2d = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, numHeads * headDim, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [normed, layer.qProj],
+        { M: seqLen, N: numHeads * headDim, K: hiddenSize, transposeB: 'auto' }
+      );
+      const k2d = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, numKVHeads * headDim, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [normed, layer.kProj],
+        { M: seqLen, N: numKVHeads * headDim, K: hiddenSize, transposeB: 'auto' }
+      );
+      const v2d = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, numKVHeads * headDim, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [normed, layer.vProj],
+        { M: seqLen, N: numKVHeads * headDim, K: hiddenSize, transposeB: 'auto' }
+      );
+
+      const q3d = createTensor(q2d.buffer, q2d.dtype, [seqLen, numHeads, headDim], `layer_${layerIdx}_q`);
+      const k3d = createTensor(k2d.buffer, k2d.dtype, [seqLen, numKVHeads, headDim], `layer_${layerIdx}_k`);
+      const v3d = createTensor(v2d.buffer, v2d.dtype, [seqLen, numKVHeads, headDim], `layer_${layerIdx}_v`);
+
+      const qRope = await tape.record(
+        OpType.ROPE,
+        (q, cos, sin) => runRoPE(q, cos, sin, seqLen, { numHeads, headDim, startPos: 0 }),
+        [q3d, ropeCos, ropeSin],
+        { seqLen, numHeads, headDim, startPos: 0 }
+      );
+      const kRope = await tape.record(
+        OpType.ROPE,
+        (k, cos, sin) => runRoPE(k, cos, sin, seqLen, { numHeads: numKVHeads, headDim, startPos: 0 }),
+        [k3d, ropeCos, ropeSin],
+        { seqLen, numHeads: numKVHeads, headDim, startPos: 0 }
+      );
+
+      const attention = await tape.record(
+        OpType.ATTENTION,
+        (q, k, v) => runAttention(q, k, v, null, numHeads, headDim, {
+          seqLen,
+          kvLen: seqLen,
+          numKVHeads,
+          causal: true,
+          startPos: 0,
+          scale: 1 / Math.sqrt(headDim),
+        }),
+        [qRope, kRope, v3d],
+        { seqLen, numHeads, headDim, scale: 1 / Math.sqrt(headDim), causal: true, recomputeForward: true }
+      );
+      const attention2d = createTensor(
+        attention.buffer,
+        attention.dtype,
+        [seqLen, hiddenSize],
+        `layer_${layerIdx}_attn_2d`
+      );
+
+      const attentionOutput = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, hiddenSize, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [attention2d, layer.oProj],
+        { M: seqLen, N: hiddenSize, K: hiddenSize, transposeB: 'auto' }
+      );
+      const postAttention = await tape.record(
+        OpType.RESIDUAL_ADD,
+        (a, b) => runResidualAdd(a, b, seqLen * hiddenSize),
+        [attentionOutput, hidden],
+        { size: seqLen * hiddenSize }
+      );
+
+      const ffnInput = layer.postAttentionNorm
+        ? await tape.record(
+          OpType.RMSNORM,
+          (x, gamma) => runRMSNorm(x, gamma, rmsNormEps, {
+            batchSize: seqLen,
+            hiddenSize,
+            rmsNormWeightOffset: modelConfig.rmsNormWeightOffset === true,
+          }),
+          [postAttention, layer.postAttentionNorm],
+          { numTokens: seqLen, hiddenSize, eps: rmsNormEps }
+        )
+        : postAttention;
+      const gateUp = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, intermediateSize * 2, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [ffnInput, layer.gateUp],
+        { M: seqLen, N: intermediateSize * 2, K: hiddenSize, transposeB: 'auto' }
+      );
+      const activated = await tape.record(
+        OpType.SILU_ROWSPLIT,
+        (x) => runSiLURowSplit(x, {
+          numTokens: seqLen,
+          dim: intermediateSize,
+          activation: hiddenActivation === 'gelu' ? 'gelu' : 'silu',
+          swigluLimit: hiddenActivation === 'gelu' ? null : swigluLimit,
+        }),
+        [gateUp],
+        {
+          numTokens: seqLen,
+          dim: intermediateSize,
+          activation: hiddenActivation === 'gelu' ? 'gelu' : 'silu',
+          swigluLimit: hiddenActivation === 'gelu' ? 0 : swigluLimit,
+        }
+      );
+      const ffnOutput = await tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, seqLen, hiddenSize, intermediateSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [activated, layer.down],
+        { M: seqLen, N: hiddenSize, K: intermediateSize, transposeB: 'auto' }
+      );
+      hidden = await tape.record(
+        OpType.RESIDUAL_ADD,
+        (a, b) => runResidualAdd(a, b, seqLen * hiddenSize),
+        [ffnOutput, postAttention],
+        { size: seqLen * hiddenSize }
+      );
+    }
+
+    const finalHidden = await tape.record(
+      OpType.RMSNORM,
+      (x, gamma) => runRMSNorm(x, gamma, rmsNormEps, {
+        batchSize: seqLen,
+        hiddenSize,
+        rmsNormWeightOffset: modelConfig.rmsNormWeightOffset === true,
+      }),
+      [hidden, finalNormWeight],
+      { numTokens: seqLen, hiddenSize, eps: rmsNormEps }
+    );
+    const lastHidden = await tape.record(
+      OpType.ROW_SLICE,
+      (x) => createRowSliceTensor(x, seqLen, hiddenSize, seqLen - 1, 'distill_last_hidden'),
+      [finalHidden],
+      { rows: seqLen, cols: hiddenSize, rowIndex: seqLen - 1 }
+    );
+    return tape.record(
+      OpType.MATMUL,
+      (x, w) => runMatmul(x, w, 1, vocabSize, hiddenSize, {
+        transposeB: 'auto',
+        outputDtype: 'f32',
+      }),
+      [lastHidden, lmHeadWeight],
+      { M: 1, N: vocabSize, K: hiddenSize, transposeB: 'auto' }
+    );
+  }
+
+  const model = {
+    async forward(inputTensor, tape) {
+      return tape.record(
+        OpType.MATMUL,
+        (x, w) => runMatmul(x, w, 1, vocabSize, hiddenSize, {
+          transposeB: 'auto',
+          outputDtype: 'f32',
+        }),
+        [inputTensor, lmHeadWeight],
+        { M: 1, N: vocabSize, K: hiddenSize, transposeB: 'auto' }
+      );
+    },
+    async forwardDistill(batch, tape, forwardOptions = {}) {
+      const requestedPhase = String(forwardOptions?.phase || 'anchor').trim();
+      const phase = requestedPhase === 'positive'
+        ? 'positive'
+        : (requestedPhase === 'negative' ? 'negative' : 'anchor');
+      const prompts = resolvePhasePrompts(batch, phase);
+      if (prompts.length !== 1) {
+        throw new Error(
+          `Distill full-graph student currently requires batchSize=1, got ${prompts.length}.`
+        );
+      }
+      const logits = await runTransformerPrompt(prompts[0], tape);
+      return { logits };
+    },
+    cleanupDistillStep() {
+      for (const tensor of temporaryInputs) {
+        releaseTensor(tensor);
+      }
+      temporaryInputs.clear();
+    },
+    loraParams() {
+      return decoderParams;
+    },
+    paramGroups() {
+      return {
+        encoder: encoderParams,
+        prior: [],
+        decoder: decoderParams,
+        base: baseParams,
+        lora: [],
+      };
+    },
+  };
+
+  return {
+    config,
+    model,
+    outputDim: vocabSize,
+    embeddingDim: hiddenSize,
+    cleanup() {
+      model.cleanupDistillStep();
+      for (const tensor of ownedTrainables) {
+        releaseTensor(tensor);
+      }
+      ownedTrainables.clear();
+    },
+  };
+}
+
+export async function createDistillStudentRuntimeModelFixture(overrides = {}, options = {}) {
+  const distillRuntime = options.distillRuntime && typeof options.distillRuntime === 'object'
+    ? options.distillRuntime
+    : null;
+  const graphMode = normalizeDistillStudentGraphMode(
+    options.studentGraphMode
+    ?? distillRuntime?.studentGraphMode
+    ?? overrides?.training?.distill?.studentGraphMode
+  );
+  if (graphMode === DISTILL_STUDENT_GRAPH_PROJECTION) {
+    return createDistillStudentProjectionModelFixture(overrides, options);
+  }
+  return createDistillStudentTransformerModelFixture(overrides, options);
+}