diff --git a/docker/Makefile b/docker/Makefile index 658aebfbfcc9..838473ea38bb 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -121,6 +121,7 @@ else endif SOURCE_DIR ?= $(shell readlink -f ..) CODE_DIR ?= /code/tensorrt_llm +EXTRA_VOLUMES ?= CCACHE_DIR ?= ${CODE_DIR}/cpp/.ccache CONAN_DIR ?= ${CODE_DIR}/cpp/.conan RUN_CMD ?= @@ -138,6 +139,7 @@ endif docker run $(DOCKER_RUN_OPTS) $(DOCKER_RUN_ARGS) \ $(GPU_OPTS) \ --volume $(SOURCE_DIR):$(CODE_DIR) \ + $(EXTRA_VOLUMES) \ --env "CCACHE_DIR=${CCACHE_DIR}" \ --env "CCACHE_BASEDIR=${CODE_DIR}" \ --env "CONAN_HOME=${CONAN_DIR}" \ diff --git a/docker/README.md b/docker/README.md index d986b8c84931..e319de865e6e 100644 --- a/docker/README.md +++ b/docker/README.md @@ -44,6 +44,11 @@ Containers can be started with the local user instead of `root` by appending `LO make -C docker devel_run LOCAL_USER=1 ``` +Extra docker volumes can be mounted in addition to the source code by appending `EXTRA_VOLUMES=` to the run target: +```bash +make -C docker devel_run LOCAL_USER=1 EXTRA_VOLUMES="--volume /pathA:/pathA --volume /pathB:/pathB" +``` + Specific CUDA architectures supported by the `wheel` can be specified WITH `CUDA_ARCHS`: ```bash diff --git a/examples/pytorch/README.md b/examples/pytorch/README.md index 19ed848a4123..40b01449e325 100644 --- a/examples/pytorch/README.md +++ b/examples/pytorch/README.md @@ -35,14 +35,15 @@ python3 quickstart_advanced.py --model_dir nvidia/Nemotron-H-8B-Base-8K --disabl ```bash # default inputs -python3 quickstart_multimodal.py --model_dir Efficient-Large-Model/NVILA-8B --modality image [--use_cuda_graph] +TLLM_MULTIMODAL_DISAGGREGATED=1 python3 quickstart_multimodal.py --model_dir llava-hf/llava-v1.6-mistral-7b-hf --modality image [--use_cuda_graph] # user inputs # supported modes: # (1) N prompt, N media (N requests are in-flight batched) # (2) 1 prompt, N media # Note: media should be either image or video. Mixing image and video is not supported. -python3 quickstart_multimodal.py --model_dir Efficient-Large-Model/NVILA-8B --modality video --prompt "Tell me what you see in the video briefly." "Describe the scene in the video briefly." --media "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4" "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/world.mp4" --max_tokens 128 [--use_cuda_graph] +TLLM_MULTIMODAL_DISAGGREGATED=0 python3 quickstart_multimodal.py --model_dir llava-hf/llava-v1.6-mistral-7b-hf --modality video --prompt "Tell me what you see in the video briefly." "Describe the scene in the video briefly." --media "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4" "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/world.mp4" --max_tokens 64 [--use_cuda_graph] [--enable_overlap_scheduler] +# use TLLM_MULTIMODAL_DISAGGREGATED to control vision+LLM in a single forward pass (0) or in separate forward pass (1) ``` ### Supported Models diff --git a/examples/pytorch/quickstart_multimodal.py b/examples/pytorch/quickstart_multimodal.py index fec48d1aebd1..edce69f320e4 100644 --- a/examples/pytorch/quickstart_multimodal.py +++ b/examples/pytorch/quickstart_multimodal.py @@ -1,7 +1,8 @@ import argparse import json import os -from typing import Any, Dict, List +from functools import partial +from typing import Any, Callable, Dict, List from quickstart_advanced import add_llm_args, setup_llm @@ -28,24 +29,24 @@ ] -def prepare_multimodal_inputs(model_dir: str, - model_type: str, - modality: str, - prompts: List[str], - media: List[str], - image_data_format: str = "pt", - num_frames: int = 8) -> List[Dict[str, Any]]: +def prepare_multimodal_inputs( + model_dir: str, + modality: str, + prompts: List[str], + media: List[str], + input_formatter: Callable, + mm_loader: Callable, + data_format: str = "pt", # Options: "pt" or "pil" + device: str = "cuda") -> List[Dict[str, Any]]: inputs = [] - if modality == "image": - inputs = default_image_loader(prompts, media, image_data_format) - elif modality == "video": - inputs = default_video_loader(prompts, media, image_data_format, - num_frames) + if modality in ["image", "video"]: + assert mm_loader, "multimodal data loader is required for image/video modality" + inputs = mm_loader(prompts, media, data_format, device) else: raise ValueError(f"Unsupported modality: {modality}") - inputs = INPUT_FORMATTER_MAP[model_type](model_dir, inputs) + inputs = input_formatter(model_dir, inputs) return inputs @@ -95,17 +96,23 @@ def main(): llm, sampling_params = setup_llm(args) - image_format = "pt" # ["pt", "pil"] - if args.model_type is not None: - model_type = args.model_type - else: - model_type = json.load( - open(os.path.join(llm._hf_model_dir, 'config.json')))['model_type'] + # feel free to override the default formatter and loaders based on your applications + model_type = args.model_type if args.model_type else json.load( + open(os.path.join(llm._hf_model_dir, 'config.json')))['model_type'] assert model_type in INPUT_FORMATTER_MAP, f"Unsupported model_type: {model_type}" - - inputs = prepare_multimodal_inputs(args.model_dir, model_type, - args.modality, args.prompt, args.media, - image_format, args.num_frames) + input_formatter = INPUT_FORMATTER_MAP[model_type] + mm_loader = None + if args.modality == "image": + mm_loader = default_image_loader + elif args.modality == "video": + mm_loader = partial(default_video_loader, num_frames=args.num_frames) + data_format = "pt" # ["pt", "pil"] + media_input_device = "cpu" + + inputs = prepare_multimodal_inputs(args.model_dir, args.modality, + args.prompt, args.media, input_formatter, + mm_loader, data_format, + media_input_device) outputs = llm.generate(inputs, sampling_params) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 5fa9cc7f3cda..3bc38af118e7 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -53,12 +53,6 @@ class AttentionMetadata: # Whether CUDA graph is enabled. is_cuda_graph: bool = field(default=False, repr=False) - # The length of each sequence in the batch for query. - # The shape is (batch_size), and located on CPU memory. - # For sub metadata of cross attention, it's automatically - # initialized to seq_lens of parent metadata. - seq_lens: Optional[torch.Tensor] # Implemented using property - # The number of context-phase sequences in the batch. num_contexts: int # Implemented using property @@ -71,6 +65,12 @@ class AttentionMetadata: # The parameters for the KV cache. kv_cache_params: Optional[KVCacheParams] = None + # The length of each sequence in the batch for query. + # The shape is (batch_size), and located on CPU memory. + # For sub metadata of cross attention, it's automatically + # initialized to seq_lens of parent metadata. + seq_lens: Optional[torch.Tensor] # Implemented using property + # The length of each sequence in the batch for key and value. # The shape is (batch_size), and located on CPU memory. # It defaults to seq_lens if not set. @@ -110,6 +110,8 @@ class AttentionMetadata: # For generation-phase sequence, the value is the token number of its context phase. # The shape is (batch_size) if provided. prompt_lens: Optional[List[int]] = None + # The "original" prompt length of each sequence in the batch. Useful in cases where the prompt tokens can be updated thus become different from the original user input prompt, e.g. multimodal. + orig_prompt_lens: Optional[List[int]] = None # These fields indicate whether the runtime can use various features. # The kernels may or may not have different behaviors when these @@ -148,6 +150,35 @@ def on_update(self): elif self._seq_lens is not None: self._num_tokens = self._seq_lens.sum().item() + def on_update_gpu(self, key="all"): + ''' + Update underlying GPU buffers when seq_lens or seq_lens_kv is updated. + ''' + if key in ["seq_lens", "all"]: + # The model executor sets seq_lens to None initially. + if self._seq_lens is not None: + self._seq_lens = self._seq_lens.pin_memory() + + if self.is_cuda_graph and self._seq_lens_cuda is not None: + # Very important: do not reallocate if we are using CUDA graphs. + # This copy is safe because the batch size is guaranteed to not + # change in the CUDA graph case. The seqlens can change if we + # are doing spec decode. + self._seq_lens_cuda.copy_(self._seq_lens, non_blocking=True) + else: + self._seq_lens_cuda = self._seq_lens.cuda(non_blocking=True) + + if self.has_cross_sub_metadata: + self.cross._seq_lens = self._seq_lens + self.cross._seq_lens_cuda = self._seq_lens_cuda + + if key in ["seq_lens_kv", "all"]: + # The model executor sets seqlens to None initially. + if self._seq_lens_kv is not None: + self._seq_lens_kv = self._seq_lens_kv.pin_memory() + self._seq_lens_kv_cuda = self._seq_lens_kv.cuda( + non_blocking=True) + @property def seq_lens(self) -> Optional[torch.Tensor]: return self._seq_lens @@ -158,23 +189,26 @@ def seq_lens(self, value: Optional[torch.Tensor]): value = value if value is not AttentionMetadata.seq_lens else None self._seq_lens = value self.on_update() + self.on_update_gpu("seq_lens") - # The model executor sets seq_lens to None initially. - if self._seq_lens is not None: - self._seq_lens = self._seq_lens.pin_memory() + @property + def seq_lens_cuda(self): + return self._seq_lens_cuda - if self.is_cuda_graph and self._seq_lens_cuda is not None: - # Very important: do not reallocate if we are using CUDA graphs. - # This copy is safe because the batch size is guaranteed to not - # change in the CUDA graph case. The seqlens can change if we - # are doing spec decode. - self._seq_lens_cuda.copy_(self._seq_lens, non_blocking=True) - else: - self._seq_lens_cuda = self._seq_lens.cuda(non_blocking=True) + @property + def seq_lens_kv(self) -> Optional[torch.Tensor]: + return self._seq_lens_kv if self._seq_lens_kv is not None else self._seq_lens - if self.has_cross_sub_metadata: - self.cross._seq_lens = self._seq_lens - self.cross._seq_lens_cuda = self._seq_lens_cuda + @seq_lens_kv.setter + def seq_lens_kv(self, value: Optional[torch.Tensor]): + value = value if value is not AttentionMetadata.seq_lens_kv else None + self._seq_lens_kv = value + self.on_update() + self.on_update_gpu("seq_lens_kv") + + @property + def seq_lens_kv_cuda(self): + return self._seq_lens_kv_cuda if self._seq_lens_kv_cuda is not None else self._seq_lens_cuda @property def num_contexts(self) -> int: @@ -196,28 +230,6 @@ def num_generations(self, value: int): self._num_generations = value self.on_update() - @property - def seq_lens_cuda(self): - return self._seq_lens_cuda - - @property - def seq_lens_kv(self) -> Optional[torch.Tensor]: - return self._seq_lens_kv if self._seq_lens_kv is not None else self._seq_lens - - @seq_lens_kv.setter - def seq_lens_kv(self, value: Optional[torch.Tensor]): - value = value if value is not AttentionMetadata.seq_lens_kv else None - self._seq_lens_kv = value - self.on_update() - # The model executor sets seqlens to None initially. - if self._seq_lens_kv is not None: - self._seq_lens_kv = self._seq_lens_kv.pin_memory() - self._seq_lens_kv_cuda = self._seq_lens_kv.cuda(non_blocking=True) - - @property - def seq_lens_kv_cuda(self): - return self._seq_lens_kv_cuda if self._seq_lens_kv_cuda is not None else self._seq_lens_cuda - @property def context_lens(self) -> torch.Tensor: """ diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index ef86ed06844c..d9a4dd702936 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -135,6 +135,8 @@ def from_pretrained(cls, model_dir = Path( transformers.utils.hub.cached_file(checkpoint_dir, 'config.json')).parent + pretrained_config.checkpoint_dir = model_dir + quant_config = QuantConfig() layer_quant_config = None # quantized ckpt in modelopt format diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py index cc64b6374c49..ed12dd5b813b 100644 --- a/tensorrt_llm/_torch/models/modeling_llava_next.py +++ b/tensorrt_llm/_torch/models/modeling_llava_next.py @@ -20,14 +20,52 @@ from ..model_config import ModelConfig from .modeling_auto import AutoModelForCausalLM from .modeling_clip import CLIPVisionModel -from .modeling_multimodal_utils import fuse_input_embeds +from .modeling_multimodal_utils import fuse_input_embeds, prepare_multimodal_ifb from .modeling_utils import ModelConfig, filter_weights, register_auto_model +DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1' + class LlavaNextInputProcessor(InputProcessor): def __init__(self, model_path, model_config, tokenizer): self.tokenizer = tokenizer + + if DISAGG: + self.mm_encoder = LlavaNextVisionModel(model_path, model_config) + + @torch.inference_mode() + def __call__( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + # text input - tokenization + text_prompt = inputs.get("prompt") + input_ids = self.tokenizer(text_prompt, + return_tensors="pt").input_ids[0] + + # multimodal input - either mm encoder or plain concatenation the raw media tensors + mm_data = inputs.get("multi_modal_data", None) + assert mm_data is not None and 'image' in mm_data, "Multimodal inputs must be a dictionary with 'image' field" + if DISAGG: + # disaggregated batching mode: multimodal encoder runs as a per-request input processor, always BS=1 + # returning expanded input_ids & multimodal embedding + mm_embeds, mm_embed_lengths = self.mm_encoder.forward( + mm_data['image']) + fused_input_ids, mm_embeds = self.mm_encoder._postprocess( + input_ids.to(mm_embeds.device), mm_embeds, mm_embed_lengths) + return fused_input_ids.tolist(), {"mm_embedding": mm_embeds} + else: + # inflight batching mode: multimodal encoder runs as part of the model forward + # returning raw input_ids & raw media + # within each request, input images/frames are usually of the same size; across requests, the sizes may vary + # InputProcessor is per-request call, so we can stack all images/frames as [num_images, C, H, W]. still on CPU. + mm_data = torch.stack(mm_data['image']) + return input_ids.tolist(), {"mm_embedding": mm_data} + + +class LlavaNextVisionModel: + + def __init__(self, model_path, model_config): self.processor = AutoProcessor.from_pretrained(model_path, use_fast=True) self.model_config = model_config @@ -67,6 +105,8 @@ def __init__(self, model_path, model_config, tokenizer): # Use HF multi-modal projector self.mm_projector = hf_mm_projector + self.max_batch_size = 8 + @nvtx_range("[Vision] preprocess") def _preprocess(self, images): return [ @@ -91,7 +131,7 @@ def _process(self, pixel_values): return image_features.reshape(-1, image_features.shape[-1]) @nvtx_range("[Vision] postprocess") - def _postprocess(self, input_ids, mm_features): + def _postprocess(self, input_ids, mm_features, mm_feature_lengths): # Define model specific variables here before shared logic mm_tokens = torch.tensor([self.model_config.image_token_index ]).to(input_ids.device) @@ -100,6 +140,9 @@ def _postprocess(self, input_ids, mm_features): start_len = end_len = 0 # for llava, need not append start/end token around each image token # End model specific variables + mm_features = mm_features.reshape(len(mm_feature_lengths), -1, + mm_features.shape[-1]) + ## find mm token positions in input_ids mm_token_positions = torch.where(torch.isin(input_ids, mm_tokens))[0] num_medias = num_mm_tokens = len(mm_token_positions) @@ -158,7 +201,7 @@ def _postprocess(self, input_ids, mm_features): assert mm_split_idx == len( mm_ids_splits), "All mm_ids_splits should be consumed" - ## concat text & mm input_ids, wrap mm feature in prompt tuning config + ## concat text & mm input_ids fused_input_ids = torch.cat(input_ids_splits).to( device=input_ids.device) fused_length = len(input_ids) + mm_total_length + num_frames * ( @@ -169,26 +212,24 @@ def _postprocess(self, input_ids, mm_features): # [num_frames, feature_length, hidden_dim] -> [num_frames * feature_length, hidden_dim] mm_features = mm_features.view(-1, mm_features.shape[-1]) - return fused_input_ids, mm_features + return fused_input_ids.to(torch.int32), mm_features @torch.inference_mode() - def __call__( - self, inputs: TextPrompt, sampling_params: SamplingParams - ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: - text_prompt, mm_data = inputs.get("prompt"), inputs.get( - "multi_modal_data", {}) - assert 'image' in mm_data + def forward(self, + images: List[torch.Tensor]) -> Tuple[torch.Tensor, List[int]]: + + # List[raw image tensors (CPU)] for N requests, could have different sizes --> CPU compute-vision ops e.g. resize, crop, etc. --> List[preprocessed tensors (GPU)] for N requests + mm_preprocessed = self._preprocess(images) + mm_preprocessed_lengths = [t.size(0) for t in mm_preprocessed] - input_ids = self.tokenizer( - text_prompt, return_tensors="pt").input_ids[0].to(self.device) + # concatenated preprocessed tensors (GPU) --> GPU model forward --> mm embedding tensors (GPU), [total_num_media_tokens, hidden_dim] + mm_embeds = self._process(torch.cat(mm_preprocessed)) + mm_embed_lengths = [ + mm_embeds.shape[0] // sum(mm_preprocessed_lengths) * l + for l in mm_preprocessed_lengths + ] # calculate num_media_tokens for each request, which is proportional to the preprocessed length - mm_tensor = self._preprocess(mm_data['image']) - mm_features = torch.stack( - [self._process(tensor) for tensor in mm_tensor]) - fused_input_ids, mm_features = self._postprocess(input_ids, mm_features) - return fused_input_ids.to(torch.int32).tolist(), { - "mm_embedding": mm_features - } + return mm_embeds, mm_embed_lengths @register_auto_model("LlavaNextForConditionalGeneration") @@ -219,8 +260,13 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, torch.float16) logger.info(f"{self.dtype=} {self.model_dtype=}") + if not DISAGG: + self.mm_encoder = LlavaNextVisionModel(config.checkpoint_dir, + config) + self.mm_tokens = torch.tensor([config.image_token_index], + device='cuda') + self.post_config() - self.is_loaded = True def load_weights(self, weights): @@ -247,13 +293,39 @@ def forward( num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations logger.debug(f"{num_context_requests=}, {num_generation_requests=}") - mm_embed = kwargs.get("multi_modal_data", []) - assert mm_embed == [] or len( - mm_embed - ) == num_context_requests, "Number of multimodal features (if provided) should be equal to number of context requests" + mm_data = kwargs.get("multi_modal_data", []) + if DISAGG: + # disaggregated batching mode: multimodal context phase is decoupled from LLM forward. mm_data is processed multimodal embedding + logger.warning( + "No multimodal encoder found in model definition. You might be in disaggregated inference mode. Skipping multimodal processing. It's expected that a expanded input_ids and mm_embed are provided as inputs." + ) + mm_embed = mm_data + assert mm_embed == [] or len( + mm_embed + ) == num_context_requests, f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch." + input_ids, inputs_embeds = fuse_input_embeds( + self.llm.model.embed_tokens, input_ids, mm_embed) + else: + # inflight batching mode: multimodal context phase is fused in LLM forward. mm_data is raw media tensors + mm_embeds, mm_embed_lengths = None, None + if len(mm_data) > 0: + assert len( + mm_data + ) == num_context_requests, f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch." + mm_embeds, mm_embed_lengths = self.mm_encoder.forward(mm_data) + + attn_metadata, input_ids, position_ids, inputs_embeds = prepare_multimodal_ifb( + self.llm.model.embed_tokens, + attn_metadata, + input_ids, + position_ids, + self.mm_tokens, + mm_embeds, + mm_embed_lengths, + prepare_position_ids=self.llm.model.layers[0].self_attn. + apply_rotary_emb + ) # position_ids is only relevant when RoPE needs to be explicitly applied outside the fused attention op. - input_ids, inputs_embeds = fuse_input_embeds( - self.llm.model.embed_tokens, input_ids, mm_embed) logits = self.llm.forward(attn_metadata, input_ids, position_ids, inputs_embeds, return_context_logits) return logits diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index e97dcc6e07b7..b304e76fba22 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @@ -27,7 +27,158 @@ from tensorrt_llm._torch.modules.embedding import Embedding +from ..._utils import nvtx_range + + +@nvtx_range("prepare_multimodal_ifb") +def prepare_multimodal_ifb(embedding_layer, + attn_metadata, + input_ids, + position_ids, + mm_tokens=None, + mm_embeds=None, + mm_embed_lengths=None, + prepare_position_ids=False): + """ + Prepare for inflight batching LLM forward in multimodal mode. When Encoder + LLM is in the same forward, this function should be called between the encoder forward and the LLM forward. + + Challenges: + in multimodal, the actual seq_lens (text + media) is only known AFTER the encoder forward (there are VLM models that vision seq_len is deterministic and independent of the media size; but a generic solution shouldn't base on such assumption). We need a way to update from the text-only seq_lens to the actual (text + media) seq_lens, which is not easy. + - approach 1: update the input_ids field in llm request, such as using a dummy input_ids that matches the real seq_lens. + Pros: worry-free about setting the correct lengths during scheduling and metadata init. + Cons: (i) inside forward, we have no handle to the request object (ii) input_ids field in LlmRequest class is immutable (iii) wasted space to store the dummy input_ids + - approach 2: update the actual lengths in model forward. + Pros: (i) model-specific handling, doesn't affect other models (ii) no wasted space. + Cons: (i) a bit hacky, modifying the metadata is error-prone (ii) in theory, the IFB scheduling may over-shoot because it's based on the text-only seq_lens + The current implementation is approach 2. + The ideal approach is to add a new seq_len field in LlmRequest. Let this field be mutable & all seq_len getter routes to this field rather than calculating len(input_ids). + TODO: handle the hybrid case of both text-only & multimodal requests + """ + + num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations + + input_embeds = None + + # multimodal context requests + if num_context_requests > 0 and mm_embeds is not None: # skip dummy requests + ## Fuse input embeddings + # 1. remove mm tokens from input_ids + raw_ctx_tokens, raw_gen_tokens = input_ids[:attn_metadata. + num_ctx_tokens], input_ids[ + attn_metadata. + num_ctx_tokens:] + raw_text_mask = ~torch.isin(raw_ctx_tokens, mm_tokens) + input_ids = torch.cat([raw_ctx_tokens[raw_text_mask], raw_gen_tokens]) + input_embeds = torch.empty(input_ids.shape[0] + mm_embeds.shape[0], + mm_embeds.shape[-1], + device=mm_embeds.device, + dtype=mm_embeds.dtype) + fused_text_mask = torch.full((input_embeds.shape[0], ), False) + if raw_gen_tokens.shape[0] > 0: + fused_text_mask[-raw_gen_tokens.shape[0]:] = True + + # 2. calculate the text token indices in the fused input_embeds + raw_text_masks = list( + raw_text_mask.split(attn_metadata.context_lens.tolist())) + mm_embed_splits = list(mm_embeds.split(mm_embed_lengths)) + start_idx, last_start_idx = 0, 0 + fused_lengths = [] + for text_mask, mm_embed in zip(raw_text_masks, + mm_embed_splits): # per request + mm_positions = torch.where(text_mask == False)[0].tolist() + num_medias = len(mm_positions) + media_length = len(mm_embed) // num_medias + + # Diagram for 2 media, length 3 & length 4 each. After processing the 1st media: + # index: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 + # raw tokens: T T T M T T M T T (T - text, M - media) + # mm_positions: ^ ^ + # ^pos + # ^last_pos + # fused_text mask: T T T F F F T T F F F F T T (T-True, F-False) + # ^start_idx + last_pos = 0 + for pos in mm_positions: # per media in each request + text_length = pos - last_pos + fused_text_mask[start_idx:start_idx + text_length] = True + start_idx += text_length + media_length + last_pos = pos + 1 + + if last_pos < len(text_mask): + # text between last media token & the end + text_length = len(text_mask) - last_pos + fused_text_mask[start_idx:start_idx + text_length] = True + start_idx += text_length + + fused_lengths.append(start_idx - last_start_idx) + last_start_idx = start_idx + + # 3. fuse embeddings + input_embeds[fused_text_mask] = embedding_layer(input_ids) + input_embeds[~fused_text_mask] = mm_embeds + input_ids = None # use input_embeds mode for multimodal + + ## Update metadata + # 1. Update attn_metadata for the following LLM forward. This can be done in-place in this forward(). NOTE: we cannot simply do attn_metadata.seq_lens[:num_context_requests] = xx, because this won't trigger the setter calls (see interface.py) thus the underlying seq_lens_cuda buffer won't get updated. We MUST do an explicit setter. + attn_metadata.prompt_lens[:num_context_requests] = fused_lengths + attn_metadata.seq_lens[:num_context_requests] = torch.tensor( + fused_lengths, dtype=torch.int32) + attn_metadata.seq_lens_kv[:num_context_requests] = torch.tensor( + fused_lengths, dtype=torch.int32) + attn_metadata.on_update_gpu() + # 2. Update request data for the generation phase. It's not straightforward to propagate the updated info back to request fields. A viable solution is to temporarily store the info in KVCacheManager, and leverage its update_resources() -- which is invoked after the forward() step -- to update the request data. see resource_manager.py. + attn_metadata.kv_cache_manager.extra_info_for_update_resources[ + 'updated_seq_lens_ctx'] = attn_metadata.seq_lens[: + num_context_requests].tolist( + ) + # 3. Update KV cache size allocation + # TODO: is there a better way to add multiple tokens or extend the sequence in one go? + for i, req in enumerate( + attn_metadata.request_ids[:num_context_requests]): + for _ in range(fused_lengths[i] - + attn_metadata.orig_prompt_lens[i]): + attn_metadata.kv_cache_manager.impl.add_token(req) + + # multimodal generation requests + if num_generation_requests > 0: + # 4. Update KV cache length for generation requests + # num_cached_tokens_per_seq is counting based on the original prompt length (because llmRequest class only stores the original text input ids). Number of generated tokens is the delta between the two. + attn_metadata.kv_cache_params.num_cached_tokens_per_seq[ + num_context_requests:] = list( + map( + lambda x, y, z: x + y - z, + attn_metadata.prompt_lens[num_context_requests:], + attn_metadata.kv_cache_params. + num_cached_tokens_per_seq[num_context_requests:], + attn_metadata.orig_prompt_lens[num_context_requests:])) + # TODO: (1) could just save the delta between original & current prompt lens to save a subtract op exposed in step latency (2) more performant impl than list map + + attn_metadata.prepare() # must update internal buffers + + if prepare_position_ids: + position_ids_list = [] + for i in range(num_context_requests + num_generation_requests): + if i < num_context_requests: + position_ids_list.append( + torch.arange(start=0, + end=attn_metadata.seq_lens[i], + dtype=torch.int, + device='cuda')) + else: + position_ids_list.append( + torch.tensor([ + attn_metadata.kv_cache_params. + num_cached_tokens_per_seq[i] + ], + dtype=torch.int, + device='cuda')) + if len(position_ids_list) > 0: + position_ids = torch.cat(position_ids_list).unsqueeze(0) + + return attn_metadata, input_ids, position_ids, input_embeds + +@nvtx_range("fuse_input_embeds") def fuse_input_embeds( embedding_layer: Embedding, input_ids: torch.LongTensor, diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index 3515dab60f0e..b1e2df24ca6c 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -19,7 +19,7 @@ class PyTorchConfig: """ # Extra resource managers to use in addition to the KV cache manager. - # Each manager's prepare_resources method is called before the forward pass, + # Each manager's prepare_resources() is called before the forward pass, # and update_resources() is called after the pass finishes. free_resources() # is called when a request finishes. # The KV cache manager is guaranteed to be invoked after all of these extra diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8dcfe25dfbb0..0ad3b2741e05 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1025,6 +1025,7 @@ def _prepare_tp_inputs( input_ids = [] sequence_lengths = [] prompt_lengths = [] + orig_prompt_lengths = [] request_ids = [] gather_ids = [] position_ids = [] @@ -1050,6 +1051,7 @@ def _prepare_tp_inputs( gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) prompt_lengths.append(len(prompt_tokens)) + orig_prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = request.context_current_position num_cached_tokens_per_seq.append(past_seen_token_num) multimodal_embedding = request.multimodal_embedding() @@ -1108,6 +1110,7 @@ def _prepare_tp_inputs( position_ids.append(past_seen_token_num) draft_lens.append(num_draft_tokens) prompt_lengths.append(num_draft_tokens + 1) + orig_prompt_lengths.append(num_draft_tokens + 1) # draft tokens input_ids.extend(request.py_draft_tokens) gather_ids.extend( @@ -1148,7 +1151,7 @@ def _prepare_tp_inputs( num_cached_tokens_per_seq.append(past_seen_token_num + self.max_draft_len + 1) prompt_lengths.append(request.py_prompt_len) - + orig_prompt_lengths.append(request.py_orig_prompt_len) request_ids.append(request.py_request_id) sequence_lengths.extend([1] * len(generation_requests)) @@ -1176,6 +1179,7 @@ def _prepare_tp_inputs( position_ids.append(past_seen_token_num) num_cached_tokens_per_seq.append(past_seen_token_num) prompt_lengths.append(request.py_prompt_len) + orig_prompt_lengths.append(request.py_orig_prompt_len) draft_lens.append(0) request.py_batch_idx = batch_idx @@ -1264,6 +1268,7 @@ def _prepare_tp_inputs( attn_metadata.request_ids = request_ids attn_metadata.prompt_lens = prompt_lengths + attn_metadata.orig_prompt_lens = orig_prompt_lengths attn_metadata.num_contexts = len(scheduled_requests.context_requests) if self.is_spec_decode and self.spec_config.spec_dec_mode.extend_ctx( self.attn_backend): @@ -1476,6 +1481,7 @@ def _prepare_star_attention_inputs(self, sequence_lengths = [] input_ids = [] prompt_lengths = [] + orig_prompt_lengths = [] request_ids = [] gather_ids = [] position_ids = [] @@ -1486,7 +1492,7 @@ def _prepare_star_attention_inputs(self, for request in scheduled_requests.context_requests: request_ids.append(request.py_request_id) prompt_lengths.append(request.py_prompt_len) - + orig_prompt_lengths.append(request.py_orig_prompt_len) ctx_iter = request.ctx_iters ctx_blocks = request.ctx_blocks ctx_position_blocks = request.ctx_position_blocks @@ -1599,7 +1605,7 @@ def _prepare_star_attention_inputs(self, for request in generation_requests: request_ids.append(request.py_request_id) prompt_lengths.append(request.py_prompt_len) - + orig_prompt_lengths.append(request.py_orig_prompt_len) input_token_id = request.get_token(0, request.get_num_tokens(0) - 1) input_ids.append(input_token_id) gather_ids.append(len(input_ids) - 1) @@ -1667,6 +1673,7 @@ def _prepare_star_attention_inputs(self, attn_metadata.request_ids = request_ids attn_metadata.prompt_lens = prompt_lengths + attn_metadata.orig_prompt_lens = orig_prompt_lengths attn_metadata.num_contexts = num_contexts attn_metadata.num_queries = num_queries diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 1a9da786dba0..ec2e22f06b7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -235,6 +235,10 @@ def __init__( self.max_blocks_per_seq = self.impl.max_blocks_per_seq self.enable_block_reuse = kv_cache_config.enable_block_reuse + # store extra info saved during model forward() that may be needed by update_resources() + # e.g. updated seq_lens in multimodal cases + self.extra_info_for_update_resources = {} + def shutdown(self): self.impl.release_pools() @@ -353,6 +357,18 @@ def add_dummy_requests( return requests def update_resources(self, scheduled_batch: ScheduledRequests): + # update seq_lens after multimodal context phase + updated_seq_lens_ctx = self.extra_info_for_update_resources.get( + 'updated_seq_lens_ctx', None) + if len(scheduled_batch.context_requests + ) > 0 and updated_seq_lens_ctx is not None: + assert len(updated_seq_lens_ctx) == len( + scheduled_batch.context_requests + ), f"updated_seq_lens_ctx length {len(updated_seq_lens_ctx)} != context_requests length {len(scheduled_batch.context_requests)}" + for i, request in enumerate(scheduled_batch.context_requests): + request.py_prompt_len = updated_seq_lens_ctx[i] + self.extra_info_for_update_resources = {} # reset + # rewind kv cache for request in scheduled_batch.generation_requests: if request.state != LlmRequestState.GENERATION_COMPLETE: diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 0b9be7cfb48c..6047605bf0bd 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -368,10 +368,7 @@ def _enqueue_request(self, request: GenerationRequest) -> int: prompt_token_ids = copy.deepcopy(request.prompt_token_ids) prompt_tuning_config = None - multimodal_embedding = None mrope_config = None - if request.multimodal_embedding is not None: - multimodal_embedding = request.multimodal_embedding if request.prompt_adapter_request is not None: self._load_prompt_adapter(request.prompt_adapter_request) uid = str(request.prompt_adapter_request.adapter_id) @@ -450,7 +447,8 @@ def _deduce_max_tokens(request: GenerationRequest, embedding_bias=request.sampling_params.embedding_bias, lora_config=lora_config, prompt_tuning_config=prompt_tuning_config, - multimodal_embedding=multimodal_embedding, + multimodal_embedding=request.multimodal_embedding + if request.multimodal_embedding is not None else None, mrope_config=mrope_config, logits_post_processor_name=( tllm.Request.BATCHED_POST_PROCESSOR_NAME diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 6cd97a91b721..ef3de8b1939a 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -276,7 +276,8 @@ def apply_template(prompt, multimodal_data): def default_image_loader(prompts: List[str], images: Union[List[List[str]], List[str]], - image_data_format: str = "pt"): + image_data_format: str = "pt", + device: str = "cuda"): if len(images) > len(prompts) and len(prompts) == 1: # 1 prompt + N media images = [images] @@ -285,10 +286,10 @@ def default_image_loader(prompts: List[str], "prompt": prompt, "multi_modal_data": { "image": [ - load_image(i, format=image_data_format, device="cuda") + load_image(i, format=image_data_format, device=device) for i in image ] if isinstance(image, list) else - [load_image(image, format=image_data_format, device="cuda")] + [load_image(image, format=image_data_format, device=device)] } } for prompt, image in zip(prompts, images)] return inputs @@ -297,6 +298,7 @@ def default_image_loader(prompts: List[str], def default_video_loader(prompts: List[str], videos: Union[List[List[str]], List[str]], image_data_format: str = "pt", + device: str = "cuda", num_frames: int = 8): if len(videos) > len(prompts) and len(prompts) == 1: # 1 prompt + N media @@ -307,11 +309,11 @@ def default_video_loader(prompts: List[str], "multi_modal_data": { "video": [ load_video( - i, num_frames, format=image_data_format, device="cuda") + i, num_frames, format=image_data_format, device=device) for i in video ] if isinstance(video, list) else [ load_video( - video, num_frames, format=image_data_format, device="cuda") + video, num_frames, format=image_data_format, device=device) ] } } for prompt, video in zip(prompts, videos)]