[TRTLLM-4983] feat: enable overlap scheduler between draft forwards (NVIDIA#4802)

Signed-off-by: Fanrong Li <[email protected]>

Author: lfr-0531

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -247,6 +247,7 @@ def forward(
         hidden_states: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
+        hidden_states = self.apply_eagle3_fc(spec_metadata.get_hidden_states())
         output, _ = self.model(
             input_ids=input_ids,
             attn_metadata=attn_metadata,

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -57,7 +57,6 @@ def __init__(
                                         device=device,
                                         dtype=torch.int32)
 
-        self.extra_model_inputs = {}
         self.attn_metadata = attn_metadata
         self.spec_metadata = spec_metadata
         self._output = None
@@ -70,22 +69,7 @@ def capture(
         self,
         forward_fn: Callable[[Dict[str, Any]], torch.Tensor],
         pool: Optional[Tuple[int, int]] = None,
-        extra_model_inputs: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Tuple[int, int]:
-        """
-        Captures a CUDA graph by calling forward_fn(inputs),
-        where inputs is extra_model_inputs + this graph runner's
-        input_ids, position_ids, spec_metadata and attn_metadata.
-
-        Extra model inputs have the following semantics if
-        the extra input is a tensor (or collection of
-        tensors). The CUDA graph runner will create a buffer
-        of the same shape/dtype/device, and subsequent calls to run() will
-        require this extra model input. Input tensors will be
-        copied into the buffer that this CUDA graph runner owns.
-        This implies that these buffers *must* have static shapes for
-        this CUDA graph's batch size.
-        """
         self._graph = torch.cuda.CUDAGraph()
         inputs = {
             "attn_metadata": self.attn_metadata,
@@ -94,11 +78,6 @@ def capture(
             "inputs_embeds": None,
             "spec_metadata": self.spec_metadata,
         }
-        if extra_model_inputs is not None:
-            for key, tensor in extra_model_inputs.items():
-                new_tensor = tensor.clone()
-                inputs[key] = new_tensor
-                self.extra_model_inputs[key] = new_tensor
 
         # We have to do warm up runs to initialize PyTorch's
         # internal states according to the docs:
@@ -119,11 +98,7 @@ def capture(
     def needs_capture(self) -> bool:
         return self._output is None
 
-    def run(
-        self,
-        inputs: Dict[str, Any],
-        extra_model_inputs: Optional[Dict[str, torch.Tensor]] = None
-    ) -> torch.Tensor:
+    def run(self, inputs: Dict[str, Any]) -> torch.Tensor:
         assert "input_ids" in inputs
         assert "position_ids" in inputs
         assert "attn_metadata" in inputs
@@ -145,13 +120,6 @@ def run(
         self.input_ids[:seqlen].copy_(input_ids)
         self.position_ids[:, :seqlen].copy_(position_ids)
 
-        if self.extra_model_inputs:
-            assert extra_model_inputs is not None, "Model was captured with extra model inputs, so extra_model_inputs must be provided to run()"
-            for key in self.extra_model_inputs:
-                assert key in extra_model_inputs, f"Graph runner is missing extra input {key}"
-                dst_tensor = self.extra_model_inputs[key]
-                dst_tensor.copy_(extra_model_inputs[key])
-
         assert self._output is not None and self._graph is not None
         self._graph.replay()
         return self._output

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -281,6 +281,7 @@ def __init__(
         self.py_rewind_len = 0
         self.py_draft_tokens = [] if self.draft_tokens is None else self.draft_tokens
         self.py_last_draft_tokens = None
+        self.py_num_accepted_draft_tokens = 0
         self.py_decoding_iter = 0
         self.is_attention_dp_dummy = False
         self.is_cuda_graph_dummy = False

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -77,7 +77,6 @@ def forward(self,
                 scheduled_requests: ScheduledRequests,
                 resource_manager: ResourceManager,
                 new_tensors_device: Optional[SampleStateTensors],
-                extra_model_inputs: Optional[Dict[str, Any]],
                 gather_context_logits: bool = False):
         raise NotImplementedError
 
@@ -338,6 +337,7 @@ def __init__(
         spec_config: Optional[SpecConfig] = None,
         guided_decoding_config: Optional[GuidedDecodingConfig] = None,
         lora_config: Optional[LoraConfig] = None,
+        is_draft_model: bool = False,
     ):
         self.ub_buffers = None
         self.batch_size = batch_size
@@ -353,10 +353,7 @@ def __init__(
         self.pytorch_backend_config = pytorch_backend_config
         self.spec_config = spec_config
         self.is_spec_decode = spec_config is not None
-        # We keep a reference to the last used spec metadata to
-        # accommodate certain target/draft model use cases. See
-        # py_executor.py for how this is used.
-        self.last_spec_metadata = None
+        self.is_draft_model = is_draft_model
 
         self.in_warmup = False
 
@@ -530,6 +527,15 @@ def wrapper(self, *args, **kwargs):
 
         return wrapper
 
+    @contextlib.contextmanager
+    def no_cuda_graph(self):
+        _run_cuda_graphs = self._run_cuda_graphs
+        self._run_cuda_graphs = False
+        try:
+            yield
+        finally:
+            self._run_cuda_graphs = _run_cuda_graphs
+
     @with_warmup_flag
     def warmup(self, resource_manager: ResourceManager) -> None:
         kv_cache_manager = resource_manager.get_resource_manager(
@@ -654,7 +660,7 @@ def get_autotune_warmup_request():
             result.context_requests = requests
             result.generation_requests = []
 
-            return result, _create_extra_inputs(1, maximum_tunable_num_tokens)
+            return result
 
         @contextlib.contextmanager
         def release_batch(result):
@@ -668,29 +674,6 @@ def release_batch(result):
                         if spec_resource_manager is not None:
                             spec_resource_manager.free_resources(req)
 
-        @contextlib.contextmanager
-        def no_cuda_graph():
-            _run_cuda_graphs = self._run_cuda_graphs
-            self._run_cuda_graphs = False
-            try:
-                yield
-            finally:
-                self._run_cuda_graphs = _run_cuda_graphs
-
-        def _create_extra_inputs(bs, num_tokens_per_request):
-            if self.spec_config is None:
-                extra_model_inputs = None
-            else:
-                warmup_inputs_creator = getattr(self.model,
-                                                "get_warmup_extra_inputs", None)
-                if callable(warmup_inputs_creator):
-                    extra_model_inputs = warmup_inputs_creator(
-                        bs, num_tokens_per_request)
-                else:
-                    extra_model_inputs = None
-
-            return extra_model_inputs
-
         # TODO: current warmup_request is not suitable for star attention
         cp_type = self.mapping.cp_config.get('cp_type', None)
         if cp_type == 'star_attention':
@@ -712,7 +695,7 @@ def disable_optimization(backend: Backend):
                 set_enable_piecewise_cuda_graph_capture_flag(True)
 
                 # Disable cuda graph capture here so that we can properly capture it later
-                with no_cuda_graph():
+                with self.no_cuda_graph():
                     available_tokens = kv_cache_manager.get_num_available_tokens(
                         self.max_draft_len)
                     warmup_batch_size = [1, self.batch_size // 2]
@@ -733,17 +716,14 @@ def disable_optimization(backend: Backend):
                                 logger.info(
                                     f"Run warmup for batch size={bs}, pure {'context' if num_tokens_per_request > 1 else 'generation'} phase"
                                 )
-                                self.forward(
-                                    batch,
-                                    new_tensors_device=None,
-                                    resource_manager=resource_manager,
-                                    extra_model_inputs=_create_extra_inputs(
-                                        bs, num_tokens_per_request))
+                                self.forward(batch,
+                                             new_tensors_device=None,
+                                             resource_manager=resource_manager)
                                 torch.cuda.synchronize()
 
             if self.pytorch_backend_config.autotuner_enabled:
-                with no_cuda_graph(), autotune():
-                    result, extra_model_inputs = get_autotune_warmup_request()
+                with self.no_cuda_graph(), autotune():
+                    result = get_autotune_warmup_request()
                     with release_batch(result) as batch:
                         if batch is None:
                             # No KV cache space!
@@ -753,8 +733,7 @@ def disable_optimization(backend: Backend):
                                 f"Run autotuning warmup for batch size={1}")
                             self.forward(batch,
                                          new_tensors_device=None,
-                                         resource_manager=resource_manager,
-                                         extra_model_inputs=extra_model_inputs)
+                                         resource_manager=resource_manager)
                             torch.cuda.synchronize()
 
                     logger.info(f"Autotuner Cache size after warmup " +
@@ -783,12 +762,11 @@ def disable_optimization(backend: Backend):
                     )
                     self.forward(batch,
                                  new_tensors_device=None,
-                                 resource_manager=resource_manager,
-                                 extra_model_inputs=_create_extra_inputs(bs, 1))
+                                 resource_manager=resource_manager)
                     torch.cuda.synchronize()
 
                 if self._torch_compile_piecewise_cuda_graph:
-                    with no_cuda_graph():
+                    with self.no_cuda_graph():
                         with release_batch(
                                 get_torch_compile_warmup_request(1,
                                                                  bs)) as batch:
@@ -797,17 +775,12 @@ def disable_optimization(backend: Backend):
                             )
 
                             for _ in range(3):
-                                self.forward(
-                                    batch,
-                                    new_tensors_device=None,
-                                    resource_manager=resource_manager,
-                                    extra_model_inputs=_create_extra_inputs(
-                                        1, bs))
-                            self.forward(
-                                batch,
-                                new_tensors_device=None,
-                                resource_manager=resource_manager,
-                                extra_model_inputs=_create_extra_inputs(1, bs))
+                                self.forward(batch,
+                                             new_tensors_device=None,
+                                             resource_manager=resource_manager)
+                            self.forward(batch,
+                                         new_tensors_device=None,
+                                         resource_manager=resource_manager)
                             torch.cuda.synchronize()
                             gc.collect()
                             torch.cuda.empty_cache()
@@ -851,15 +824,17 @@ def _set_up_spec_metadata(
                 self.spec_config,
                 self.batch_size,
                 max_num_tokens=self.max_num_tokens,
-                spec_resource_manager=spec_resource_manager)
+                spec_resource_manager=spec_resource_manager,
+                is_draft_model=self.is_draft_model)
 
         if self.spec_metadata is not None:
             return self.spec_metadata
         self.spec_metadata = get_spec_metadata(
             self.spec_config,
             self.batch_size,
             max_num_tokens=self.max_num_tokens,
-            spec_resource_manager=spec_resource_manager)
+            spec_resource_manager=spec_resource_manager,
+            is_draft_model=self.is_draft_model)
         return self.spec_metadata
 
     def _get_padded_batch(self, scheduled_requests: ScheduledRequests,
@@ -1155,7 +1130,6 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                 inputs['attn_metadata'].kv_lens_cuda[
                     num_ctx_requests:num_seqs] += (
                         self.previous_kv_lens_offsets_cuda[:num_gen_requests])
-
         return inputs
 
     def _prepare_tp_inputs(
@@ -1476,6 +1450,7 @@ def _prepare_tp_inputs(
         lora_params = self._get_lora_params_from_requests(
             scheduled_requests, attn_metadata)
 
+        # Prepare inputs
         inputs = {
             'attn_metadata': attn_metadata,
             'input_ids': self.input_ids_cuda[:total_num_tokens],
@@ -2027,7 +2002,6 @@ def forward(self,
                 scheduled_requests: ScheduledRequests,
                 resource_manager: ResourceManager,
                 new_tensors_device: Optional[SampleStateTensors] = None,
-                extra_model_inputs: Optional[Dict[str, Any]] = None,
                 gather_context_logits: bool = False):
 
         kv_cache_manager = resource_manager.get_resource_manager(
@@ -2055,9 +2029,6 @@ def forward(self,
         if kv_cache_manager is None:
             inputs, gather_ids = self._prepare_tp_inputs_no_cache(
                 scheduled_requests, attn_metadata, spec_metadata)
-            if extra_model_inputs is not None:
-                inputs.update(extra_model_inputs)
-            self.last_spec_metadata = spec_metadata
 
             with MoeLoadBalancerIterContext(moe_load_balancer):
                 return self._forward_step(inputs, gather_ids,
@@ -2081,9 +2052,6 @@ def forward(self,
                                                       attn_metadata,
                                                       spec_metadata,
                                                       new_tensors_device)
-            if extra_model_inputs is not None:
-                inputs.update(extra_model_inputs)
-            self.last_spec_metadata = spec_metadata
 
             self.iter_counter += 1
 
@@ -2104,16 +2072,15 @@ def capture_forward_fn(inputs: Dict[str, Any]):
                     pool = maybe_graph.capture(
                         capture_forward_fn,
                         self._cuda_graph_mem_pool,
-                        extra_model_inputs,
                     )
                     self._cuda_graph_mem_pool = pool
 
                     # here we don't need to use context since cuda graph capture didn't run kernel.
                     # maybe we need a cleaner way to do this.
-                    outputs = maybe_graph.run(inputs, extra_model_inputs)
+                    outputs = maybe_graph.run(inputs)
                 else:
                     with MoeLoadBalancerIterContext(moe_load_balancer):
-                        outputs = maybe_graph.run(inputs, extra_model_inputs)
+                        outputs = maybe_graph.run(inputs)
 
             # Note: To overlap the CPU and GPU computation as much as possible,
             # guided_decoder.build should be called immediately after the launch of the single step;
@@ -2148,6 +2115,8 @@ def _forward_step(self,
                       gather_ids: Optional[torch.Tensor],
                       gather_context_logits: bool = False) -> Dict[str, Any]:
         inputs = self._preprocess_inputs(inputs)
+        if inputs.get('spec_metadata', None):
+            gather_ids = inputs['spec_metadata'].gather_ids
         if self.without_logits:
             outputs = self.model_forward(**inputs)
             return outputs