[VLM] Fix CUDA IPC OOM (#16118)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-01-07 11:30:35 +08:00
committed by GitHub
parent 534ac384db
commit 53846746bf
2 changed files with 7 additions and 0 deletions

View File

@@ -1574,6 +1574,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mm_item.feature = pixel_values.reconstruct_on_target_device(
torch.cuda.current_device()
)
# The reference by CudaIpcTensorTransportProxy was cut off,
# proactively delete to avoid slow gc.
del pixel_values
self.multimodal_inputs = multimodal_inputs
self.token_type_ids = token_type_ids_tensor
self.seq_lens_sum = sum(seq_lens)

View File

@@ -880,6 +880,8 @@ class BaseMultimodalProcessor(ABC):
info_data=item.feature,
sync_buffer_meta=sync_flag,
)
elif not self.server_args.keep_mm_feature_on_device:
item.feature = item.feature.cpu()
elif (
isinstance(item.precomputed_embeddings, torch.Tensor)
and item.precomputed_embeddings.is_cuda
@@ -900,5 +902,7 @@ class BaseMultimodalProcessor(ABC):
info_data=item.precomputed_embeddings,
sync_buffer_meta=sync_flag,
)
elif not self.server_args.keep_mm_feature_on_device:
item.precomputed_embeddings = item.precomputed_embeddings.cpu()
return all_collected_items, input_ids, ret