feat: Offloading Multimodal embedding table to CPU in Chunked Prefill Mode

[katec846]

Optimization for Large Embedding Tables in Multimodal Models

Problem Statement

In multimodal models, large embedding tables generated by the multimodal encoder consume significant GPU memory. Currently, these tables are stored entirely in GPU memory during processing, which can be inefficient and limiting, especially for models processing large batches or long sequences.

Solution

This PR implements a memory optimization strategy that offloads the multimodal model embedding table (mm_embedding_table) to CPU memory and uses a chunked prefetching mechanism during processing. This approach is only available when operating in context chunk mode, allowing for efficient memory usage without compromising performance through strategic prefetching.

Key Changes

Core Implementation

• Added new buffer management system in promptTuningBuffers.h for handling prompt tuning ping-pong buffers.
• Implemented chunk-based prefetching and processing in trtGptModelInflightBatching.cpp.
• Added support for mm_embedding_table offloading in Python interface layers.

File Changes

1. C++ Changes

   • promptTuningBuffers.h: Added new buffer management system with ping-pong buffers for efficient chunk handling.
   • trtGptModelInflightBatching.cpp: Implemented asynchronous prefetching and chunk processing mechanisms.
   • Various other C++ files: Added supporting functionality for buffer management and data transfer.

2. Python Changes

   • model_runner_cpp.py: Added multimodal model embedding table offloading support with pinned memory allocation.
   • multimodal_model_runner.py: Updated model generation process and move the embedding table to pinned memory.
   • utils.py: Added new --mm_embedding_offloading argument with automatic default configuration.

Usage

The feature can be enabled using the new --mm_embedding_offloading argument. When not specified, it defaults to True if using a multimodal model with chunked context. This optimization is particularly useful for large-scale multimodal applications where GPU memory is a constraint.

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1507 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1507 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1124 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1572 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1572 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1174 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1821 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1821 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1347 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1841 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1841 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1363 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1863 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1863 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1381 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #1972 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #1972 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1454 completed with status: 'SUCCESS'

[tensorrt-cicd]

PR_Github #2790 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1973 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #2800 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2800 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1978 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #2815 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2815 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1989 completed with status: 'FAILURE'

[katec846]

/bot run

[tensorrt-cicd]

PR_Github #2871 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2871 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2027 completed with status: 'SUCCESS'

[symphonylyh]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #2905 [ reuse-pipeline ] triggered by Bot

[tensorrt-cicd]

PR_Github #2905 [ reuse-pipeline ] completed with state SUCCESS
Reusing PR_Github #2871 for commit 8348f4a