[TRTLLM-7384][feat] enable rejection sampling for CDL (NVIDIA#7731)

Signed-off-by: linquanh <linquanh@nvidia.com>

Author: kris1025

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -456,6 +456,8 @@ def __init__(
         self.use_draft_model = is_draft
         # Whether the request is for the first forward of the draft model.
         self.py_is_first_draft = is_first_draft
+        self.d2t = None
+        self.py_draft_use_greedy_sampling = False
 
         # Chunked logits parameters
         self.py_use_chunked_generation_logits = use_chunked_generation_logits

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -31,7 +31,6 @@
 from ..distributed import MPIDist, TorchDist
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager)
-from ..utils import _get_allow_chain_drafter
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
                     validate_feature_combination)
@@ -344,13 +343,11 @@ def create_py_executor(
                 _ExecutorCreationStage.MODEL_ENGINE_DRAFT):
             draft_spec_config = copy.copy(spec_config)
 
-            if _get_allow_chain_drafter():
-                use_chain_drafter = (
-                    guided_decoding_config is None
-                    and draft_spec_config._allow_greedy_draft_tokens
-                    and pytorch_backend_config.attn_backend == "TRTLLM")
-            else:
-                use_chain_drafter = False
+            use_chain_drafter = (
+                guided_decoding_config is None
+                and draft_spec_config._allow_chain_drafter
+                and draft_spec_config._allow_greedy_draft_tokens
+                and pytorch_backend_config.attn_backend == "TRTLLM")
 
             logger.debug(f"USE CHAIN DRAFTER: {use_chain_drafter}")
             if use_chain_drafter:

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -310,10 +310,15 @@ def greedy_search_sampling_batch(
     softmax_indices: Optional[torch.IntTensor] = None
 ) -> tuple[torch.Tensor, torch.Tensor]:
     next_tokens = torch.argmax(logits, dim=-1)
+    index_to_scatter = next_tokens
     if softmax_indices is not None:
-        logits = logits[softmax_indices.to(logits.device, non_blocking=True)]
-    softmax = torch.softmax(logits, dim=-1)
-    return next_tokens, softmax
+        logits = logits[softmax_indices]
+        index_to_scatter = next_tokens[softmax_indices]
+    probs = torch.zeros_like(logits)
+    probs.scatter_(dim=-1,
+                   index=index_to_scatter.unsqueeze(-1),
+                   src=torch.ones_like(logits))
+    return next_tokens, probs
 
 
 def get_rejected_indices(draft_probs: torch.Tensor, target_probs: torch.Tensor,
@@ -1127,20 +1132,38 @@ def _tree_sampling_batch(self, requests: list[LlmRequest],
 
         return new_draft_tokens_host
 
+    @torch.inference_mode()
     def _process_draft_tokens_rejection_sampling(
             self, request: LlmRequest, new_tokens: list[list[list[int]]],
             new_tokens_tensor: torch.Tensor) -> int:
         # FIXME: Passing a dummy vocab_size could result in unnecessary
         #        filtering of vocab_size logits, out of vocab_size in
         #        total. The 'sample' below should generally be avoided
         #        by retaining the draft_probs during drafting (TRTLLM-7772).
-        sampling_strategy = _request_strategy(request, vocab_size=2**31)
+        draft_sampling_strategy = (
+            "greedy", None
+        ) if request.py_draft_use_greedy_sampling else _request_strategy(
+            request, vocab_size=2**31)
         generator = self.get_generator(request.py_draft_logits.device)
-        _, draft_probs = sample(sampling_strategy,
+        _, draft_probs = sample(draft_sampling_strategy,
                                 request.py_draft_logits,
                                 generator=generator)
-        draft_probs = draft_probs.squeeze(0)
         target_probs = request.py_target_probs
+        d2t = getattr(request, "d2t", None)
+        if d2t is not None:
+            vocab_d = draft_probs.shape[-1]
+            vocab_t = target_probs.shape[-1]
+            assert d2t.numel(
+            ) == vocab_d, f"d2t size mismatch: {d2t.numel()} != {vocab_d}"
+            assert d2t.device == draft_probs.device, f"d2t device mismatch: {d2t.device} != {draft_probs.device}"
+            aligned_draft_probs = torch.zeros(
+                (*draft_probs.shape[:-1], vocab_t),
+                device=draft_probs.device,
+                dtype=draft_probs.dtype)
+            source_indices = torch.arange(vocab_d, device=draft_probs.device)
+            target_indices = (source_indices + d2t) % vocab_t
+            aligned_draft_probs[..., target_indices] = draft_probs
+            draft_probs = aligned_draft_probs
         rejected_indices = get_rejected_indices(draft_probs, target_probs,
                                                 generator,
                                                 request.py_draft_tokens)
@@ -1181,7 +1204,8 @@ def process_draft_tokens(
             new_tokens: list[list[list[int]]],
             new_tokens_tensor: torch.Tensor,
             resource_manager: Optional[ResourceManager] = None) -> int:
-        if request.py_draft_logits is None:
+        if _request_strategy(request, vocab_size=2**
+                             31) == GREEDY or request.py_draft_logits is None:
             spec_tree_manager = self.get_spec_tree_manager(resource_manager)
             if spec_tree_manager is not None:
                 num_accepted = self._process_draft_tokens_tree(

tensorrt_llm/_torch/speculative/drafting_loops.py

@@ -116,8 +116,7 @@ def __init__(self, max_draft_len: int, draft_model: torch.nn.Module):
 
     def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
                 attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata,
-                **kwargs) -> torch.Tensor:
-
+                **kwargs) -> dict[str, torch.Tensor]:
         logits = self.draft_model.forward(input_ids=input_ids,
                                           position_ids=position_ids,
                                           attn_metadata=attn_metadata,
@@ -126,6 +125,7 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
         logits = logits[spec_metadata.gather_ids]
 
         new_draft_tokens = [self.sample(logits)]
+        draft_logits = [logits]
         with save_metadata_state(attn_metadata, spec_metadata):
             batch_size = attn_metadata.num_seqs
 
@@ -139,13 +139,17 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
                     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])
 
-        return torch.stack(new_draft_tokens)
+        return {
+            "new_draft_tokens": torch.stack(new_draft_tokens),
+            "draft_logits": torch.stack(draft_logits)
+        }
 
     def sample(self, logits: torch.Tensor) -> torch.Tensor:
         # TODO: inject the sampler here so we can support non-greedy

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -231,6 +231,12 @@ def _prepare_draft_batch(
             ScheduledRequests: The prepared draft batch
         """
         try:
+            for req in scheduled_requests.all_requests():
+                draft_model = self.draft_model_engine.model.draft_model if self.use_static_draft_loop else self.draft_model_engine.model
+                if hasattr(draft_model.model, "d2t"):
+                    req.d2t = draft_model.model.d2t.data
+                req.py_draft_use_greedy_sampling = self.use_static_draft_loop
+
             draft_batch = ScheduledRequests()
 
             for request in scheduled_requests.context_requests:
@@ -530,7 +536,8 @@ def _setup_draft_batch_and_resources(
         return draft_batch, req_id_to_old_request
 
     def process_static_draft_outputs(
-            self, outputs: torch.Tensor | SampleState,
+            self,
+            outputs: dict[str, torch.Tensor] | tuple[torch.Tensor, SampleState],
             draft_batch: ScheduledRequests,
             req_id_to_old_request: Dict[int, LlmRequest]) -> None:
         """
@@ -541,23 +548,26 @@ def process_static_draft_outputs(
             draft_batch: The draft batch that was processed
             req_id_to_old_request: Mapping from draft request ID to original request
         """
-        if isinstance(outputs, torch.Tensor):
-            # For non-overlap scheduler path.
-            outputs_host = outputs.cpu()
+
+        if isinstance(outputs, dict):
+            draft_tokens_host = outputs["new_draft_tokens"].cpu()
+            draft_logits = outputs["draft_logits"]
         else:
-            outputs_host = outputs.host.new_tokens
-            outputs.sampler_event.synchronize()
-
-        for token_idx in range(self.max_draft_tokens):
-            for req_idx, req in enumerate(draft_batch.all_requests()):
-                target_model_req = req_id_to_old_request[req.py_request_id]
-                if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS:
-                    # Chunked prefill request in progress; no need to append draft tokens
-                    continue
+            draft_logits = outputs[0]
+            draft_tokens_host = outputs[1].host.new_tokens
+            outputs[1].sampler_event.synchronize()
 
-                target_req = req_id_to_old_request[req.py_request_id]
-                target_req.py_draft_tokens.append(
-                    outputs_host[token_idx][req_idx])
+        for req_idx, req in enumerate(draft_batch.all_requests()):
+            target_model_req = req_id_to_old_request[req.py_request_id]
+            if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS:
+                # Chunked prefill request in progress; no need to append draft tokens
+                continue
+            py_draft_logits = []
+            for token_idx in range(self.max_draft_tokens):
+                target_model_req.py_draft_tokens.append(
+                    draft_tokens_host[token_idx][req_idx])
+                py_draft_logits.append(draft_logits[token_idx][req_idx])
+            target_model_req.py_draft_logits = torch.stack(py_draft_logits)
 
         # Clean up draft resources
         for req in draft_batch.all_requests():
@@ -710,23 +720,26 @@ def generate_draft_tokens_with_overlap(
             # Only update target inputs, cleanup will be done in executor loop
             self._update_target_inputs_with_draft_tokens(
                 target_inputs,
-                outputs,
+                outputs["new_draft_tokens"],
                 draft_position=0,
                 draft_length=self.max_draft_tokens,
                 draft_batch=draft_batch,
                 req_id_to_old_request=req_id_to_old_request)
 
-            new_tokens_host = outputs.to(device='cpu', non_blocking=True)
+            new_tokens_host = outputs["new_draft_tokens"].to(device='cpu',
+                                                             non_blocking=True)
             sampler_event = torch.cuda.Event()
             sampler_event.record()
 
-            outputs = SampleState(
+            sample_state = SampleState(
                 scheduled_requests=draft_batch,
-                device=SampleStateTensors(new_tokens=outputs),
+                device=SampleStateTensors(
+                    new_tokens=outputs["new_draft_tokens"]),
                 host=SampleStateTensors(new_tokens=new_tokens_host),
                 sampler_event=sampler_event)
 
-            return target_inputs, outputs, draft_batch
+            return target_inputs, (outputs["draft_logits"],
+                                   sample_state), draft_batch
 
         # Handle guided decoder and sampling for non-static loop
         if self.guided_decoder is not None:

tensorrt_llm/_torch/utils.py

@@ -308,12 +308,6 @@ def create_lm_head_tp_mapping(mapping: Mapping, token_count: int) -> Mapping:
     )
 
 
-# Development function to control chain drafter feature.
-# It's here so that unit tests can mock it and turn it off.
-def _get_allow_chain_drafter() -> bool:
-    return True
-
-
 def get_device_uuid(device_idx: int) -> str:
     """Get the UUID of a CUDA device using torch cuda api"""
 

tensorrt_llm/llmapi/llm_args.py

@@ -366,6 +366,8 @@ class DecodingBaseConfig(StrictBaseModel):
 
     load_format: Optional[str] = None
 
+    # If set, drafting is allowed to use chain drafter.
+    _allow_chain_drafter: bool = PrivateAttr(True)
     # If set, drafting uses greedy sampling, irrespective of sampling parameters.
     _allow_greedy_draft_tokens: bool = PrivateAttr(True)