[VLM] Fix CUDA IPC OOM (#16118)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
2026-01-07 11:30:35 +08:00
parent 534ac384db
commit 53846746bf
2 changed files with 7 additions and 0 deletions
@@ -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