Enable compilation

text-generation-inference/server/text_generation_server/generator.py

@@ -307,12 +307,6 @@ def __init__(
         tokenizer: PreTrainedTokenizerBase,
     ):
         self.model = model
-        if model.device.type == "xla" and "DBG_COMPILE" in os.environ:
-            self.model_one_token = torch.compile(model, backend="openxla")
-            logger.debug("Model compiled for decoding")
-        else:
-            self.model_one_token = model
-
         # Specify padding options for decoder-only architecture
         tokenizer.pad_token_id = tokenizer.eos_token_id
         tokenizer.padding_side = "left"
@@ -328,6 +322,12 @@ def __init__(
                 f"Static cache not available for {self.model.__class__.__name__}. Performance will be affected"
             )
             self.use_static_cache = False
+        # compile model when possible to accelerate decoding
+        if model.device.type == "xla" and ("DBG_COMPILE" in os.environ or self.use_static_cache):
+            self.model_one_token = torch.compile(model, backend="openxla")
+            logger.debug("Model compiled for decoding")
+        else:
+            self.model_one_token = model
 
     @property
     def info(self) -> InfoResponse:
@@ -425,7 +425,6 @@ def prefill(self, batch: Batch) -> Tuple[List[Generation], CachedBatch]:
         self.past_key_values = None
         # Obtain position ids using attention mask.
         position_ids = (attention_mask.cumsum(-1) - 1).masked_fill(attention_mask == 0, 0)
-        position_ids = position_ids[:, -input_ids.shape[-1] :]
 
         extra_args = {}
         if self.use_static_cache:
@@ -483,7 +482,7 @@ def decode(self, batches: List[CachedBatch]) -> Tuple[List[Generation], CachedBa
                         device=self.model.device,
                     )
                 # input_ids are simply the tokens generated by the last decode or prefill requests (other tokens are cached)
-                input_ids[i, 0] = slot.next_token
+                input_ids.index_put_([torch.tensor([i])], slot.next_token)
                 if not self.use_static_cache:
                     # When using dynamic cache, the whole attention mask needs to be passed over to the model at each iteration.
                     if attention_mask is None:
@@ -493,8 +492,8 @@ def decode(self, batches: List[CachedBatch]) -> Tuple[List[Generation], CachedBa
                             dtype=torch.int64,
                             device=self.model.device,
                         )
-                    attention_mask[i, :] = slot.attention_mask
-                position_ids[i, 0] = slot.cur_position
+                    attention_mask.index_put_([torch.tensor([i])], slot.attention_mask)
+                position_ids.index_put_([torch.tensor([i])], torch.tensor(slot.cur_position))
         if input_ids is None:
             raise ValueError("Unable to decode tokens for non-prefilled batches (probably due to a previous failure)")
         extra_args = {}