diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 14ba2e3a6ab2..2c6d2f54a221 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -610,6 +610,23 @@ class SpecMetadata: request_temperatures: Optional[torch.Tensor] = None request_top_ks: Optional[torch.Tensor] = None request_top_ps: Optional[torch.Tensor] = None + # Describe what the sampling-parameter buffers currently hold, so a step + # that reproduces them can skip the refill. Two entries because the two + # buffer groups depend on different things: + # [0] request_* -- the per-request values, in batch order. + # [1] the expanded per-token buffers, which additionally depend on each + # request's token count (a context request contributing one row + # instead of draft_len + 1 shifts every later request's offset). + # A context->generation transition therefore invalidates [1] while leaving + # [0] valid. See _sampling_params_buffers_need_update. + # + # Held in a list rather than plain fields because + # create_cuda_graph_metadata shallow-copies this object: the graph views + # and the eager view write the *same* tensors, so they must agree on what + # those tensors hold. Plain fields would give each view its own stale + # answer and let one skip a fill another view invalidated. + _sampling_params_signature: list = field( + default_factory=lambda: [None, None], repr=False) # Whether to use sampling parameters when sampling draft tokens. use_sampling_params_for_draft_tokens: bool = False # Vocab size used for draft_probs buffer allocation. @@ -744,6 +761,10 @@ def create_cuda_graph_metadata(self, max_batch_size: int): cuda_graph_metadata = copy.copy(self) cuda_graph_metadata.is_cuda_graph = True cuda_graph_metadata.max_num_requests = max_batch_size + # NB: the shallow copy deliberately keeps sharing + # _sampling_params_signature with this object. Both views write the + # same sampling-parameter tensors, so the record of what those tensors + # hold has to be shared too. cuda_graph_metadata.__post_init__() return cuda_graph_metadata @@ -973,6 +994,8 @@ def populate_sampling_params_for_one_model( if self.temperatures is None or self.temperatures.numel( ) < required_flat_size: + # Fresh tensors hold none of the recorded values. + self.invalidate_sampling_params_cache() # Allocate once; the captured graph reads from these stable addresses. self.temperatures = torch.ones(required_flat_size, dtype=torch.float32, @@ -1011,53 +1034,125 @@ def populate_sampling_params_for_one_model( return # Phase 2: build per-token / per-request lists and copy to GPU. - temperatures: list[float] = [] - top_ks: list[int] = [] - top_ps: list[float] = [] - request_temperatures: list[float] = [] - request_top_ks: list[int] = [] - request_top_ps: list[float] = [] - for temp_val, tk_val, tp_val, num_tokens in per_request_normalized: - request_temperatures.append(temp_val) - request_top_ks.append(tk_val) - request_top_ps.append(tp_val) - temperatures.extend(temp_val for _ in range(num_tokens)) - top_ks.extend(tk_val for _ in range(num_tokens)) - top_ps.extend(tp_val for _ in range(num_tokens)) - - self.temperatures[:len(temperatures)].copy_(torch.tensor( - temperatures, dtype=torch.float32, pin_memory=prefer_pinned()), - non_blocking=True) - self.top_ks[:len(top_ks)].copy_(torch.tensor( - top_ks, dtype=torch.int32, pin_memory=prefer_pinned()), - non_blocking=True) - self.top_ps[:len(top_ps)].copy_(torch.tensor( - top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), - non_blocking=True) - self.request_temperatures[:len(request_temperatures)].copy_( - torch.tensor(request_temperatures, - dtype=torch.float32, - pin_memory=prefer_pinned()), - non_blocking=True) - self.request_top_ks[:len(request_top_ks)].copy_( - torch.tensor(request_top_ks, - dtype=torch.int32, - pin_memory=prefer_pinned()), - non_blocking=True, - ) - self.request_top_ps[:len(request_top_ps)].copy_( - torch.tensor(request_top_ps, - dtype=torch.float32, - pin_memory=prefer_pinned()), - non_blocking=True, - ) + # + # Sampling params are fixed for a request's lifetime, so a steady-state + # decode batch reproduces the buffers it already holds. Both the host + # expansion and the copies sit on the critical path ahead of the + # forward, so skip whichever group is already current. + need_update_sampler_param, need_update_expanded_sampler_param = ( + self._sampling_params_buffers_need_update(per_request_normalized)) + if not (need_update_sampler_param + or need_update_expanded_sampler_param): + return + + if need_update_sampler_param: + request_temperatures: list[float] = [] + request_top_ks: list[int] = [] + request_top_ps: list[float] = [] + for temp_val, tk_val, tp_val, _ in per_request_normalized: + request_temperatures.append(temp_val) + request_top_ks.append(tk_val) + request_top_ps.append(tp_val) + + self.request_temperatures[:len(request_temperatures)].copy_( + torch.tensor(request_temperatures, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True) + self.request_top_ks[:len(request_top_ks)].copy_( + torch.tensor(request_top_ks, + dtype=torch.int32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) + self.request_top_ps[:len(request_top_ps)].copy_( + torch.tensor(request_top_ps, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) - # Pre-compute top_k_max on the CPU so CUDA-graph capture does not - # encounter boolean-tensor indexing (dynamic size) or .item() calls. - # DISABLE_TOPK_VAL (INT32_MAX) is the sentinel for "top-k disabled". - _disable_topk = torch.iinfo(torch.int32).max - self.top_k_max = max( - (tk for tk in request_top_ks if 0 < tk < _disable_topk), default=0) + # Pre-compute top_k_max on the CPU so CUDA-graph capture does not + # encounter boolean-tensor indexing (dynamic size) or .item() + # calls. DISABLE_TOPK_VAL (INT32_MAX) is the "top-k disabled" + # sentinel. Derived from the same values as the per-request + # buffers, so it is refreshed exactly when they are. + _disable_topk = torch.iinfo(torch.int32).max + self.top_k_max = max( + (tk for tk in request_top_ks if 0 < tk < _disable_topk), + default=0) + + if need_update_expanded_sampler_param: + temperatures: list[float] = [] + top_ks: list[int] = [] + top_ps: list[float] = [] + for temp_val, tk_val, tp_val, num_tokens in per_request_normalized: + temperatures.extend(temp_val for _ in range(num_tokens)) + top_ks.extend(tk_val for _ in range(num_tokens)) + top_ps.extend(tp_val for _ in range(num_tokens)) + + self.temperatures[:len(temperatures)].copy_(torch.tensor( + temperatures, dtype=torch.float32, pin_memory=prefer_pinned()), + non_blocking=True) + self.top_ks[:len(top_ks)].copy_(torch.tensor( + top_ks, dtype=torch.int32, pin_memory=prefer_pinned()), + non_blocking=True) + self.top_ps[:len(top_ps)].copy_(torch.tensor( + top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), + non_blocking=True) + + def _sampling_params_buffers_need_update( + self, per_request_normalized: list[tuple[float, int, float, int]] + ) -> tuple[bool, bool]: + """Report which sampling-parameter buffers this step has to refill. + + Returns ``(need_update_sampler_param, + need_update_expanded_sampler_param)`` for the per-request buffers and + the expanded per-token buffers respectively, recording the new + signatures as a side effect. + + Both signatures are built from ``per_request_normalized``, which every + consumer reads by batch position -- so its order already encodes the + batch ordering and a reshuffle changes the signature on its own. Slot + ids are deliberately absent: they index ``batch_slot_ids`` (copied + separately for the rejection path), never these buffers, so including + them would only force refills when a slot changes hands between + requests that happen to sample identically. + + The expanded buffers additionally depend on each request's token count, + which sets their layout, so they can need a refill while the + per-request buffers stay valid -- a context request becoming a + generation request grows its span from one row to ``draft_len + 1`` and + shifts every later request. Whenever the per-request buffers need an + update the expanded ones do too. + + ``top_k_max`` derives from the same values as the per-request buffers, + so it stays valid for as long as they do. + """ + values = tuple((temp, top_k, top_p) + for temp, top_k, top_p, _ in per_request_normalized) + num_tokens = tuple(n for *_, n in per_request_normalized) + + request_signature = values + expanded_signature = (values, num_tokens) + + need_update_sampler_param = (self._sampling_params_signature[0] + != request_signature) + need_update_expanded_sampler_param = (self._sampling_params_signature[1] + != expanded_signature) + + self._sampling_params_signature[0] = request_signature + self._sampling_params_signature[1] = expanded_signature + return need_update_sampler_param, need_update_expanded_sampler_param + + def invalidate_sampling_params_cache(self) -> None: + """Force the next populate call to refill both buffer groups. + + Needed whenever the buffers stop reflecting the recorded signatures, + e.g. after reallocating them. + """ + self._sampling_params_signature[0] = None + self._sampling_params_signature[1] = None class SpecWorkerBase(nn.Module, ABC):