Fix OOM by offloading multimodal features to CPU after embedding (#16018)

2025-12-31 23:02:34 +08:00
parent c1dfbc777b
commit abdf65d4f3
1 changed files with 15 additions and 2 deletions
@@ -1077,8 +1077,21 @@ def general_mm_embed_routine(
            # add for qwen3_vl deepstack
            if use_deepstack:
                kwargs["input_deepstack_embeds"] = other_info["input_deepstack_embeds"]
-            # once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models
+            # Offload GPU features to CPU instead of discarding them to balance memory
-            # just being defensive here
+            # efficiency and data persistence.
            # In chunked-prefill, a request is processed across multiple batches, and
            # the original multimodal data must remain accessible until the entire
            # prefill phase is complete. Since the multimodal embedding cache is
            # best-effort, offloading to CPU ensures we have a reliable fallback
            # if a cache miss occurs in subsequent chunks, while still freeing up
            # critical GPU memory.
            if mm_inputs_list:
                for mm_input_obj in mm_inputs_list:
                    if mm_input_obj and hasattr(mm_input_obj, "mm_items"):
                        for mm_item in mm_input_obj.mm_items:
                            feature = getattr(mm_item, "feature", None)
                            if isinstance(feature, torch.Tensor) and feature.is_cuda:
                                mm_item.feature = feature.to("cpu", non_blocking=True)
            forward_batch.mm_inputs = None
        else:
            input_embeds = embed_tokens(input_ids)