@simulatte/doppler 0.1.8 → 0.1.9

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

@@ -1,6 +1,6 @@
 import { getDevice, setTrackSubmits } from '../../../gpu/device.js';
 import { releaseBuffer, readBuffer } from '../../../memory/buffer-pool.js';
-import { runArgmax, runGPUSample, recordArgmax, recordGPUSample, isGPUSamplingAvailable } from '../../../gpu/kernels/sample.js';
+import { recordArgmax, recordGPUSample, isGPUSamplingAvailable } from '../../../gpu/kernels/sample.js';
 import { recordCheckStop } from '../../../gpu/kernels/check-stop.js';
 import { resetSubmitStats, logSubmitStats } from '../../../gpu/submit-tracker.js';
 import { createCommandRecorder, createProfilingRecorder, CommandRecorder } from '../../../gpu/command-recorder.js';
@@ -20,6 +20,7 @@ import { decodeReadback } from './debug-utils/index.js';
 import { getFinalNormWeights, extractEmbeddingFromHidden } from './generator-runtime.js';
 import { parseFinitenessStatusWords } from './finiteness-guard-status.js';
 import { hasLinearAttentionLayers } from './linear-attention.js';
+import { hasConvLayers } from './layer.js';
 
 const UNKNOWN_TOKEN_TEXT = '<unknown>';
 
@@ -91,6 +92,13 @@ export function shouldUseBatchDecode(config) {
   return isBatchDecodeEnabled(config);
 }
 
+export function shouldUseFusedDecodeSampling(config) {
+  return config.recorderEnabled === true
+    && config.gpuSamplingEnabled === true
+    && config.fusedDecodeDisabled !== true
+    && !hasConvLayers(config.layerTypes ?? []);
+}
+
 export function resolveBatchStop(tokens, stopFlags, stopTokenIds, eosTokenId) {
   let actualCount = tokens.length;
   if (stopFlags) {
@@ -403,7 +411,12 @@ export async function decodeStep(state, currentIds, opts, helpers) {
   const padTokenId = state.tokenizer?.getSpecialTokens?.()?.pad ?? null;
   const lmHeadIsCpu = isCpuWeightBuffer(state.weights.get('lm_head'));
   const useGPUSampling = state.useGPU && isGPUSamplingAvailable() && !lmHeadIsCpu;
-  const useFusedDecode = recorder && useGPUSampling && !state.disableFusedDecode;
+  const useFusedDecode = shouldUseFusedDecodeSampling({
+    recorderEnabled: Boolean(recorder),
+    gpuSamplingEnabled: useGPUSampling,
+    fusedDecodeDisabled: state.disableFusedDecode,
+    layerTypes: config.layerTypes,
+  });
 
   if (useFusedDecode) {
     const ring = state.decodeRing;
@@ -631,36 +644,35 @@ export async function decodeStep(state, currentIds, opts, helpers) {
     );
     if (logitsResult) {
       const { logitsBuffer, vocabSize, logitsDtype } = logitsResult;
+      const logitsBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: logitsDtype });
+      const logitsData = await readBuffer(logitsBuffer, numTokens * vocabSize * logitsBytes);
+      releaseBuffer(logitsBuffer);
 
-      const nextToken = opts.temperature < samplingDefaults.greedyThreshold
-        ? await runArgmax(logitsBuffer, vocabSize, { padTokenId, logitSoftcap, logitsDtype, outputIndex: 0 })
-        : await runGPUSample(logitsBuffer, vocabSize, {
-          temperature: opts.temperature,
-          topK: opts.topK,
-          padTokenId,
-          logitSoftcap,
-          logitsDtype,
-          outputIndex: 0,
-          greedyThreshold: samplingDefaults.greedyThreshold,
-          randomSeed: opts.seed,
-        });
+      const rawLogits = decodeReadback(logitsData, logitsDtype);
+      const finalizedLogits = await finalizeLogits(
+        rawLogits,
+        numTokens,
+        vocabSize,
+        config.vocabSize,
+        config,
+        state.runtimeConfig.shared.debug.probes
+      );
+      const sampledLogits = extractLastPositionLogits(finalizedLogits, numTokens, config.vocabSize);
 
-      releaseBuffer(logitsBuffer);
-      const invalidGpuToken = nextToken >= config.vocabSize
-        || (padTokenId != null && nextToken === padTokenId)
-        || (padTokenId == null && nextToken === 0);
-      if (!invalidGpuToken) {
-        if (!context.decodeBuffers?.ownsBuffer(hiddenStates)) {
-          releaseBuffer(hiddenStates);
-        }
-        state.currentSeqLen++;
-        return nextToken;
+      applyRepetitionPenalty(sampledLogits, currentIds, opts.repetitionPenalty);
+      const nextToken = sample(sampledLogits, {
+        temperature: opts.temperature,
+        topP: opts.topP,
+        topK: opts.topK,
+        padTokenId,
+        seed: opts.seed,
+      });
+
+      if (!context.decodeBuffers?.ownsBuffer(hiddenStates)) {
+        releaseBuffer(hiddenStates);
       }
-      state.disableFusedDecode = true;
-      log.warn(
-        'Decode',
-        `GPU sampling produced invalid token ${nextToken} (vocabSize=${config.vocabSize}, step=${state.decodeStepCount}); falling back to CPU sampling.`
-      );
+      state.currentSeqLen++;
+      return nextToken;
     }
   }
 
@@ -887,6 +899,11 @@ export async function generateNTokensGPU(state, startToken, N, currentIds, opts,
       '[Pipeline] Batch decode path is disabled for linear_attention models; use single-token decode.'
     );
   }
+  if (hasConvLayers(config.layerTypes)) {
+    throw new Error(
+      '[Pipeline] Batch decode path is disabled for conv models; use single-token decode.'
+    );
+  }
   const samplingDefaults = state.runtimeConfig.inference.sampling;
   const executionPlan = opts.executionPlan;
   const batchSize = executionPlan?.batchSize ?? opts.batchSize ?? state.runtimeConfig.inference.batching.batchSize;

package/src/inference/pipelines/text/generator.d.ts

@@ -27,6 +27,11 @@ export declare class PipelineGenerator {
    * Batching and readback cadence are controlled by runtime.inference.batching.
    */
   generate(prompt: PromptInput, options?: GenerateOptions): AsyncGenerator<string, void, void>;
+  generateTokens(prompt: PromptInput, options?: GenerateOptions): AsyncGenerator<number, void, void>;
+  generateTokenIds(
+    prompt: PromptInput,
+    options?: GenerateOptions
+  ): Promise<{ tokenIds: number[]; stats: import('./types.js').PipelineStats }>;
   prefillKVOnly(prompt: PromptInput, options?: GenerateOptions): Promise<KVCacheSnapshot>;
   prefillWithEmbedding(prompt: PromptInput, options?: GenerateOptions): Promise<PrefillEmbeddingResult>;
   prefillWithLogits(prompt: PromptInput, options?: GenerateOptions): Promise<PrefillResult>;