[CUDAGraph]Add support for custom all-reduce operators under SOT mode (#4386)

2025-12-24 13:28:13 +08:00 · 2025-10-16 19:31:19 +08:00
parent 26ff2f8683
commit b87e2c6184
3 changed files with 5 additions and 3 deletions
--- a/fastdeploy/worker/gpu_model_runner.py
+++ b/fastdeploy/worker/gpu_model_runner.py
@@ -1575,6 +1575,7 @@ class GPUModelRunner(ModelRunnerBase):
                self.proposer.update_task_chunk_prefill(task)
            task.chunk_idx += 1

+    @sot_warmup_guard(True)
    def capture_model(self) -> None:
        """
        Trigger CUDA Graph capture for all shapes in cuda graph capture list