Method append

tensorrt_llm/_torch/pyexecutor/llm_request.py:101–127

(self, logits: torch.Tensor)

            requires_grad=False)

    def append(self, logits: torch.Tensor):
        if logits.ndim == 2:
            logits = logits.unsqueeze(1)
        assert logits.ndim == 3, f"Bad logits shape, expect [num_tokens, beam_width, vocab_size], got {logits.shape}"

        if self.use_chunked_generation_logits:
            if self.beam_width == -1:
                self._init_chunked_storage(logits)
            self._add_fragment(logits)
        else:
            if self.beam_width == -1:
                self._init(logits)

            assert logits.size(1) == self.beam_width, "Beam width mismatch"

            position = 0 if not self._logits_indices else self._logits_indices[
                -1][1]
            new_position = logits.size(0) + position
            if new_position > self.seq_length:
                raise ValueError(
                    f"LogitsStorage overflow. This storage can only hold {self.seq_length} logits "
                    f"({position} already filled) but trying to append {logits.size(0)} more logits"
                )

            self._storage[position:new_position].copy_(logits,
                                                       non_blocking=True)
            self._logits_indices.append((position, new_position))

    def get(self, all_logits: bool, exclude_last: bool) -> torch.Tensor | None:
        """Returns the used logits storage if there are any, otherwise, returns None.