[ROCm] Revert Part of `[ROCm] Fix pooling startup workspace lock` #47912 (#48154)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jul 9, 2026 at 15:34 UTC
766469a4c460043ae52cda19b1c52f0dc87e555c
1 file changed
+3
-22
vllm/v1/worker/gpu_worker.py
+3
-22
@@ -729,11 +729,6 @@ class Worker(WorkerBase):
729
@instrument(span_name="Warmup (GPU)")
730
def compile_or_warm_up_model(self) -> CompilationTimes:
731
warmup_sizes: list[int] = []
732
- cg_capture_sizes: list[int] = []
733
-
734
- if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
735
- cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
736
- cg_capture_sizes = [] if cg_sizes is None else cg_sizes
732
733
if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
734
# warm up sizes that are not in cudagraph capture sizes,
@@ -741,8 +736,11 @@ class Worker(WorkerBase):
736
# e.g. for the max-num-batched token size in chunked prefill.
737
compile_sizes = self.vllm_config.compilation_config.compile_sizes
738
warmup_sizes = compile_sizes.copy() if compile_sizes is not None else [] # type: ignore[assignment]
739
+ cg_capture_sizes: list[int] = []
740
741
if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
742
+ cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
743
+ cg_capture_sizes = [] if cg_sizes is None else cg_sizes
744
warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes]
745
746
compile_ranges = self.vllm_config.compilation_config.get_compile_ranges()
@@ -755,23 +753,6 @@ class Worker(WorkerBase):
753
if not any(x in compile_range for x in all_sizes):
754
warmup_sizes.append(compile_range.end)
755
758
- # TODO(LucasWilkinson, akaratza): Remove when MRV1 is deprecated
759
- if (
760
- current_platform.is_rocm()
761
- and not self.use_v2_model_runner
762
- and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
763
- and get_pp_group().is_last_rank
764
- ):
765
- max_num_reqs = min(
766
- self.scheduler_config.max_num_seqs,
767
- self.scheduler_config.max_num_batched_tokens,
768
- )
769
- if (
770
- max_num_reqs not in cg_capture_sizes
771
- and max_num_reqs not in warmup_sizes
772
- ):
773
- warmup_sizes.append(max_num_reqs)
774
-
756
# We skip EPLB here since we don't want to record dummy metrics
757
for size in sorted(warmup_sizes, reverse=True):
758
logger.info("Compile and warming up model for size %d", size)