[generation] Fix per_batch_line sampling param lookup for CB

GitOrigin-RevId: 5039ec6b067d1431201563659bb05b9c90e79bb0

Author: aws-patlange

src/transformers_neuronx/decoder.py

@@ -698,7 +698,6 @@ def _hlo_eagle_draft_unroll(self, hidden, tensors, layers_caches, layers_weights
         return logits, hidden, out_caches
 
     def _hlo_fully_unrolled(self, n_positions, batch_size):
-
         self.builder.n_positions = n_positions
         if self.neuron_config.optimized_paged_attention and self.n_active_tokens == 1:
             self.builder.num_active_blocks = batch_size
@@ -733,7 +732,7 @@ def fully_unrolled(scribe):
             else:
                 logits, out_caches = self._hlo_unroll(hidden, tensors, in_caches, layers_weights, pre_layer_params, lm_head_params)
             self._hlo_cache_aliases(in_caches, out_caches)
-            output = self._hlo_generation(logits, generation_params)
+            output = self._hlo_generation(logits, generation_params, start_ids=tensors[1])
 
             # Set the output
             out_caches = itertools.chain(*out_caches)
@@ -1006,11 +1005,12 @@ def ln_lm_head(scribe):
                 next_tok_id = scribe.s32[batch_size].Parameter(parameter_number=1)
             else:
                 next_tok_id = scribe.s32[1].Parameter(parameter_number=1)
-            param_builder = DecoderParameterBuilder(scribe, 2)
+            start_ids = scribe.s32[batch_size].Parameter(parameter_number=2)
+            param_builder = DecoderParameterBuilder(scribe, 3)
             ln_f_weight, ln_f_bias, head_weight, head_bias = self._hlo_lm_head_params(param_builder)
-            gneration_params = self._hlo_generation_params(param_builder)
+            generation_params = self._hlo_generation_params(param_builder)
             logits = self.ln_lm_head_builder(hidden, next_tok_id, ln_f_weight, ln_f_bias, head_weight, head_bias, return_all_outputs=self.return_all_outputs)
-            output = self._hlo_generation(logits, gneration_params)
+            output = self._hlo_generation(logits, generation_params, start_ids=start_ids)
             if self.neuron_config.log_softmax_scores:
                 logits, scores = self._hlo_post_layer(logits)
                 outputs = [logits, scores]
@@ -1034,7 +1034,7 @@ def _hlo_generation_params(self, param_builder):
                 params.append(param)
         return params
 
-    def _hlo_generation(self, logits, params, early_return=False, return_probs=False):
+    def _hlo_generation(self, logits, params, early_return=False, return_probs=False, start_ids=None):
         generation_config = self.neuron_config.on_device_generation
         if generation_config is None:
             return logits
@@ -1045,13 +1045,16 @@ def _hlo_generation(self, logits, params, early_return=False, return_probs=False
             self.neuron_config.on_device_generation.top_p = top_p
             self.neuron_config.on_device_generation.temperature = temperature
             self.neuron_config.on_device_generation.top_p_min_tokens = top_p_min_tokens
+        
+        seq_ids = start_ids if self.neuron_config.continuous_batching is not None else None
         return generation.generate(
             logits,
             logits_indices,
             config=generation_config,
             tp_degree=self.tp_degree,
             early_return=early_return,
             return_probs=return_probs,
+            seq_ids=seq_ids,
         )
 
     # Mainly used for serialization purposes.
@@ -2458,10 +2461,12 @@ def setup(self, layers, pre_layer_params, ln_lm_head_params):
 
         hidden_buffers = list()
         last_token_id_buffers = list()
+        start_ids_buffers = list()
         for input_buffer in self.input_buffers:
-            hidden_buffer, _, _, last_token_id_buffer, *_ = input_buffer
+            hidden_buffer, _, start_ids_buffer, last_token_id_buffer, *_ = input_buffer
             hidden_buffers.append(hidden_buffer)
             last_token_id_buffers.append(last_token_id_buffer)
+            start_ids_buffers.append(start_ids_buffer)
 
         multi_layer_starts = range(0, len(layers), self.unroll)
         multi_layers = [layers[start:start+self.unroll] for start in multi_layer_starts]
@@ -2480,7 +2485,7 @@ def setup(self, layers, pre_layer_params, ln_lm_head_params):
         if self.neuron_config.is_valid_lm_head():
             for head_idx in range(0,len(self.ln_lm_head_kernels)):
                 output_tensors = [*self.logits_buffer[head_idx]] if self.neuron_config.log_softmax_scores or self.neuron_config.is_eagle_target else [self.logits_buffer[head_idx]]
-                self.ln_lm_head_memories[head_idx].setup([hidden_buffers[head_idx], last_token_id_buffers[head_idx], *ln_lm_head_params], output_tensors)
+                self.ln_lm_head_memories[head_idx].setup([hidden_buffers[head_idx], last_token_id_buffers[head_idx], start_ids_buffers[head_idx], *ln_lm_head_params], output_tensors)
                 self.ln_lm_head_kernels[head_idx].build()
                 self.ln_lm_head_kernels[head_idx].load()
 

src/transformers_neuronx/layers/generation.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from transformers_neuronx import hlo, config
 
-def generate(logits, logits_indices, config: config.GenerationConfig, tp_degree=1, early_return=False, return_probs=False):
+def generate(logits, logits_indices, config: config.GenerationConfig, tp_degree=1, early_return=False, return_probs=False, seq_ids=None):
     logits = mask_logits(logits, logits_indices, config.vocab_size)
     if not config.dynamic and not config.do_sample:
         tokens = greedy_search(logits, tp_degree=tp_degree)
@@ -62,10 +62,9 @@ def generate(logits, logits_indices, config: config.GenerationConfig, tp_degree=
         logits_slice = hlo.slice_along(logits, 0, start=batch_line, limit=batch_line+1)
         indices_slice = None if indices is None else hlo.slice_along(indices, 0, start=batch_line, limit=batch_line+1)
 
-        batch_line_top_k = config.top_k if hlo._is_hlo_scalar(config.top_k) else hlo.get_hlo_scalar_by_index(config.top_k, batch_line)
-        batch_line_top_p = config.top_p if hlo._is_hlo_scalar(config.top_p) else hlo.get_hlo_scalar_by_index(config.top_p, batch_line)
-        batch_line_temperature = config.temperature if hlo._is_hlo_scalar(config.temperature) else hlo.get_hlo_scalar_by_index(config.temperature, batch_line)
-        batch_line_top_p_min_tokens =  config.top_p_min_tokens if hlo._is_hlo_scalar(config.top_p_min_tokens) else  hlo.get_hlo_scalar_by_index(config.top_p_min_tokens, batch_line)
+        batch_line_top_k, batch_line_top_p, batch_line_temperature, batch_line_top_p_min_tokens = sampling_params_for_batch_line(
+            seq_ids, batch_line, config
+        )
 
         token = sample(
             logits_slice,
@@ -104,6 +103,22 @@ def generate(logits, logits_indices, config: config.GenerationConfig, tp_degree=
         returned_tokens = hlo.concatenate(tokens, dimension=0)
         return returned_tokens
 
+
+def sampling_params_for_batch_line(seq_ids, batch_line: int, config: config.GenerationConfig):
+    if seq_ids is not None:
+        seq_id_for_batch = hlo.slice_along(seq_ids, 0, start=batch_line, limit=batch_line+1)
+        batch_line_top_k = hlo.reshape(hlo.index_select(config.top_k, 0, seq_id_for_batch), [])
+        batch_line_top_p = hlo.reshape(hlo.index_select(config.top_p, 0, seq_id_for_batch), [])
+        batch_line_temperature = hlo.reshape(hlo.index_select(config.temperature, 0, seq_id_for_batch), [])
+        batch_line_top_p_min_tokens = hlo.reshape(hlo.index_select(config.top_p_min_tokens, 0, seq_id_for_batch), [])
+    else:
+        batch_line_top_k = config.top_k if hlo._is_hlo_scalar(config.top_k) else hlo.get_hlo_scalar_by_index(config.top_k, batch_line)     
+        batch_line_top_p = config.top_p if hlo._is_hlo_scalar(config.top_p) else hlo.get_hlo_scalar_by_index(config.top_p, batch_line)
+        batch_line_temperature = config.temperature if hlo._is_hlo_scalar(config.temperature) else hlo.get_hlo_scalar_by_index(config.temperature, batch_line)
+        batch_line_top_p_min_tokens =  config.top_p_min_tokens if hlo._is_hlo_scalar(config.top_p_min_tokens) else  hlo.get_hlo_scalar_by_index(config.top_p_min_tokens, batch_line)
+    return (batch_line_top_k, batch_line_top_p, batch_line_temperature, batch_line_top_p_min_tokens)
+
+
 def mask_logits(logits, indices, model_vocab_size):
     vocab_size, n_active_tokens, _ = logits.sizes
     indices_br = hlo.broadcast(indices, (logits.sizes), broadcast_dimensions=(0,))