Skip to content

GitLab

Unverified Commit 13cc7f53 authored by Lucas Wilkinson's avatar Lucas Wilkinson Committed by GitHub
Browse files

[BugFix] Fix DBO hang (#25625) 


Signed-off-by: default avatarLucas Wilkinson <lwilkins@redhat.com>
parent 916bd920
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
vllm/v1/worker/gpu_ubatch_wrapper.py
View file @ 13cc7f53
...@@ -366,7 +366,8 @@ class UBatchWrapper: ...@@ -366,7 +366,8 @@ class UBatchWrapper:
cudagraph_runtime_mode=CUDAGraphMode.NONE) cudagraph_runtime_mode=CUDAGraphMode.NONE)
with self.sm_control: with self.sm_control:
return self._capture_ubatches(ubatch_metadata, self.model) return self._capture_ubatches(ubatch_metadata, self.model)
elif num_tokens in self.cudagraphs: elif num_tokens in self.cudagraphs \
and cudagraph_runtime_mode is CUDAGraphMode.FULL:
cudagraph_metadata = self.cudagraphs[num_tokens] cudagraph_metadata = self.cudagraphs[num_tokens]
cudagraph_metadata.cudagraph.replay() cudagraph_metadata.cudagraph.replay()
return cudagraph_metadata.outputs return cudagraph_metadata.outputs
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment