Method capture_model

fastdeploy/worker/xpu_model_runner.py:1483–1512

Trigger CUDA Graph capture for all shapes in 'CudaGraphConfig.cudagraph_capture_sizes'

(self)

            paddle.device.xpu.set_debug_level(debug_level)

    def capture_model(self) -> None:
        """
        Trigger CUDA Graph capture for all shapes in 'CudaGraphConfig.cudagraph_capture_sizes'
        """
        time_before_capture = time.perf_counter()
        expected_decode_len = 1
        capture_sizes = self.cudagraph_capture_sizes.copy()

        try:
            for batch_size in sorted(capture_sizes, reverse=True):
                self._dummy_run(
                    num_tokens=self.scheduler_config.max_num_batched_tokens,
                    batch_size=batch_size,
                    expected_decode_len=expected_decode_len,
                    in_capturing=True,
                )
                logger.info(f"Warm up the model with the batch size:{batch_size}, num tokens:{expected_decode_len}")
        except RuntimeError as e:
            if "out of memory" in str(e):
                raise RuntimeError(
                    "CUDA out of memory occurred when warming up CUDAGraph "
                    f"with the capture sizes {capture_sizes}. Please try "
                    "lowering `max_num_seqs` or `gpu_memory_utilization` when "
                    "initializing the engine."
                ) from e
            else:
                raise e

        time_after_capture = time.perf_counter()
        logger.info(f"Cuda Graph capturing took {time_after_capture - time_before_capture} seconds")

    @sot_warmup_guard(True)
    def sot_warmup(self) -> None: