[feat] Improve 2-model perf

Signed-off-by: Mike Iovine <[email protected]>

Author: mikeiovine

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -293,18 +293,6 @@ def load_weights_from_target_model(self,
         if self.load_lm_head_from_target:
             self.lm_head = target_model.lm_head
 
-    # TODO: should input/position IDs be included in this? Keeping it implicit
-    # for now since the shapes/dtypes are the same across all models we have.
-    def get_warmup_extra_inputs(self, batch_size: int,
-                                num_tokens: int) -> Dict[str, Any]:
-
-        hidden_states = torch.empty(batch_size * num_tokens,
-                                    self.model.hidden_size,
-                                    dtype=self.model.dtype,
-                                    device='cuda')
-
-        return {'hidden_states': hidden_states}
-
     def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """
         Hack for eagle3. We might need to run a matmul to reduce

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -9,7 +9,7 @@
 import weakref
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, Callable
 
 import torch
 import torch._dynamo.config
@@ -274,6 +274,8 @@ def __init__(
         spec_config: Optional["DecodingBaseConfig"] = None,
         lora_config: Optional[LoraConfig] = None,
         is_draft_model: bool = False,
+        model_wrapper: Optional[Callable[[torch.nn.Module],
+                                         torch.nn.Module]] = None,
     ):
         self.ub_buffers = None
         self.batch_size = batch_size
@@ -309,7 +311,8 @@ def __init__(
             max_num_tokens=max_num_tokens,
             moe_max_num_tokens=pytorch_backend_config.moe_max_num_tokens,
             moe_load_balancer=pytorch_backend_config.moe_load_balancer,
-            lora_config=lora_config)
+            lora_config=lora_config,
+            model_wrapper=model_wrapper)
         # In case that some tests use stub models and override `_load_model`.
         if not hasattr(self.model, 'extra_attrs'):
             self.model.extra_attrs = {}
@@ -387,7 +390,7 @@ def __init__(
             self.spec_metadata = None
             update_spec_config_from_model_config(self.spec_config,
                                                  self.model.config)
-            max_num_draft_tokens = self.spec_config.max_draft_len * batch_size
+            max_num_draft_tokens = self.spec_config.max_draft_len * batch_size if not self.is_draft_model else 0
             self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ),
                                                  dtype=torch.int,
                                                  device='cuda')
@@ -401,9 +404,11 @@ def __init__(
             self.previous_kv_lens_offsets_cuda = torch.zeros((batch_size, ),
                                                              dtype=torch.int,
                                                              device='cuda')
+            # TODO undo this hack
             self.without_logits = self.spec_config.spec_dec_mode.without_logits(
-            )
-            self.max_draft_len = spec_config.max_draft_len
+            ) or (self.is_draft_model
+                  and self.spec_config.spec_dec_mode.is_eagle3())
+            self.max_draft_len = spec_config.max_draft_len if not self.is_draft_model else 0
         else:
             self.without_logits = False
             self.max_draft_len = 0
@@ -464,7 +469,7 @@ def __init__(
 
     @property
     def runtime_draft_len(self):
-        return self.max_draft_len if self.enable_spec_decode else 0
+        return self.max_draft_len if self.enable_spec_decode and not self.is_draft_model else 0
 
     def set_lora_model_config(self,
                               lora_target_modules: list[str],
@@ -902,6 +907,8 @@ def _load_model(self,
                     moe_max_num_tokens: Optional[int] = None,
                     moe_load_balancer: Optional[MoeLoadBalancerConfig] = None,
                     lora_config: Optional[LoraConfig] = None,
+                    model_wrapper: Optional[Callable[[torch.nn.Module],
+                                                     torch.nn.Module]] = None,
                     **kwargs) -> DecoderModelForCausalLM:
         config = checkpoint_loader.load_config(
             checkpoint_dir,
@@ -1005,6 +1012,10 @@ def init_meta_tensor(t: torch.Tensor):
                 logger.info("moe_load_balancer finalize model done")
 
             torch.cuda.current_stream().synchronize()
+
+        if model_wrapper is not None:
+            model = model_wrapper(model)
+
         return model
 
     def _call_load_weights(self, load_method, weights, weight_mapper):
@@ -1257,7 +1268,7 @@ def _prepare_tp_inputs(
                 past_seen_token_num = request.max_beam_num_tokens - 1
                 draft_lens.append(num_draft_tokens)
 
-                if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                if self.enable_spec_decode and not self.is_draft_model and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend):
                     # We're treating the prompt lengths as context requests here, so
                     # the the prompt lens should not include the cached tokens.

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -252,13 +252,19 @@ def create_py_executor(
         with mem_monitor.observe_creation_stage(
                 _ExecutorCreationStage.MODEL_ENGINE_DRAFT):
             draft_spec_config = copy.copy(spec_config)
-            draft_pytorch_backend_config = copy.copy(pytorch_backend_config)
-            if spec_config.load_format == "dummy":
-                draft_pytorch_backend_config.load_format = LoadFormat.DUMMY
             # The draft model won't have any draft tokens attached to
             # generation requests when we invoke it autoregressively
             draft_spec_config.max_draft_len = 0
 
+            def wrap_model(model):
+                from tensorrt_llm._torch.speculative.eagle3 import ChainDrafter
+
+                return ChainDrafter(spec_config.max_draft_len, model)
+
+            draft_pytorch_backend_config = copy.copy(pytorch_backend_config)
+            if spec_config.load_format == "dummy":
+                draft_pytorch_backend_config.load_format = LoadFormat.DUMMY
+
             draft_model_engine = PyTorchModelEngine(
                 model_path=spec_config.speculative_model_dir,
                 pytorch_backend_config=draft_pytorch_backend_config,
@@ -274,6 +280,7 @@ def create_py_executor(
                 spec_config=draft_spec_config,
                 checkpoint_loader=executor_config.checkpoint_loader,
                 is_draft_model=True,
+                model_wrapper=wrap_model,
             )
             draft_model_engine.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER
             draft_model_engine.load_weights_from_target_model(

tensorrt_llm/_torch/speculative/eagle3.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import List, Optional, Set
+from typing import List, Optional, Set, Tuple, Any, Dict
 
 import torch
 from torch import nn
@@ -503,3 +503,86 @@ def prepare_1st_drafter_inputs(
             "attn_metadata": attn_metadata,
             "spec_metadata": spec_metadata,
         }
+
+
+class ChainDrafter(torch.nn.Module):
+
+    def __init__(self, max_draft_len: int, draft_model: torch.nn.Module):
+        super().__init__()
+        self.draft_model = draft_model
+        self.config = self.draft_model.config
+        self.model_config = self.draft_model.model_config
+        self.max_draft_len = max_draft_len
+
+    def forward(self, input_ids, position_ids, attn_metadata, spec_metadata,
+                **kwargs):
+        batch_size = attn_metadata.num_seqs
+
+        logits = self.draft_model.forward(input_ids=input_ids,
+                                          position_ids=position_ids,
+                                          attn_metadata=attn_metadata,
+                                          spec_metadata=spec_metadata)
+
+        new_draft_tokens = [self.sample(logits)]
+
+        if attn_metadata.is_cuda_graph:
+            seq_len = attn_metadata._seq_lens[:batch_size].clone()
+            seq_len_cuda = attn_metadata._seq_lens_cuda[:batch_size].clone()
+
+        last_tokens_idx = torch.cumsum(
+            attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
+        new_position_ids = position_ids[0, last_tokens_idx] + 1
+
+        attn_metadata._seq_lens[:batch_size].fill_(1)
+        attn_metadata._seq_lens_cuda[:batch_size].fill_(1)
+        attn_metadata.on_update()
+        attn_metadata.kv_lens_cuda[:batch_size] += 1
+
+        attn_metadata.host_request_types[:attn_metadata.num_contexts].fill_(1)
+        attn_metadata.num_contexts = 0
+
+        spec_metadata.eagle3_resource_manager.is_first_draft = False
+        spec_metadata.is_first_draft = False
+
+        old_write_indices = spec_metadata.hidden_states_write_indices
+
+        spec_metadata.hidden_states_read_indices[:batch_size].copy_(
+            old_write_indices[last_tokens_idx])
+        spec_metadata.hidden_states_write_indices[:batch_size].copy_(
+            torch.arange(
+                batch_size,
+                dtype=spec_metadata.hidden_states_write_indices.dtype,
+                device=spec_metadata.hidden_states_write_indices.device))
+        spec_metadata.num_tokens = batch_size
+
+        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))
+            new_position_ids += 1
+            attn_metadata.kv_lens_cuda[:batch_size] += 1
+            if i == 0:
+                spec_metadata.hidden_states_read_indices[:batch_size].copy_(
+                    spec_metadata.hidden_states_write_indices[:batch_size])
+
+        spec_metadata.is_first_draft = True
+        spec_metadata.eagle3_resource_manager.is_first_draft = True
+
+        if attn_metadata.is_cuda_graph:
+            attn_metadata._seq_lens[:batch_size].copy_(seq_len[:batch_size])
+            attn_metadata._seq_lens_cuda[:batch_size].copy_(
+                seq_len_cuda[:batch_size])
+
+        return torch.stack(new_draft_tokens)
+
+    def sample(self, logits: torch.Tensor) -> torch.Tensor:
+        tokens = torch.argmax(logits, dim=-1)
+        d2t = self.draft_model.model.d2t.data
+
+        return tokens + d2t[tokens]
+
+    def load_weights_from_target_model(self,
+                                       target_model: torch.nn.Module) -> None:
+        self.draft_model.load_weights_from_target_model(target_model)

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -72,6 +72,8 @@ def __init__(
             self._request_draft_logits = sampler.enable_mixed_sampler
         self.guided_decoder = guided_decoder
 
+        self.use_static_draft_loop = True
+
     def _create_draft_request(self, request: LlmRequest,
                               input_tokens: Optional[List]) -> LlmRequest:
         """Create a draft request with common parameters."""
@@ -237,6 +239,8 @@ def _should_disable_cuda_graph(
         """Check if CUDA graph should be disabled for the current forward pass."""
         if previous_batch is not None:
             return False
+        if self.use_static_draft_loop:
+            return False
         return self.spec_config.spec_dec_mode.needs_kv_cache_recompute()
 
     def _forward_draft_model(
@@ -257,8 +261,8 @@ def _forward_draft_model(
                 new_tensors_device=new_tensors_device)
 
         # Handle d2t data if available
-        if hasattr(self.draft_model_engine.model.model, 'd2t'):
-            outputs['d2t'] = self.draft_model_engine.model.model.d2t.data
+        # if hasattr(self.draft_model_engine.model.model, 'd2t'):
+        #     outputs['d2t'] = self.draft_model_engine.model.model.d2t.data
 
         return outputs
 
@@ -377,6 +381,17 @@ def prepare_draft_tokens(
 
             # Initial forward pass
             outputs = self._forward_draft_model(draft_batch, resource_manager)
+
+            if self.use_static_draft_loop:
+                outputs_host = outputs.cpu()
+                for token_idx in range(self.max_draft_tokens):
+                    for req_idx, req in enumerate(draft_batch.all_requests()):
+                        target_req = req_id_to_old_request[req.py_request_id]
+                        target_req.py_draft_tokens.append(
+                            outputs_host[token_idx][req_idx])
+
+                return
+
             self._execute_guided_decoder(draft_batch,
                                          outputs['logits'],
                                          d2t=outputs.get('d2t'))