Ensure num_logits_to_keep is always set to 1 for generate calls (#1681)

* support scalar tensor tolist()

* gemma support num_logits_to_keep

* set num_logits_to_keep in phi3_v and chatterbox too

* set num_logits_to_keep in qwen2_vl and multi_modality

* allow num_logits_to_keep to be used in encoder-decoder models

* Deduplicate set logic

Author: xenova

packages/transformers/src/models/chatterbox/modeling_chatterbox.js

@@ -49,6 +49,7 @@ export class ChatterboxModel extends ChatterboxPreTrainedModel {
         // Generic generation parameters
         generation_config = null,
         logits_processor = null,
+        num_logits_to_keep = null,
 
         // Speaker embeddings/features (useful for re-using pre-computed speaker data)
         audio_features = null, // float32[batch_size,sequence_length,1024]
@@ -118,6 +119,7 @@ export class ChatterboxModel extends ChatterboxPreTrainedModel {
                 attention_mask,
                 generation_config,
                 logits_processor,
+                num_logits_to_keep,
             },
             false,
         );

packages/transformers/src/models/florence2/modeling_florence2.js

@@ -79,6 +79,7 @@ export class Florence2ForConditionalGeneration extends Florence2PreTrainedModel
 
         inputs_embeds,
         decoder_inputs_embeds,
+        num_logits_to_keep = null,
     }) {
         if (!inputs_embeds) {
             ({ inputs_embeds, attention_mask } = await this._prepare_inputs_embeds({
@@ -108,6 +109,7 @@ export class Florence2ForConditionalGeneration extends Florence2PreTrainedModel
             encoder_attention_mask: attention_mask,
             encoder_hidden_states: encoder_outputs,
             past_key_values,
+            num_logits_to_keep,
         };
         return await decoder_forward(this, decoderFeeds, true);
     }

packages/transformers/src/models/gemma3n/modeling_gemma3n.js

@@ -38,6 +38,7 @@ export class Gemma3nForConditionalGeneration extends Gemma3nPreTrainedModel {
         // Generic generation parameters
         generation_config = null,
         logits_processor = null,
+        num_logits_to_keep = null,
 
         // TODO: needed?
         ...kwargs
@@ -84,6 +85,7 @@ export class Gemma3nForConditionalGeneration extends Gemma3nPreTrainedModel {
                 position_ids,
                 generation_config,
                 logits_processor,
+                num_logits_to_keep,
             },
             true,
         );

packages/transformers/src/models/modeling_utils.js

@@ -1294,6 +1294,31 @@ export function addPastKeyValues(self, decoderFeeds, pastKeyValues) {
     return new DynamicCache(entries);
 }
 
+/**
+ * Sets `num_logits_to_keep` on `model_inputs` if the decoder session declares it as an input
+ * and it has not already been set.
+ *
+ * `num_logits_to_keep` specifies how many trailing prompt logits the model computes:
+ * - `0n` (or unset) computes logits for the entire sequence — used for prefill/scoring.
+ * - `1n` computes only the last token's logits — used during autoregressive generation,
+ *   since only the last prompt token's logits are needed to sample the next token. For long
+ *   sequences, computing all logits uses a lot of memory, so `1n` significantly reduces the
+ *   memory footprint.
+ * - Any other positive integer keeps the last `num_logits_to_keep` logits.
+ *
+ * @param {PreTrainedModel} self The model instance.
+ * @param {Record<string, any>} model_inputs The model inputs to mutate.
+ * @param {bigint} value The value to set (typically `1n` for generation, `0n` as a fallback).
+ * @private
+ */
+export function setNumLogitsToKeep(self, model_inputs, value) {
+    if (model_inputs.num_logits_to_keep) return;
+    const session = self.sessions['decoder_model_merged'] ?? self.sessions['model'];
+    if (session?.inputNames.includes('num_logits_to_keep')) {
+        model_inputs.num_logits_to_keep = new Tensor('int64', [value], []);
+    }
+}
+
 /**
  * Forward pass of a decoder model.
  * @param {Object} self The decoder model.
@@ -1321,14 +1346,8 @@ export async function decoder_forward(self, model_inputs, is_encoder_decoder = f
         new_model_inputs.position_ids = create_position_ids(new_model_inputs, past_key_values, start_index);
     }
 
-    if (session.inputNames.includes('num_logits_to_keep') && !new_model_inputs.num_logits_to_keep) {
-        // `num_logits_to_keep` specifies the number of prompt logits to calculate during generation.
-        // If unset (or 0), all logits will be calculated. If an integer value, only last `num_logits_to_keep`
-        // logits will be calculated. During generation, the default is 1 because only the logits of the last
-        // prompt token are needed for generation. For long sequences, the logits for the entire sequence may
-        // use a lot of memory so, setting `num_logits_to_keep=1` will reduce memory footprint significantly.
-        new_model_inputs.num_logits_to_keep = new Tensor('int64', [0n], []);
-    }
+    // Fallback for non-generation forward calls (e.g. prefill scoring): compute all logits.
+    setNumLogitsToKeep(self, new_model_inputs, 0n);
 
     // Unpack the `past_key_values` object into model inputs
     addPastKeyValues(self, new_model_inputs, past_key_values);
@@ -1353,6 +1372,7 @@ export async function decoder_forward(self, model_inputs, is_encoder_decoder = f
  * @param {DynamicCache} [params.past_key_values=null]
  * @param {Object} [params.generation_config=null]
  * @param {Object} [params.logits_processor=null]
+ * @param {Tensor} [params.num_logits_to_keep=null]
  * @returns {Promise<Tensor>} The model's output tensor
  * @private
  */
@@ -1377,6 +1397,7 @@ export async function generic_text_to_text_forward(
         // Generic generation parameters
         generation_config = null,
         logits_processor = null,
+        num_logits_to_keep = null,
 
         // Additional parameters
         ...kwargs
@@ -1454,6 +1475,7 @@ export async function generic_text_to_text_forward(
             position_ids,
             generation_config,
             logits_processor,
+            num_logits_to_keep,
         },
         true,
     );
@@ -1552,12 +1574,7 @@ export function create_position_ids(model_inputs, past_key_values = null, start_
 export function decoder_prepare_inputs_for_generation(self, input_ids, model_inputs, generation_config) {
     const past_length = model_inputs.past_key_values ? model_inputs.past_key_values.get_seq_length() : 0;
 
-    // During generation, only the last token's logits are needed. Setting num_logits_to_keep=1
-    // avoids computing logits for the entire sequence, significantly reducing memory usage.
-    const session = self.sessions['decoder_model_merged'] ?? self.sessions['model'];
-    if (session?.inputNames.includes('num_logits_to_keep') && !model_inputs.num_logits_to_keep) {
-        model_inputs.num_logits_to_keep = new Tensor('int64', [1n], []);
-    }
+    setNumLogitsToKeep(self, model_inputs, 1n);
 
     if (!model_inputs.attention_mask) {
         // If the attention mask is not provided, we attempt to infer based on provided inputs
@@ -1606,6 +1623,8 @@ export function encoder_decoder_prepare_inputs_for_generation(self, input_ids, m
         input_ids = input_ids.map((x) => [x.at(-1)]);
     }
 
+    setNumLogitsToKeep(self, model_inputs, 1n);
+
     return {
         ...model_inputs,
         decoder_input_ids: toI64Tensor(input_ids),

packages/transformers/src/models/multi_modality/modeling_multi_modality.js

@@ -1,4 +1,4 @@
-import { PreTrainedModel, decoder_forward } from '../modeling_utils.js';
+import { PreTrainedModel, decoder_forward, setNumLogitsToKeep } from '../modeling_utils.js';
 import { sessionRun } from '../session.js';
 import { pick } from '../../utils/core.js';
 import { RawImage } from '../../utils/image.js';
@@ -73,6 +73,8 @@ export class MultiModalityCausalLM extends MultiModalityPreTrainedModel {
     prepare_inputs_for_generation(input_ids, model_inputs, generation_config) {
         const has_past_key_values = !!model_inputs.past_key_values;
 
+        setNumLogitsToKeep(this, model_inputs, 1n);
+
         if (generation_config.guidance_scale !== null && generation_config.guidance_scale > 1) {
             if (has_past_key_values) {
                 model_inputs.input_ids = cat([model_inputs.input_ids, model_inputs.input_ids], 0);

packages/transformers/src/models/phi3_v/modeling_phi3_v.js

@@ -29,6 +29,7 @@ export class Phi3VForCausalLM extends Phi3VPreTrainedModel {
         // Generic generation parameters
         generation_config = null,
         logits_processor = null,
+        num_logits_to_keep = null,
 
         // TODO: needed?
         ...kwargs
@@ -65,6 +66,7 @@ export class Phi3VForCausalLM extends Phi3VPreTrainedModel {
                 position_ids,
                 generation_config,
                 logits_processor,
+                num_logits_to_keep,
             },
             false,
         );

packages/transformers/src/models/qwen2_vl/modeling_qwen2_vl.js

@@ -1,4 +1,9 @@
-import { PreTrainedModel, cumsum_masked_fill, default_merge_input_ids_with_image_features } from '../modeling_utils.js';
+import {
+    PreTrainedModel,
+    cumsum_masked_fill,
+    default_merge_input_ids_with_image_features,
+    setNumLogitsToKeep,
+} from '../modeling_utils.js';
 import { sessionRun } from '../session.js';
 import { stack, Tensor, ones_like, zeros } from '../../utils/tensor.js';
 import { max } from '../../utils/maths.js';
@@ -285,6 +290,8 @@ export class Qwen2VLForConditionalGeneration extends Qwen2VLPreTrainedModel {
     }
 
     prepare_inputs_for_generation(input_ids, model_inputs, generation_config) {
+        setNumLogitsToKeep(this, model_inputs, 1n);
+
         // Overwritten -- in specific circumstances we don't want to forward image inputs to the model
         if (!model_inputs.attention_mask || model_inputs.position_ids) {
             return model_inputs;

packages/transformers/src/utils/tensor.js

@@ -973,7 +973,7 @@ export class Tensor {
  */
 function reshape(data, dimensions) {
     const totalElements = data.length;
-    const dimensionSize = dimensions.reduce((a, b) => a * b);
+    const dimensionSize = dimensions.length === 0 ? 1 : dimensions.reduce((a, b) => a * b);
 
     if (totalElements !== dimensionSize) {
         throw Error(`cannot reshape array of size ${totalElements} into shape (${dimensions})`);

packages/transformers/tests/utils/tensor.test.js

@@ -369,6 +369,11 @@ describe("Tensor operations", () => {
   });
 
   describe("tolist", () => {
+    it("should work with scalar tensors", () => {
+      const t1 = new Tensor("float32", [42], []);
+      const arr = t1.tolist();
+      expect(arr).toBe(42);
+    });
     it("should return nested arrays for a 2D tensor", () => {
       const t1 = new Tensor("float32", [1, 2, 3, 4], [2, 2]);
       const arr = t1.tolist();