[feat] Improve 2-model perf

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

Author: mikeiovine

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -388,7 +388,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')
@@ -402,9 +402,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
@@ -1153,6 +1155,9 @@ def init_meta_tensor(t: torch.Tensor):
                 logger.info("moe_load_balancer finalize model done")
 
             torch.cuda.current_stream().synchronize()
+        if self.spec_config is not None and self.is_draft_model:
+            model = self.spec_config.get_draft_model_wrapper(model) or model
+
         return model
 
     def _call_load_weights(self, load_method, weights, weight_mapper):

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -252,9 +252,6 @@ def create_py_executor(
         with mem_monitor.observe_creation_stage(
                 _ExecutorCreationStage.MODEL_ENGINE_DRAFT):
             draft_spec_config = copy.copy(spec_config)
-            # 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
 
             draft_model_engine = PyTorchModelEngine(
                 model_path=spec_config.speculative_model_dir,

tensorrt_llm/_torch/speculative/eagle3.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Any, Dict
 
 import torch
 from torch import nn
@@ -176,8 +176,6 @@ def maybe_capture_hidden_states(
     def get_hidden_states(self):
         hidden_states = self.eagle3_resource_manager.hidden_states[
             self.hidden_states_read_indices[:self.num_tokens], :]
-        if not self.is_first_draft:
-            hidden_states = hidden_states[:, :self.hidden_size]
         return hidden_states
 
 
@@ -493,3 +491,84 @@ 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()
+
+        old_seqlens = torch.cumsum(attn_metadata.seq_lens_cuda,
+                                   dim=0,
+                                   dtype=torch.long)
+
+        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.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.hidden_states_read_indices[:batch_size].copy_(
+            old_seqlens - 1)
+        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
+
+        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 get_warmup_extra_inputs(self, batch_size: int,
+                                num_tokens: int) -> Dict[str, Any]:
+        return self.draft_model.get_warmup_extra_inputs(batch_size, num_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."""
@@ -257,8 +259,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 +379,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'))

tensorrt_llm/llmapi/llm_args.py

@@ -400,6 +400,9 @@ def spec_dec_mode(self):
         return TorchSpeculativeDecodingMode.from_string(
             self.decoding_type.upper())
 
+    def get_draft_model_wrapper(self, model):
+        return None
+
 
 class MedusaDecodingConfig(DecodingBaseConfig):
     medusa_choices: Optional[List[List[int]]] = None
@@ -443,6 +446,11 @@ def spec_dec_mode(self):
             return TorchSpeculativeDecodingMode.EAGLE3_ONE_MODEL
         return TorchSpeculativeDecodingMode.EAGLE3
 
+    def get_draft_model_wrapper(self, model):
+        from tensorrt_llm._torch.speculative.eagle3 import ChainDrafter
+
+        return ChainDrafter(self.max_draft_len, model)
+
 
 class UserProvidedDecodingConfig(DecodingBaseConfig):
     # Cannot use real type annotations due to circular imports