@simulatte/doppler 0.1.0

package/src/inference/pipelines/text/generator.js

@@ -0,0 +1,1345 @@
+
+
+import { getDevice, setTrackSubmits } from '../../../gpu/device.js';
+import { releaseBuffer, readBuffer, readBufferSlice } from '../../../memory/buffer-pool.js';
+import { isGPUSamplingAvailable } from '../../../gpu/kernels/sample.js';
+import { markWarmed as markKernelCacheWarmed } from '../../../gpu/kernel-selection-cache.js';
+import { resetSubmitStats, logSubmitStats } from '../../../gpu/submit-tracker.js';
+import { createCommandRecorder, createProfilingRecorder, CommandRecorder } from '../../../gpu/command-recorder.js';
+import { allowReadback } from '../../../gpu/perf-guards.js';
+import { log, trace } from '../../../debug/index.js';
+import { validateCallTimeOptions } from '../../../config/param-validator.js';
+import { selectRuleValue } from '../../../rules/rule-registry.js';
+
+// Pipeline sub-modules
+import { sample, applyRepetitionPenalty, logitsSanity, getTopK } from './sampling.js';
+import { enforceLogitDrift } from '../../../hotswap/intent-bundle.js';
+import { applyChatTemplate, isStopToken } from './init.js';
+import { formatChatMessages } from './chat-format.js';
+import { embed } from './embed.js';
+import { processLayer } from './layer.js';
+import { computeLogits, recordLogitsGPU, extractLastPositionLogits, applySoftcapping } from './logits/index.js';
+import { isWeightBuffer, isCpuWeightBuffer, getWeightDtype } from '../../../gpu/weight-buffer.js';
+import {
+  decodeStep,
+  decodeStepLogits,
+  advanceWithToken,
+  generateNTokensGPU,
+  shouldUseBatchDecode,
+  sumProfileTimings,
+  FinitenessError,
+  advanceWithTokenAndEmbedding as runAdvanceWithTokenAndEmbedding,
+} from './generator-steps.js';
+import { buildLayerContext, debugCheckBuffer as debugCheckBufferHelper, getLogitsConfig, getLogitsWeights } from './generator-helpers.js';
+import {
+  assertTokenIdsInRange,
+  assertTokenIdInRange,
+  resolveStepOptions,
+  resolveGenerateOptions,
+  resolvePrefillOptions,
+  resolvePrefillEmbeddingOptions,
+  resolveAdvanceEmbeddingMode,
+  getFinalNormWeights,
+  extractEmbeddingFromHidden,
+} from './generator-runtime.js';
+
+import { decodeReadback, getLogitsHealth } from './debug-utils/index.js';
+import { parseFinitenessStatusWords } from './finiteness-guard-status.js';
+import { resolveDeferredRoundingWindowTokens } from './finiteness-policy.js';
+import {
+  activateFallbackExecutionPlan,
+  rebaseExecutionSessionPlan,
+  resetActiveExecutionPlan,
+  resolveActiveExecutionPlan,
+  setActiveExecutionPlan,
+} from './execution-plan.js';
+import {
+  cloneLinearAttentionRuntime,
+  hasLinearAttentionLayers,
+  resetLinearAttentionRuntime,
+  restoreLinearAttentionRuntime,
+} from './linear-attention.js';
+
+function isStructuredChatRequest(prompt) {
+  return prompt != null
+    && typeof prompt === 'object'
+    && !Array.isArray(prompt)
+    && Array.isArray(prompt.messages);
+}
+
+function resolvePromptInput(state, prompt, useChatTemplate, contextLabel) {
+  if (typeof prompt === 'string') {
+    if (useChatTemplate && state.modelConfig.chatTemplateType) {
+      if (state.modelConfig.chatTemplateType === 'translategemma') {
+        throw new Error(
+          `[Pipeline] ${contextLabel}: translategemma chat template requires structured messages. ` +
+          'Pass { messages: [...] } instead of a plain string prompt.'
+        );
+      }
+      return applyChatTemplate(prompt, state.modelConfig.chatTemplateType);
+    }
+    return prompt;
+  }
+
+  const messages = isStructuredChatRequest(prompt)
+    ? prompt.messages
+    : (Array.isArray(prompt) ? prompt : null);
+  if (!messages) {
+    throw new Error(
+      `[Pipeline] ${contextLabel}: prompt must be a string, chat message array, or { messages: [...] }.`
+    );
+  }
+  const templateType = useChatTemplate ? state.modelConfig.chatTemplateType : null;
+  return formatChatMessages(messages, templateType);
+}
+
+function resolveTokenText(tokenizer, tokenIds, fallbackText = '?', renderTokenText, renderFallbackTokenText) {
+  const renderPrimary = typeof renderTokenText === 'function'
+    ? renderTokenText
+    : (ids) => tokenizer?.decode?.(ids);
+  const renderFallback = typeof renderFallbackTokenText === 'function'
+    ? renderFallbackTokenText
+    : (ids) => tokenizer?.decode?.(ids, false);
+
+  const primaryText = renderPrimary(tokenIds);
+  if (typeof primaryText === 'string' && primaryText.length > 0) {
+    return primaryText;
+  }
+
+  const fallback = renderFallback(tokenIds);
+  if (typeof fallback === 'string' && fallback.length > 0) {
+    // Keep skip-special behavior deterministic: if primary decoding filtered this
+    // token to empty, do not reintroduce obvious special-token text via fallback.
+    if (
+      primaryText === ''
+      && /^<[^>\n]{1,80}>$/.test(fallback.trim())
+    ) {
+      return '';
+    }
+    return fallback;
+  }
+
+  return fallbackText;
+}
+
+export class PipelineGenerator {
+
+  #state;
+  #finitenessFallbackWindow;
+
+  _assertTokenIdsInRange(tokenIds, context = 'encode') {
+    assertTokenIdsInRange(this.#state, tokenIds, context);
+  }
+
+  _assertTokenIdInRange(tokenId, context = 'token') {
+    assertTokenIdInRange(this.#state, tokenId, context);
+  }
+
+
+  constructor(state) {
+    this.#state = state;
+    this.#finitenessFallbackWindow = null;
+  }
+
+  _resolveDeferredRoundingWindowTokens() {
+    const activePlan = resolveActiveExecutionPlan(this.#state);
+    return activePlan?.deferredRoundingWindowTokens
+      ?? resolveDeferredRoundingWindowTokens(this.#state.runtimeConfig?.inference?.compute);
+  }
+
+  _getEffectiveActivationDtype() {
+    return resolveActiveExecutionPlan(this.#state).activationDtype;
+  }
+
+  _hasFinitenessFallbackWindow() {
+    return this.#finitenessFallbackWindow !== null;
+  }
+
+  _openFinitenessFallbackWindow(opts, reasonLabel, tokenCount) {
+    const normalizedCount = Number.isFinite(tokenCount)
+      ? Math.max(1, Math.floor(tokenCount))
+      : 1;
+    if (this.#finitenessFallbackWindow) {
+      this.#finitenessFallbackWindow.remainingTokens = Math.max(
+        this.#finitenessFallbackWindow.remainingTokens,
+        normalizedCount
+      );
+      return;
+    }
+    const original = this._beginFinitenessFallback(opts, reasonLabel);
+    this.#finitenessFallbackWindow = {
+      original,
+      remainingTokens: normalizedCount,
+    };
+  }
+
+  _closeFinitenessFallbackWindow(opts) {
+    if (!this.#finitenessFallbackWindow) {
+      return;
+    }
+    const original = this.#finitenessFallbackWindow.original;
+    this.#finitenessFallbackWindow = null;
+    this._endFinitenessFallback(opts, original);
+  }
+
+  _consumeFinitenessFallbackToken(opts) {
+    if (!this.#finitenessFallbackWindow) {
+      return;
+    }
+    this.#finitenessFallbackWindow.remainingTokens -= 1;
+    if (this.#finitenessFallbackWindow.remainingTokens <= 0) {
+      this._closeFinitenessFallbackWindow(opts);
+    }
+  }
+
+  _resolveStepOptions(options = {}) {
+    return resolveStepOptions(this.#state, options);
+  }
+
+  _getDecodeHelpers(debugCheckBuffer) {
+    return {
+      buildLayerContext: (recorder, isDecodeMode, debugLayers, executionPlan) =>
+        buildLayerContext(this.#state, recorder, isDecodeMode, debugLayers, debugCheckBuffer, executionPlan),
+      getLogitsWeights: () => getLogitsWeights(this.#state),
+      getLogitsConfig: () => getLogitsConfig(this.#state),
+      debugCheckBuffer,
+    };
+  }
+
+  async _getFinalNormWeights() {
+    return getFinalNormWeights(this.#state);
+  }
+
+  _extractEmbeddingFromHidden(hiddenStates, numTokens, hiddenSize, embeddingMode, finalNormWeights, config) {
+    return extractEmbeddingFromHidden(
+      hiddenStates,
+      numTokens,
+      hiddenSize,
+      embeddingMode,
+      finalNormWeights,
+      config
+    );
+  }
+
+  _beginFinitenessFallback(opts, reasonLabel) {
+    const originalPlan = resolveActiveExecutionPlan(this.#state);
+    const original = {
+      activePlanId: this.#state.executionPlanState?.activePlanId ?? 'primary',
+      seed: opts.seed,
+    };
+
+    const fallbackPlan = activateFallbackExecutionPlan(this.#state);
+    if (!fallbackPlan) {
+      throw new Error('[Pipeline] Finiteness fallback plan is unavailable for this model/runtime configuration.');
+    }
+    log.warn(
+      'Pipeline',
+      `FinitenessGuard fallback (${reasonLabel}): ` +
+      `${originalPlan.kernelPathId ?? 'none'} -> ${fallbackPlan.kernelPathId ?? 'none'}`
+    );
+
+    this.#state.decodeBuffers?.ensureBuffers({
+      hiddenSize: this.#state.modelConfig.hiddenSize,
+      intermediateSize: this.#state.modelConfig.intermediateSize,
+      activationDtype: fallbackPlan.activationDtype,
+      enablePingPong: true,
+    });
+
+    if (opts.seed == null) {
+      const fallbackSeedBase = (this.#state.decodeStepCount + this.#state.currentSeqLen + 1) >>> 0;
+      opts.seed = (fallbackSeedBase * 2654435761) >>> 0;
+    }
+    opts.executionPlan = rebaseExecutionSessionPlan(this.#state, opts.executionPlan);
+
+    return original;
+  }
+
+  _endFinitenessFallback(opts, original) {
+    opts.seed = original.seed;
+    setActiveExecutionPlan(this.#state, original.activePlanId);
+    opts.executionPlan = rebaseExecutionSessionPlan(this.#state, opts.executionPlan);
+    const nextActivationDtype = this._getEffectiveActivationDtype();
+    this.#state.decodeBuffers?.ensureBuffers({
+      hiddenSize: this.#state.modelConfig.hiddenSize,
+      intermediateSize: this.#state.modelConfig.intermediateSize,
+      activationDtype: nextActivationDtype,
+      enablePingPong: true,
+    });
+  }
+
+  async _retryWithFinitenessFallback(opts, reasonLabel, retryFn) {
+    if (this._hasFinitenessFallbackWindow()) {
+      return retryFn();
+    }
+    this.#state.kvCache?.truncate(this.#state.currentSeqLen);
+    const original = this._beginFinitenessFallback(opts, reasonLabel);
+    try {
+      return await retryFn();
+    } finally {
+      this._endFinitenessFallback(opts, original);
+    }
+  }
+
+  async _retryDecodeStepWithFinitenessWindow(generatedIds, opts, reasonLabel) {
+    const windowTokens = this._resolveDeferredRoundingWindowTokens();
+    if (windowTokens <= 1) {
+      return this._retryWithFinitenessFallback(
+        opts,
+        reasonLabel,
+        () => this._decodeStep(generatedIds, opts)
+      );
+    }
+
+    this.#state.kvCache?.truncate(this.#state.currentSeqLen);
+    this._openFinitenessFallbackWindow(opts, reasonLabel, windowTokens);
+    try {
+      return await this._decodeStep(generatedIds, opts);
+    } catch (error) {
+      this._closeFinitenessFallbackWindow(opts);
+      throw error;
+    }
+  }
+
+  // ==========================================================================
+  // Generation Public API
+  // ==========================================================================
+
+
+  async *generate(prompt, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    if (this.#state.isGenerating) throw new Error('Generation already in progress');
+
+    validateCallTimeOptions(options);
+
+    this.#state.isGenerating = true;
+    this.#state.decodeStepCount = 0;
+    this.#state.disableRecordedLogits = false;
+    this.#state.disableFusedDecode = false;
+    resetActiveExecutionPlan(this.#state);
+    this.#state.decodeRing?.reset();
+    this.#state.stats.gpuTimePrefillMs = undefined;
+    this.#state.stats.gpuTimeDecodeMs = undefined;
+    this.#state.stats.decodeRecordMs = 0;
+    this.#state.stats.decodeSubmitWaitMs = 0;
+    this.#state.stats.decodeReadbackWaitMs = 0;
+    this.#state.stats.ttftMs = 0;
+    const startTime = performance.now();
+
+    const opts = resolveGenerateOptions(this.#state, options);
+
+    if (opts.debug) {
+      log.debug('Pipeline', `ChatTemplate: options=${options.useChatTemplate}, final=${opts.useChatTemplate}`);
+    }
+
+    try {
+      const processedPrompt = resolvePromptInput(this.#state, prompt, opts.useChatTemplate, 'generate');
+      if (opts.debug && opts.useChatTemplate) {
+        log.debug('Pipeline', `Applied ${this.#state.modelConfig.chatTemplateType} chat template`);
+      }
+
+      const inputIds = this.#state.tokenizer.encode(processedPrompt);
+      this._assertTokenIdsInRange(inputIds, 'generate.encode');
+      const generatedIds = [...inputIds];
+      this.#state.stats.prefillTokens = inputIds.length;
+
+      if (opts.debug) {
+        log.debug('Pipeline', `Input: ${inputIds.length} tokens`);
+      }
+
+      const prefillStart = performance.now();
+      let prefillLogits;
+      try {
+        prefillLogits = await this._prefill(inputIds, opts);
+      } catch (error) {
+        if (error.name === 'FinitenessError') {
+          log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefill. Retrying with F32 precision.`);
+          prefillLogits = await this._retryWithFinitenessFallback(
+            opts,
+            'prefill',
+            () => this._prefill(inputIds, opts)
+          );
+        } else {
+          throw error;
+        }
+      }
+      this.#state.stats.prefillTimeMs = performance.now() - prefillStart;
+
+      const intentBundleConfig = this.#state.runtimeConfig.shared.intentBundle;
+      const intentBundle = intentBundleConfig?.bundle;
+      const expectedTopK = intentBundle?.payload?.expectedTopK
+        ?? intentBundle?.payload?.expected_top_k;
+      const maxDriftThreshold = intentBundle?.constraints?.maxDriftThreshold
+        ?? intentBundle?.constraints?.max_drift_threshold;
+
+      if (intentBundleConfig?.enabled && Array.isArray(expectedTopK) && expectedTopK.length > 0) {
+        const actualTopK = getTopK(
+          prefillLogits,
+          expectedTopK.length,
+        (tokens) => resolveTokenText(this.#state.tokenizer, tokens),
+      ).map((token) => token.token);
+        const driftResult = enforceLogitDrift(expectedTopK, actualTopK, maxDriftThreshold);
+        if (!driftResult.ok) {
+          throw new Error(`Intent bundle drift check failed: ${driftResult.reason}`);
+        }
+      }
+
+      applyRepetitionPenalty(prefillLogits, generatedIds, opts.repetitionPenalty);
+      const padTokenId = this.#state.tokenizer?.getSpecialTokens?.()?.pad;
+
+      if (opts.debug) {
+        const topAfterPenalty = getTopK(
+          prefillLogits,
+          5,
+          (tokens) => resolveTokenText(this.#state.tokenizer, tokens)
+        );
+        log.debug('Pipeline', `After rep penalty top-5: ${topAfterPenalty.map(t => `"${t.text}"(${(t.prob * 100).toFixed(1)}%)`).join(', ')}`);
+      }
+
+      const firstToken = sample(prefillLogits, {
+        temperature: opts.temperature,
+        topP: opts.topP,
+        topK: opts.topK,
+        padTokenId,
+        seed: opts.seed,
+      });
+
+      if (opts.debug) {
+        const firstTokenText = resolveTokenText(this.#state.tokenizer, [firstToken], `[${firstToken}]`, (tokens) => this.#state.tokenizer?.decode?.(tokens, true, false));
+        log.debug('Pipeline', `First token sampled: id=${firstToken} text="${firstTokenText}"`);
+      }
+
+      generatedIds.push(firstToken);
+      this.#state.stats.ttftMs = performance.now() - startTime;
+
+      const decodeToken = (tokenId) => resolveTokenText(
+        this.#state.tokenizer,
+        [tokenId],
+        `[${tokenId}]`,
+        (tokens) => this.#state.tokenizer?.decode?.(tokens, true, false),
+        (tokens) => this.#state.tokenizer?.decode?.(tokens, false, false)
+      );
+
+      const firstText = decodeToken(firstToken);
+      yield firstText;
+      if (options.onToken) options.onToken(firstToken, firstText);
+
+      yield* this._runDecodeLoop(generatedIds, opts, options, {
+        stopTokenIds: this.#state.modelConfig.stopTokenIds,
+        eosToken: this.#state.tokenizer.getSpecialTokens?.()?.eos,
+        stopSequenceStart: inputIds.length,
+        decodeToken,
+        logBatchPath: opts.debug,
+      });
+      const tokensGenerated = this.#state.stats.decodeTokens;
+      this.#state.stats.totalTimeMs = performance.now() - startTime;
+
+      if (opts.debug) {
+        log.debug('Pipeline', `Generated ${tokensGenerated} tokens in ${this.#state.stats.totalTimeMs.toFixed(0)}ms`);
+      }
+
+      const ttft = this.#state.stats.ttftMs ?? this.#state.stats.prefillTimeMs;
+      const decodeTokens = Math.max(0, tokensGenerated - 1);
+      const decodeSpeed = decodeTokens > 0 ? (decodeTokens / this.#state.stats.decodeTimeMs * 1000) : 0;
+      if (opts.benchmark) {
+        log.info('Benchmark', `TTFT: ${ttft.toFixed(0)}ms | Prefill: ${this.#state.stats.prefillTimeMs.toFixed(0)}ms | Decode: ${this.#state.stats.decodeTimeMs.toFixed(0)}ms (${decodeTokens} tokens @ ${decodeSpeed.toFixed(1)} tok/s)`);
+      } else {
+        log.info('Perf', `TTFT: ${ttft.toFixed(0)}ms | Prefill: ${this.#state.stats.prefillTimeMs.toFixed(0)}ms | Decode: ${this.#state.stats.decodeTimeMs.toFixed(0)}ms (${decodeTokens} tokens @ ${decodeSpeed.toFixed(1)} tok/s)`);
+      }
+      trace.perf('Decode summary', {
+        ttftMs: ttft,
+        prefillMs: this.#state.stats.prefillTimeMs,
+        decodeMs: this.#state.stats.decodeTimeMs,
+        decodeTokens,
+        decodeSpeed,
+        totalMs: this.#state.stats.totalTimeMs,
+      });
+    } finally {
+      this._closeFinitenessFallbackWindow(opts);
+      resetActiveExecutionPlan(this.#state);
+      this.#state.isGenerating = false;
+    }
+  }
+
+
+  async prefillKVOnly(prompt, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    resetActiveExecutionPlan(this.#state);
+    this.#state.stats.gpuTimePrefillMs = undefined;
+    const opts = resolvePrefillOptions(this.#state, options);
+
+    const processedPrompt = resolvePromptInput(this.#state, prompt, opts.useChatTemplate, 'prefillKVOnly');
+
+    const inputIds = this.#state.tokenizer.encode(processedPrompt);
+    this._assertTokenIdsInRange(inputIds, 'prefillKVOnly.encode');
+    if (opts.debug) {
+      log.debug('Pipeline', `PrefillKVOnly: ${inputIds.length} tokens`);
+    }
+
+    let prefillResult;
+    try {
+      prefillResult = await this._prefillToHidden(inputIds, opts);
+    } catch (error) {
+      if (error.name === 'FinitenessError') {
+        log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefillKVOnly. Retrying with F32 precision.`);
+        prefillResult = await this._retryWithFinitenessFallback(
+          opts,
+          'prefillKVOnly',
+          () => this._prefillToHidden(inputIds, opts)
+        );
+      } else {
+        throw error;
+      }
+    }
+
+    const {
+      numTokens,
+      startPos,
+      currentRecorder,
+      recordProfile,
+      currentHiddenBuffer,
+    } = prefillResult;
+
+    // Ensure prefill work completes before returning a usable snapshot.
+    if (currentRecorder) {
+      await currentRecorder.submitAndWait();
+      await recordProfile(currentRecorder);
+    } else {
+      const device = getDevice();
+      if (device) {
+        await device.queue.onSubmittedWorkDone();
+      }
+    }
+
+    this.#state.currentSeqLen = startPos + numTokens;
+    releaseBuffer(currentHiddenBuffer);
+
+    const snapshot = this.#state.kvCache?.clone();
+    if (!snapshot) {
+      throw new Error('KV cache unavailable after prefill');
+    }
+
+    return {
+      cache: snapshot,
+      seqLen: this.#state.currentSeqLen,
+      tokens: inputIds,
+      linearAttention: await cloneLinearAttentionRuntime(this.#state.linearAttentionRuntime),
+    };
+  }
+
+  async prefillWithEmbedding(prompt, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    resetActiveExecutionPlan(this.#state);
+    this.#state.stats.gpuTimePrefillMs = undefined;
+    const opts = resolvePrefillEmbeddingOptions(this.#state, options);
+
+    const processedPrompt = resolvePromptInput(this.#state, prompt, opts.useChatTemplate, 'prefillWithEmbedding');
+
+    const inputIds = this.#state.tokenizer.encode(processedPrompt);
+    this._assertTokenIdsInRange(inputIds, 'prefillWithEmbedding.encode');
+    if (opts.debug) {
+      log.debug('Pipeline', `PrefillWithEmbedding: ${inputIds.length} tokens (mode=${opts.embeddingMode})`);
+    }
+
+    let prefillResult;
+    try {
+      prefillResult = await this._prefillToHidden(inputIds, opts);
+    } catch (error) {
+      if (error.name === 'FinitenessError') {
+        log.warn('Pipeline', `FinitenessGuard caught NaN/Inf during prefillWithEmbedding. Retrying with F32 precision.`);
+        prefillResult = await this._retryWithFinitenessFallback(
+          opts,
+          'prefillWithEmbedding',
+          () => this._prefillToHidden(inputIds, opts)
+        );
+      } else {
+        throw error;
+      }
+    }
+
+    const {
+      numTokens,
+      config,
+      startPos,
+      activationDtype,
+      activationBytes,
+      currentRecorder,
+      recordProfile,
+      currentHiddenBuffer,
+    } = prefillResult;
+
+    // Ensure prefill work completes before readback.
+    if (currentRecorder) {
+      await currentRecorder.submitAndWait();
+      await recordProfile(currentRecorder);
+    } else {
+      const device = getDevice();
+      if (device) {
+        await device.queue.onSubmittedWorkDone();
+      }
+    }
+
+    if (!allowReadback('pipeline.prefill.embedding')) {
+      throw new Error('GPU readback disabled; cannot return embedding');
+    }
+
+    let embedding;
+    try {
+      const hiddenSize = config.hiddenSize;
+      const hiddenBytes = numTokens * hiddenSize * activationBytes;
+      const hiddenData = await readBuffer(currentHiddenBuffer, hiddenBytes);
+      if (hiddenData.byteLength === 0) {
+        throw new Error('GPU readback disabled; cannot return embedding');
+      }
+      const hiddenStates = decodeReadback(hiddenData, activationDtype);
+      const finalNormWeights = await this._getFinalNormWeights();
+      embedding = this._extractEmbeddingFromHidden(
+        hiddenStates,
+        numTokens,
+        hiddenSize,
+        opts.embeddingMode,
+        finalNormWeights,
+        config
+      );
+    } finally {
+      releaseBuffer(currentHiddenBuffer);
+    }
+
+    this.#state.currentSeqLen = startPos + numTokens;
+
+    const snapshot = this.#state.kvCache?.clone();
+    if (!snapshot) {
+      throw new Error('KV cache unavailable after prefill');
+    }
+
+    return {
+      cache: snapshot,
+      seqLen: this.#state.currentSeqLen,
+      tokens: inputIds,
+      embedding,
+      embeddingMode: opts.embeddingMode,
+      linearAttention: await cloneLinearAttentionRuntime(this.#state.linearAttentionRuntime),
+    };
+  }
+
+  async prefillWithLogits(prompt, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    resetActiveExecutionPlan(this.#state);
+    this.#state.stats.gpuTimePrefillMs = undefined;
+    const opts = resolvePrefillOptions(this.#state, options);
+
+    const processedPrompt = resolvePromptInput(this.#state, prompt, opts.useChatTemplate, 'prefillWithLogits');
+
+    const inputIds = this.#state.tokenizer.encode(processedPrompt);
+    this._assertTokenIdsInRange(inputIds, 'prefillWithLogits.encode');
+    if (opts.debug) {
+      log.debug('Pipeline', `PrefillWithLogits: ${inputIds.length} tokens`);
+    }
+
+    const logits = await this._prefill(inputIds, opts);
+
+    const snapshot = this.#state.kvCache?.clone();
+    if (!snapshot) {
+      throw new Error('KV cache unavailable after prefill');
+    }
+
+    return {
+      cache: snapshot,
+      seqLen: this.#state.currentSeqLen,
+      tokens: inputIds,
+      logits,
+      linearAttention: await cloneLinearAttentionRuntime(this.#state.linearAttentionRuntime),
+    };
+  }
+
+
+  async *generateWithPrefixKV(prefix, prompt, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    if (this.#state.isGenerating) throw new Error('Generation already in progress');
+
+    validateCallTimeOptions(options);
+
+    // Apply snapshot
+    this.#state.kvCache = prefix.cache.clone();
+    if (this.#state.useGPU && this.#state.kvCache) {
+      const device = getDevice();
+      if (device) {
+        this.#state.kvCache.setGPUContext({ device });
+      }
+    }
+    if (
+      hasLinearAttentionLayers(this.#state.modelConfig.layerTypes)
+      && prefix.linearAttention == null
+    ) {
+      throw new Error(
+        'Prefix snapshot is missing linear_attention recurrent state. ' +
+        'Regenerate the prefix snapshot using the current runtime.'
+      );
+    }
+    this.#state.linearAttentionRuntime = restoreLinearAttentionRuntime(
+      this.#state.linearAttentionRuntime,
+      prefix.linearAttention ?? null
+    );
+    this.#state.currentSeqLen = prefix.seqLen;
+
+    this.#state.isGenerating = true;
+    this.#state.decodeStepCount = 0;
+    resetActiveExecutionPlan(this.#state);
+    this.#state.stats.gpuTimePrefillMs = undefined;
+    this.#state.stats.gpuTimeDecodeMs = undefined;
+    this.#state.decodeRing?.reset();
+    this.#state.stats.decodeRecordMs = 0;
+    this.#state.stats.decodeSubmitWaitMs = 0;
+    this.#state.stats.decodeReadbackWaitMs = 0;
+    this.#state.stats.ttftMs = 0;
+    const startTime = performance.now();
+
+    const opts = resolveGenerateOptions(this.#state, options);
+
+    try {
+      const processedPrompt = resolvePromptInput(this.#state, prompt, opts.useChatTemplate, 'generateWithPrefixKV');
+
+      const inputIds = this.#state.tokenizer.encode(processedPrompt);
+      this._assertTokenIdsInRange(inputIds, 'generateWithPrefixKV.encode');
+      const generatedIds = [...prefix.tokens, ...inputIds];
+      const promptTokenCount = generatedIds.length;
+      this.#state.stats.prefillTokens = inputIds.length;
+
+      const prefillStart = performance.now();
+      const prefillLogits = await this._prefill(inputIds, opts);
+      this.#state.stats.prefillTimeMs = performance.now() - prefillStart;
+
+      applyRepetitionPenalty(prefillLogits, generatedIds, opts.repetitionPenalty);
+      const padTokenId = this.#state.tokenizer?.getSpecialTokens?.()?.pad;
+      const firstToken = sample(prefillLogits, {
+        temperature: opts.temperature,
+        topP: opts.topP,
+        topK: opts.topK,
+        padTokenId,
+        seed: opts.seed,
+      });
+
+      generatedIds.push(firstToken);
+      this.#state.stats.ttftMs = performance.now() - startTime;
+
+      const firstText = resolveTokenText(
+        this.#state.tokenizer,
+        [firstToken],
+        `[${firstToken}]`,
+        (tokens) => this.#state.tokenizer?.decode?.(tokens, true, false),
+        (tokens) => this.#state.tokenizer?.decode?.(tokens, false, false)
+      );
+      yield firstText;
+      if (options.onToken) options.onToken(firstToken, firstText);
+
+      yield* this._runDecodeLoop(generatedIds, opts, options, {
+        stopTokenIds: this.#state.modelConfig.stopTokenIds,
+        eosToken: this.#state.tokenizer.getSpecialTokens?.()?.eos,
+        stopSequenceStart: promptTokenCount,
+        decodeToken: (tokenId) => this.#state.tokenizer.decode([tokenId], true, false),
+        logBatchPath: false,
+      });
+      this.#state.stats.totalTimeMs = performance.now() - startTime;
+    } finally {
+      this._closeFinitenessFallbackWindow(opts);
+      resetActiveExecutionPlan(this.#state);
+      this.#state.isGenerating = false;
+    }
+  }
+
+  // ==========================================================================
+  // Internal Methods (Prefill, Decode, Helpers)
+  // ==========================================================================
+
+  async *_runDecodeLoop(generatedIds, opts, options, runtime) {
+    const {
+      stopTokenIds,
+      eosToken,
+      stopSequenceStart,
+      decodeToken,
+      logBatchPath = false,
+    } = runtime;
+
+    let tokensGenerated = 1;
+    markKernelCacheWarmed();
+
+    const decodeStart = performance.now();
+    const lmHead = this.#state.weights.get('lm_head');
+    const embedBuffer = this.#state.weights.get('embed');
+    const hasCpuWeights = isCpuWeightBuffer(lmHead)
+      || isCpuWeightBuffer(embedBuffer)
+      || lmHead instanceof Float32Array
+      || embedBuffer instanceof Float32Array;
+    const hasLinearLayers = hasLinearAttentionLayers(this.#state.modelConfig.layerTypes);
+    const gpuSamplingAvailable = isGPUSamplingAvailable() && !hasCpuWeights;
+    const executionPlan = opts.executionPlan;
+    let useBatchPath = shouldUseBatchDecode({
+      batchSize: executionPlan.batchSize,
+      useGPU: this.#state.useGPU,
+      gpuSamplingAvailable,
+      disableMultiTokenDecode: executionPlan.disableMultiTokenDecode,
+      disableCommandBatching: executionPlan.disableCommandBatching,
+      isBdpaPagedLayout: this.#state.kvCache?.layout === 'bdpa_paged',
+      finitenessFallbackWindowOpen: this._hasFinitenessFallbackWindow(),
+    });
+    if (hasLinearLayers) {
+      useBatchPath = false;
+    }
+    const readbackInterval = executionPlan.readbackInterval;
+    const intervalBatches = readbackInterval == null ? 1 : readbackInterval;
+
+    if (logBatchPath && useBatchPath) {
+      log.debug(
+        'Pipeline',
+        `Using batch decode path with batchSize=${executionPlan.batchSize}, stopCheckMode=${executionPlan.stopCheckMode}, readbackInterval=${readbackInterval}`
+      );
+    }
+
+    while (tokensGenerated < opts.maxTokens) {
+      if (options.signal?.aborted) break;
+      if (this._hasFinitenessFallbackWindow() && useBatchPath) {
+        useBatchPath = false;
+      }
+
+      if (useBatchPath) {
+        const remaining = opts.maxTokens - tokensGenerated;
+        const thisBatchSize = Math.min(executionPlan.batchSize * intervalBatches, remaining);
+        const lastToken = generatedIds[generatedIds.length - 1];
+
+        try {
+          const batchResult = await this._generateNTokensGPU(lastToken, thisBatchSize, generatedIds, opts);
+          let batchTokens = [];
+          for (const tokenId of batchResult.tokens) {
+            generatedIds.push(tokenId);
+            tokensGenerated++;
+            const tokenText = decodeToken(tokenId);
+            yield tokenText;
+            if (options.onToken) options.onToken(tokenId, tokenText);
+            batchTokens.push({ id: tokenId, text: tokenText });
+            if (batchTokens.length === executionPlan.batchSize) {
+              if (options.onBatch) options.onBatch(batchTokens);
+              batchTokens = [];
+            }
+          }
+          if (batchTokens.length > 0 && options.onBatch) options.onBatch(batchTokens);
+          if (batchResult.actualCount < thisBatchSize) break;
+          if (opts.stopSequences.length > 0) {
+            const fullText = this.#state.tokenizer.decode(generatedIds.slice(stopSequenceStart), false);
+            if (opts.stopSequences.some((seq) => fullText.endsWith(seq))) break;
+          }
+        } catch (error) {
+          log.warn('Pipeline', `Batch decode failed, falling back to single-token: ${error}`);
+          useBatchPath = false;
+          let nextToken;
+          try {
+            nextToken = await this._decodeStep(generatedIds, opts);
+          } catch (singleTokenError) {
+            if (singleTokenError.name === 'FinitenessError') {
+              log.warn('Pipeline', `FinitenessGuard caught NaN/Inf at batch step ${tokensGenerated}. Truncating KV cache and retrying token with F32 precision.`);
+              nextToken = await this._retryDecodeStepWithFinitenessWindow(
+                generatedIds,
+                opts,
+                `decode-batch-step-${tokensGenerated}`
+              );
+            } else {
+              throw singleTokenError;
+            }
+          }
+          generatedIds.push(nextToken);
+          tokensGenerated++;
+          const tokenText = decodeToken(nextToken);
+          yield tokenText;
+          if (options.onToken) options.onToken(nextToken, tokenText);
+          this._consumeFinitenessFallbackToken(opts);
+          if (isStopToken(nextToken, stopTokenIds, eosToken)) break;
+        }
+      } else {
+        const tokenStart = performance.now();
+        let nextToken;
+        try {
+          nextToken = await this._decodeStep(generatedIds, opts);
+        } catch (error) {
+          if (error.name === 'FinitenessError') {
+            log.warn('Pipeline', `FinitenessGuard caught NaN/Inf at step ${tokensGenerated}. Truncating KV cache and retrying token with F32 precision.`);
+            nextToken = await this._retryDecodeStepWithFinitenessWindow(
+              generatedIds,
+              opts,
+              `decode-step-${tokensGenerated}`
+            );
+          } else {
+            throw error;
+          }
+        }
+        const tokenTime = performance.now() - tokenStart;
+        generatedIds.push(nextToken);
+        tokensGenerated++;
+        const tokenText = decodeToken(nextToken);
+        yield tokenText;
+        if (options.onToken) options.onToken(nextToken, tokenText);
+        this._consumeFinitenessFallbackToken(opts);
+
+        if (opts.debug || opts.benchmark) {
+          const elapsedMs = performance.now() - decodeStart;
+          const tokPerSec = (tokensGenerated / elapsedMs) * 1000;
+          log.debug('Decode', `#${tokensGenerated} "${tokenText}" ${tokenTime.toFixed(0)}ms (${tokPerSec.toFixed(2)} tok/s avg)`);
+        }
+
+        if (isStopToken(nextToken, stopTokenIds, eosToken)) break;
+        if (opts.stopSequences.length > 0) {
+          const fullText = this.#state.tokenizer.decode(generatedIds.slice(stopSequenceStart), false);
+          if (opts.stopSequences.some((seq) => fullText.endsWith(seq))) break;
+        }
+      }
+    }
+
+    this.#state.stats.decodeTimeMs = performance.now() - decodeStart;
+    this.#state.stats.tokensGenerated = tokensGenerated;
+    this.#state.stats.decodeTokens = tokensGenerated;
+  }
+
+  async _prefillToHidden(inputIds, opts) {
+    // Internal-only: reuse the main prefill implementation but stop before logits.
+    return this._prefill(inputIds, { ...opts, _returnHidden: true });
+  }
+
+
+  async _prefill(inputIds, opts) {
+    const numTokens = inputIds.length;
+    const config = this.#state.modelConfig;
+    const startPos = this.#state.currentSeqLen;
+    const returnHidden = opts?._returnHidden === true;
+    this.#state.stats.gpuTimePrefillMs = undefined;
+
+    if (startPos === 0 && hasLinearAttentionLayers(config.layerTypes)) {
+      this.#state.linearAttentionRuntime = resetLinearAttentionRuntime(this.#state.linearAttentionRuntime);
+    }
+
+    const embedBufferRaw = this.#state.weights.get('embed');
+    if (!(embedBufferRaw instanceof GPUBuffer) && !isWeightBuffer(embedBufferRaw) && !isCpuWeightBuffer(embedBufferRaw) && !(embedBufferRaw instanceof Float32Array)) {
+      throw new Error('Embed buffer not found or not a supported buffer type');
+    }
+    const embedBuffer = isWeightBuffer(embedBufferRaw) ? embedBufferRaw.buffer : embedBufferRaw;
+    const embedDtype = isWeightBuffer(embedBufferRaw)
+      ? getWeightDtype(embedBufferRaw)
+      : isCpuWeightBuffer(embedBufferRaw)
+        ? embedBufferRaw.dtype
+        : null;
+    if (opts.debug) {
+      const embedSize = embedBuffer instanceof GPUBuffer ? embedBuffer.size : 'N/A';
+      log.debug('Pipeline', `Embed buffer: type=${embedBuffer?.constructor?.name}, size=${embedSize}, dtype=${embedDtype}`);
+    }
+
+    const device = getDevice();
+    const useCheckpoints = opts.debugLayers && opts.debugLayers.length > 0;
+    const disableCommandBatching = opts.disableCommandBatching === true
+      || opts.debug === true
+      || this.#state.kvCache?.layout === 'bdpa_paged';
+    const createRecorder = (label) => {
+      if (!device || disableCommandBatching) return undefined;
+      return opts.profile ? createProfilingRecorder(label) : createCommandRecorder(label);
+    };
+    const recorder = createRecorder('prefill');
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+    const context = buildLayerContext(
+      this.#state,
+      recorder,
+      false,
+      opts.debugLayers,
+      debugCheckBuffer,
+      opts.executionPlan
+    );
+    context.currentTokenIds = inputIds;
+    let gpuTimePrefillMs = 0;
+    let hasGpuTimePrefill = false;
+    const recordProfile = async (rec) => {
+      if (!opts.profile || !rec?.isProfilingEnabled()) return;
+      const timings = await rec.resolveProfileTimings();
+      const total = sumProfileTimings(timings);
+      if (total !== null) {
+        gpuTimePrefillMs += total;
+        hasGpuTimePrefill = true;
+      }
+      if (timings) {
+        log.warn('Profile', `Prefill (${rec.label}):`);
+        log.warn('Profile', CommandRecorder.formatProfileReport(timings));
+      }
+    };
+
+    const benchmarkSubmits = opts.debug;
+    if (benchmarkSubmits) {
+      setTrackSubmits(true);
+      resetSubmitStats();
+    }
+
+    const activationDtype = opts.executionPlan?.activationDtype ?? this._getEffectiveActivationDtype();
+    const activationBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: activationDtype });
+    let hiddenStates = await embed(inputIds, embedBuffer, {
+      hiddenSize: config.hiddenSize,
+      vocabSize: config.vocabSize,
+      scaleEmbeddings: config.scaleEmbeddings,
+      debug: opts.debug,
+      recorder,
+      transpose: this.#state.embeddingTranspose,
+      debugProbes: this.#state.runtimeConfig.shared.debug.probes,
+      activationDtype,
+      embeddingDtype: selectRuleValue('inference', 'dtype', 'f16OrF32FromDtype', { dtype: embedDtype }),
+    });
+
+    if (opts.debug && hiddenStates instanceof GPUBuffer) {
+      if (recorder) {
+        await recorder.submitAndWait();
+        await recordProfile(recorder);
+      }
+      const debugReadbackSize = this.#state.runtimeConfig.shared.debug.pipeline.readbackSampleSize;
+      const sample = await readBuffer(hiddenStates, Math.min(debugReadbackSize, hiddenStates.size));
+      const f32 = decodeReadback(sample, activationDtype);
+      const nanCount = f32.filter(x => !Number.isFinite(x)).length;
+      let maxAbs = 0;
+      for (let i = 0; i < f32.length; i++) {
+        const abs = Math.abs(f32[i]);
+        if (abs > maxAbs) maxAbs = abs;
+      }
+      const first8 = Array.from(f32).slice(0, 8).map(x => x.toFixed(4)).join(', ');
+      log.debug('Pipeline', `After embed: buffer.label=${hiddenStates.label}, buffer.size=${hiddenStates.size}, maxAbs=${maxAbs.toFixed(4)}`);
+      log.debug('Pipeline', `After embed first8=[${first8}], nan=${nanCount}/${f32.length}`);
+    }
+
+    if (opts.debug) {
+      log.debug('Pipeline', `LAYER_LOOP_START: numLayers=${config.numLayers}, useGPU=${context.useGPU}`);
+    }
+
+    if (this.#state.finitenessBuffer) {
+      const device = getDevice();
+      if (device) {
+        device.queue.writeBuffer(this.#state.finitenessBuffer, 0, new Uint32Array([0, 0, 0, 0]));
+      }
+    }
+
+    let currentRecorder = recorder;
+
+    let currentHiddenBuffer = hiddenStates.buffer;
+    for (let l = 0; l < config.numLayers; l++) {
+      context.recorder = currentRecorder;
+
+      const prevBuffer = currentHiddenBuffer;
+      const layerOutput = await processLayer(l, currentHiddenBuffer, numTokens, true, context);
+      if (!(layerOutput instanceof GPUBuffer)) throw new Error('Expected GPUBuffer from processLayer');
+      currentHiddenBuffer = layerOutput;
+
+      const isCheckpoint = useCheckpoints && opts.debugLayers?.includes(l);
+
+      if (isCheckpoint && currentRecorder) {
+        await currentRecorder.submitAndWait();
+        await recordProfile(currentRecorder);
+        currentRecorder = undefined;
+      }
+
+      const shouldDebug = opts.debug && currentHiddenBuffer && (!recorder || isCheckpoint);
+      if (shouldDebug && !currentRecorder) {
+        const device = getDevice();
+        if (device) {
+          if (allowReadback(`pipeline.prefill.layer-${l}`)) {
+            try {
+              const sampleSize = config.hiddenSize * activationBytes;
+              const staging = device.createBuffer({
+                size: sampleSize,
+                usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+              });
+              const enc = device.createCommandEncoder();
+              const lastTokenOffset = (numTokens - 1) * config.hiddenSize * activationBytes;
+              enc.copyBufferToBuffer(currentHiddenBuffer, lastTokenOffset, staging, 0, sampleSize);
+              device.queue.submit([enc.finish()]);
+              await staging.mapAsync(GPUMapMode.READ);
+              const data = decodeReadback(staging.getMappedRange().slice(0), activationDtype);
+              staging.unmap();
+              staging.destroy();
+              let min = Infinity;
+              let max = -Infinity;
+              let maxAbs = 0;
+              for (const v of data) {
+                if (!Number.isFinite(v)) continue;
+                if (v < min) min = v;
+                if (v > max) max = v;
+                const av = Math.abs(v);
+                if (av > maxAbs) maxAbs = av;
+              }
+              const sample = Array.from(data).slice(0, 3).map(x => x.toFixed(3)).join(', ');
+              log.debug('Pipeline', `LAYER_${l}_LAST[pos=${numTokens - 1}]: min=${min.toFixed(3)}, max=${max.toFixed(3)}, maxAbs=${maxAbs.toFixed(2)}, sample=[${sample}]`);
+            } catch (e) {
+              log.debug('Pipeline', `LAYER_${l}_LAST: error reading buffer: ${e}`);
+            }
+          }
+        }
+      }
+
+      if (isCheckpoint && useCheckpoints && l < config.numLayers - 1) {
+        currentRecorder = createRecorder('prefill-cont');
+      }
+
+      if (prevBuffer !== currentHiddenBuffer) {
+        if (currentRecorder) {
+          currentRecorder.trackTemporaryBuffer(prevBuffer);
+        } else {
+          releaseBuffer(prevBuffer);
+        }
+      }
+    }
+
+    if (this.#state.finitenessBuffer) {
+      if (currentRecorder) {
+        await currentRecorder.submitAndWait();
+        await recordProfile(currentRecorder);
+        currentRecorder = undefined;
+      }
+      const isInfiniteData = await readBuffer(this.#state.finitenessBuffer, 16);
+      const u32 = new Uint32Array(isInfiniteData.buffer, isInfiniteData.byteOffset, 4);
+      const finitenessStatus = parseFinitenessStatusWords(u32, 0);
+      if (finitenessStatus.triggered) {
+        if (currentHiddenBuffer instanceof GPUBuffer) {
+          releaseBuffer(currentHiddenBuffer);
+        }
+        throw new FinitenessError(`F16 bounds exceeded during prefill${finitenessStatus.metadata}`);
+      }
+    }
+
+    if (benchmarkSubmits) {
+      logSubmitStats(`Prefill (${numTokens} tokens, ${config.numLayers} layers)`);
+      setTrackSubmits(false);
+    }
+
+    if (opts.debug) {
+      log.debug('Pipeline', `LAYER_LOOP_DONE, currentHiddenBuffer type=${currentHiddenBuffer?.constructor?.name}`);
+      if (currentHiddenBuffer && allowReadback('pipeline.prefill.final-hidden')) {
+        const device = getDevice();
+        const lastTokenOffset = (numTokens - 1) * config.hiddenSize * activationBytes;
+        const sampleSize = config.hiddenSize * activationBytes;
+        const staging = device.createBuffer({
+          size: sampleSize,
+          usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+        });
+        const enc = device.createCommandEncoder();
+        enc.copyBufferToBuffer(currentHiddenBuffer, lastTokenOffset, staging, 0, sampleSize);
+        device.queue.submit([enc.finish()]);
+        await staging.mapAsync(GPUMapMode.READ);
+        const data = decodeReadback(staging.getMappedRange().slice(0), activationDtype);
+        staging.unmap();
+        staging.destroy();
+        const nanCount = Array.from(data).filter(x => !Number.isFinite(x)).length;
+        const nonZero = Array.from(data).filter(x => Number.isFinite(x) && x !== 0).slice(0, 5);
+        log.debug('Pipeline', `FINAL_HIDDEN[pos=${numTokens - 1}]: nan=${nanCount}/${data.length}, sample=[${nonZero.map(x => x.toFixed(4)).join(', ')}]`);
+      }
+    }
+
+    if (hasGpuTimePrefill) {
+      this.#state.stats.gpuTimePrefillMs = gpuTimePrefillMs;
+    }
+
+    if (returnHidden) {
+      return {
+        numTokens,
+        config,
+        startPos,
+        activationDtype,
+        activationBytes,
+        currentRecorder,
+        recordProfile,
+        debugCheckBuffer,
+        currentHiddenBuffer,
+      };
+    }
+
+
+    let lastLogits;
+    let logitsVocabSize = config.vocabSize;
+    let usedRecordedLogits = false;
+    const lmHead = this.#state.weights.get('lm_head');
+    const canRecordLogits = !!currentRecorder
+      && !!lmHead
+      && !isCpuWeightBuffer(lmHead)
+      && !this.#state.disableRecordedLogits
+      && numTokens === 1;
+    if (currentRecorder && canRecordLogits) {
+      const recorded = await recordLogitsGPU(
+        currentRecorder,
+        currentHiddenBuffer,
+        numTokens,
+        getLogitsWeights(this.#state),
+        getLogitsConfig(this.#state)
+      );
+      logitsVocabSize = recorded.vocabSize;
+      usedRecordedLogits = true;
+
+      await currentRecorder.submitAndWait();
+      await recordProfile(currentRecorder);
+
+      const logitsBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: recorded.logitsDtype });
+      const lastLogitsSize = logitsVocabSize * logitsBytes;
+      const lastLogitsOffset = (numTokens - 1) * lastLogitsSize;
+      const logitsData = await readBufferSlice(recorded.logitsBuffer, lastLogitsOffset, lastLogitsSize);
+      releaseBuffer(recorded.logitsBuffer);
+      lastLogits = decodeReadback(logitsData, recorded.logitsDtype);
+
+      const health = getLogitsHealth(lastLogits);
+      if (health.nanCount > 0 || health.infCount > 0 || health.nonZeroCount === 0) {
+        log.warn(
+          'Logits',
+          `Recorded logits invalid (nan=${health.nanCount} inf=${health.infCount} nonZero=${health.nonZeroCount}, maxAbs=${health.maxAbs.toFixed(3)}); recomputing without recorder.`
+        );
+        this.#state.disableRecordedLogits = true;
+        this.#state.disableFusedDecode = true;
+        const fallbackLogits = await computeLogits(
+          currentHiddenBuffer,
+          numTokens,
+          getLogitsWeights(this.#state),
+          getLogitsConfig(this.#state),
+          this.#state.useGPU,
+          this.#state.debugFlags,
+          undefined,
+          debugCheckBuffer,
+          this.#state.runtimeConfig.shared.debug.probes,
+          { lastPositionOnly: true }
+        );
+        const fallbackHealth = getLogitsHealth(fallbackLogits);
+        if (fallbackHealth.nanCount > 0 || fallbackHealth.infCount > 0 || fallbackHealth.nonZeroCount === 0) {
+          throw new Error(
+            `[Logits] Fallback logits invalid (nan=${fallbackHealth.nanCount} inf=${fallbackHealth.infCount} nonZero=${fallbackHealth.nonZeroCount}, maxAbs=${fallbackHealth.maxAbs.toFixed(3)}). ` +
+            'This indicates upstream kernel output is NaN/Inf (often prefill attention/matmul).'
+          );
+        }
+        logitsVocabSize = config.vocabSize;
+        usedRecordedLogits = false;
+        lastLogits = fallbackLogits.length === logitsVocabSize
+          ? fallbackLogits
+          : extractLastPositionLogits(fallbackLogits, numTokens, logitsVocabSize);
+      }
+
+      releaseBuffer(currentHiddenBuffer);
+    } else {
+      if (currentRecorder) {
+        await currentRecorder.submitAndWait();
+        await recordProfile(currentRecorder);
+      }
+      const logits = await computeLogits(
+        currentHiddenBuffer,
+        numTokens,
+        getLogitsWeights(this.#state),
+        getLogitsConfig(this.#state),
+        this.#state.useGPU,
+        this.#state.debugFlags,
+        undefined,
+        debugCheckBuffer,
+        this.#state.runtimeConfig.shared.debug.probes,
+        { lastPositionOnly: true }
+      );
+
+      lastLogits = logits.length === logitsVocabSize
+        ? logits
+        : extractLastPositionLogits(logits, numTokens, logitsVocabSize);
+      releaseBuffer(currentHiddenBuffer);
+    }
+
+    this.#state.currentSeqLen = startPos + numTokens;
+
+    if (usedRecordedLogits) {
+      if (logitsVocabSize < config.vocabSize) {
+        const padded = new Float32Array(config.vocabSize);
+        padded.set(lastLogits);
+        padded.fill(-Infinity, logitsVocabSize);
+        lastLogits = padded;
+      }
+      if (config.finalLogitSoftcapping != null) {
+        applySoftcapping(lastLogits, config.finalLogitSoftcapping);
+      }
+    }
+
+    if (opts.debug) {
+      logitsSanity(lastLogits, 'Prefill', (tokens) => resolveTokenText(this.#state.tokenizer, tokens));
+    }
+
+    if (opts.debug) {
+      if (this.#state.kvCache?.hasGPUCache?.()) {
+        log.debug('Pipeline', `KV cache active after prefill: seqLen=${this.#state.kvCache.getKeyCache(0)?.constructor.name ?? '?'}`);
+      } else {
+        log.warn('Pipeline', `KV cache NOT active after prefill! hasGPUCache=${this.#state.kvCache?.hasGPUCache?.()}`);
+      }
+    }
+
+    return lastLogits;
+  }
+
+
+  async _decodeStep(currentIds, opts) {
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+    return decodeStep(this.#state, currentIds, opts, this._getDecodeHelpers(debugCheckBuffer));
+  }
+
+  async decodeStepLogits(currentIds, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    if (this.#state.isGenerating) throw new Error('Generation already in progress');
+    resetActiveExecutionPlan(this.#state);
+
+    validateCallTimeOptions(options);
+
+    const opts = this._resolveStepOptions(options);
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+
+    return decodeStepLogits(this.#state, currentIds, opts, this._getDecodeHelpers(debugCheckBuffer));
+  }
+
+  async advanceWithToken(tokenId, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    if (this.#state.isGenerating) throw new Error('Generation already in progress');
+    resetActiveExecutionPlan(this.#state);
+
+    validateCallTimeOptions(options);
+
+    const opts = this._resolveStepOptions(options);
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+
+    this._assertTokenIdInRange(tokenId, 'advanceWithToken');
+    await advanceWithToken(this.#state, tokenId, opts, this._getDecodeHelpers(debugCheckBuffer));
+  }
+
+  async advanceWithTokenAndEmbedding(tokenId, options = {}) {
+    if (!this.#state.isLoaded) throw new Error('Model not loaded');
+    if (this.#state.isGenerating) throw new Error('Generation already in progress');
+    resetActiveExecutionPlan(this.#state);
+
+    validateCallTimeOptions(options);
+
+    const opts = this._resolveStepOptions(options);
+    const embeddingMode = resolveAdvanceEmbeddingMode(this.#state, options);
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+
+    this._assertTokenIdInRange(tokenId, 'advanceWithTokenAndEmbedding');
+    return runAdvanceWithTokenAndEmbedding(
+      this.#state,
+      tokenId,
+      opts,
+      this._getDecodeHelpers(debugCheckBuffer),
+      embeddingMode
+    );
+  }
+
+  async _generateNTokensGPU(startToken, N, currentIds, opts) {
+    const debugCheckBuffer = this.#state.debug
+      ? (buffer, label, numTokens, expectedDim) =>
+        debugCheckBufferHelper(this.#state, buffer, label, numTokens, expectedDim)
+      : undefined;
+    return generateNTokensGPU(this.#state, startToken, N, currentIds, opts, this._getDecodeHelpers(debugCheckBuffer));
+  }
+}