Remove torch.compile

Signed-off-by: ziyixiong-nv <[email protected]>

Author: ziyixiong-nv

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1474,7 +1474,6 @@ def _prepare_incremental_update_metadata(
 
         return lora_params
 
-    @torch.compile(options={"max-autotune": True})
     def _update_draft_input_tensors(self,
                                     num_accepted_tokens_device: torch.Tensor,
                                     new_tokens_device: torch.Tensor,
@@ -1599,7 +1598,6 @@ def _apply_incremental_update_draft(
 
         return inputs, self.gather_ids_cuda[:num_generation_tokens]
 
-    @torch.compile(options={"max-autotune": True})
     def _update_target_input_tensors(
             self, num_accepted_tokens_device: torch.Tensor,
             new_tokens_device: torch.Tensor,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1687,7 +1687,6 @@ def _executor_loop_overlap(self):
                 self.iter_counter += 1
 
     @nvtx_range("_accept_draft_tokens")
-    @torch.compile(options={"max-autotune": True})
     def _accept_draft_tokens(
         self, scheduled_batch: ScheduledRequests,
         target_outputs: SampleStateTensors,

tensorrt_llm/_torch/speculative/drafting_loops.py

@@ -120,24 +120,27 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
 
         new_draft_tokens = [self.sample(logits)]
         draft_logits = [logits]
-        with save_metadata_state(attn_metadata, spec_metadata):
-            batch_size = attn_metadata.num_seqs
-
-            new_position_ids = self.prepare_for_generation(
-                attn_metadata, spec_metadata, position_ids)
-            for i in range(self.max_draft_len - 1):
-                logits = self.draft_model.forward(
-                    input_ids=new_draft_tokens[-1],
-                    position_ids=new_position_ids,
-                    attn_metadata=attn_metadata,
-                    spec_metadata=spec_metadata)
-                new_draft_tokens.append(self.sample(logits))
-                draft_logits.append(logits)
-                new_position_ids += 1
-                attn_metadata.kv_lens_cuda[:batch_size] += 1
-                if i == 0 and isinstance(spec_metadata, Eagle3SpecMetadata):
-                    spec_metadata.hidden_states_read_indices[:batch_size].copy_(
-                        spec_metadata.hidden_states_write_indices[:batch_size])
+        if self.max_draft_len > 1:
+            is_eagle3 = isinstance(spec_metadata, Eagle3SpecMetadata)
+            with save_metadata_state(attn_metadata, spec_metadata):
+                batch_size = attn_metadata.num_seqs
+
+                new_position_ids = self.prepare_for_generation(
+                    attn_metadata, spec_metadata, position_ids)
+                for i in range(self.max_draft_len - 1):
+                    logits = self.draft_model.forward(
+                        input_ids=new_draft_tokens[-1],
+                        position_ids=new_position_ids,
+                        attn_metadata=attn_metadata,
+                        spec_metadata=spec_metadata)
+                    new_draft_tokens.append(self.sample(logits))
+                    draft_logits.append(logits)
+                    new_position_ids += 1
+                    attn_metadata.kv_lens_cuda[:batch_size] += 1
+                    if i == 0 and is_eagle3:
+                        spec_metadata.hidden_states_read_indices[:batch_size].copy_(
+                            spec_metadata.
+                            hidden_states_write_indices[:batch_size])
 
         return {
             "new_draft_tokens": torch.stack(new_draft_tokens),
@@ -153,7 +156,6 @@ def sample(self, logits: torch.Tensor) -> torch.Tensor:
 
         return tokens
 
-    @torch.compile(options={'max-autotune': True})
     def prepare_for_generation(self, attn_metadata: AttentionMetadata,
                                spec_metadata: SpecMetadata,
                                position_ids: torch.Tensor) -> torch.Tensor: