Method get_buffer

tensorrt_llm/_torch/memory_buffer_utils.py:52–122

Return a reusable buffer view for the requested shape/dtype. The returned tensor is backed by an underlying `torch.uint8` buffer. When no suitable buffer exists in the pool, a new tensor is created via `torch.empty`, so its contents are uninitialized. Overwrite the data befor

(self, tensor_shape: list[int], dtype: torch.dtype,
                   buffer_name: str, reserve_buffer: bool)

            target_shape)

    def get_buffer(self, tensor_shape: list[int], dtype: torch.dtype,
                   buffer_name: str, reserve_buffer: bool):
        """Return a reusable buffer view for the requested shape/dtype.
        The returned tensor is backed by an underlying `torch.uint8` buffer. When
        no suitable buffer exists in the pool, a new tensor is created via
        `torch.empty`, so its contents are uninitialized. Overwrite the data before use if needed.
        """

        # all buffers are allocated with 1 byte element size
        required_memory_size = math.prod(tensor_shape) * dtype.itemsize

        candidate_blocks = self.buffers.get(buffer_name, [])

        # Find the best-fit available buffer.
        best_fit_block: Optional[BufferBlock] = None
        smallest_sufficient_size = float('inf')
        for block in candidate_blocks:
            # Skip buffers that are too small.
            if block.buffer.numel() < required_memory_size:
                continue

            # Find the smallest buffer that is still large enough (best-fit).
            if block.buffer.numel() < smallest_sufficient_size:
                # Use reserved block if find one.
                if best_fit_block is not None and best_fit_block.is_reserved and not block.is_reserved:
                    continue

                best_fit_block = block
                smallest_sufficient_size = block.buffer.numel()

        if best_fit_block is not None:
            if reserve_buffer:
                best_fit_block.is_reserved = True
            # A suitable buffer was found, so reuse it.
            return self._view_as(best_fit_block.buffer, tensor_shape, dtype)

        for block in list(candidate_blocks):
            if not block.is_reserved:
                # Need to call del BufferBlock.buffer, otherwise memory isn't
                # released and OOM may happen.
                buffer_size = block.buffer.numel()
                del block.buffer
                if buffer_size >= 1024 * 1024 * 1024:
                    torch.cuda.empty_cache()
                candidate_blocks.remove(block)

        # No suitable buffer was found, so allocate a new one.
        # The new buffer is created with uint8 to represent raw bytes.
        new_buffer_tensor = None
        try:
            with torch.cuda.memory.use_mem_pool(get_shared_pool()):
                new_buffer_tensor = torch.empty((required_memory_size, ),
                                                device='cuda',
                                                dtype=torch.uint8)
        except Exception as ex:
            # Need to check if this is an OOM exception
            logger.debug(
                f"Exception happened to create tensor from given memory pool: {str(ex)}"