diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 301ef4bb7..28677efea 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -128,7 +128,6 @@ class ImageInputs: image_hashes: Optional[list] = None image_sizes: Optional[list] = None image_offsets: Optional[list] = None - image_pad_len: Optional[list] = None pad_values: Optional[list] = None modalities: Optional[list] = None num_image_tokens: Optional[int] = None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 24a285952..74a7d1fc5 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -111,20 +111,15 @@ class ModelRunner: ) if self.is_multimodal: + logger.info( + "Automatically turn off --chunked-prefill-size and adjust --mem-fraction-static for multimodal models." + ) + server_args.chunked_prefill_size = -1 self.mem_fraction_static *= 0.95 - if self.model_config.hf_config.architectures == [ - "MllamaForConditionalGeneration" - ]: - logger.info("Automatically turn off --chunked-prefill-size for mllama.") - server_args.chunked_prefill_size = -1 # TODO: qwen2-vl does not support radix cache now, set disable_radix_cache=True automatically if self.model_config.hf_config.architectures == [ "Qwen2VLForConditionalGeneration" ]: - logger.info( - "Automatically turn off --chunked-prefill-size and disable radix cache for qwen2-vl." - ) - server_args.chunked_prefill_size = -1 server_args.disable_radix_cache = True # Global vars diff --git a/python/sglang/srt/models/llava.py b/python/sglang/srt/models/llava.py index c8ce9302b..4c62dbb25 100644 --- a/python/sglang/srt/models/llava.py +++ b/python/sglang/srt/models/llava.py @@ -57,7 +57,6 @@ class LlavaBaseForCausalLM(nn.Module): else: image_aspect_ratio = "anyres" offset_list = [] - image_inputs.image_pad_len = [] for image_idx, image_s in enumerate(image_sizes): if len(image_sizes) > 16: # 2x2 pooling with stride 2 @@ -104,7 +103,6 @@ class LlavaBaseForCausalLM(nn.Module): + input_ids[offset + 1 :] ) offset_list.append(offset) - image_inputs.image_pad_len.append(new_image_feature_len) image_inputs.image_offsets = offset_list return input_ids @@ -136,14 +134,6 @@ class LlavaBaseForCausalLM(nn.Module): image_inputs = forward_batch.image_inputs if forward_batch.forward_mode.is_extend(): - # Clamp input ids. This is because the input_ids for the image tokens are - # filled with the hash values of the image for the prefix matching in the radix attention. - # There values are useless because their embeddings will be replaced by vision embeddings anyway. - input_ids.clamp_(min=0, max=self.config.vocab_size - 1) - - # Embed text inputs - input_embeds = self.language_model.model.embed_tokens(input_ids) - # Got List[List[str]] extend it to List[str] # The length of the List should be equal to batch size modalities_list = [] @@ -152,12 +142,18 @@ class LlavaBaseForCausalLM(nn.Module): if im and im.modalities is not None: modalities_list.extend(im.modalities) if im and im.image_offsets: - max_image_offset.append( - np.max(np.array(im.image_offsets) + np.array(im.image_pad_len)) - ) + max_image_offset.append(max(im.image_offsets)) else: max_image_offset.append(-1) + # Clamp input ids. This is because the input_ids for the image tokens are + # filled with the hash values of the image for the prefix matching in the radix attention. + # There values are useless because their embeddings will be replaced by vision embeddings anyway. + input_ids.clamp_(min=0, max=self.config.vocab_size - 1) + + # Embed text inputs + input_embeds = self.language_model.model.embed_tokens(input_ids) + start_positions = positions[forward_batch.extend_start_loc].cpu().numpy() need_vision = start_positions <= np.array(max_image_offset) @@ -354,7 +350,6 @@ class LlavaBaseForCausalLM(nn.Module): # Fill in the placeholder for the image extend_start_loc_cpu = forward_batch.extend_start_loc.cpu().numpy() - extend_seq_lens = forward_batch.extend_seq_lens.cpu().numpy() prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu pt = 0 for i in range(bs): @@ -362,36 +357,18 @@ class LlavaBaseForCausalLM(nn.Module): continue start_idx = extend_start_loc_cpu[i] - seq_len = extend_seq_lens[i] prefix_len = prefix_lens_cpu[i] # Multiple images - for image_idx, image_offset in enumerate( - image_inputs[i].image_offsets - ): - if ( - image_offset + image_inputs[i].image_pad_len[image_idx] - <= prefix_len - ): + for j, image_offset in enumerate(image_inputs[i].image_offsets): + if image_offset < prefix_len: continue - if image_offset >= prefix_len + seq_len: - break - tmp_image_feature = image_features[pt][image_idx] + tmp_image_feature = image_features[pt][j] pad_len = tmp_image_feature.shape[0] - input_offset = image_offset - prefix_len - left_idx = start_idx + input_offset - right_idx = left_idx + pad_len - assert right_idx > start_idx - if input_offset < 0: - left_idx = start_idx - tmp_image_feature = tmp_image_feature[-input_offset:] - if right_idx > start_idx + seq_len: - tmp_image_feature = tmp_image_feature[ - : start_idx + seq_len - right_idx - ] - right_idx = start_idx + seq_len + left_idx = start_idx + (image_offset - prefix_len) + right_idx = start_idx + (image_offset - prefix_len) + pad_len try: input_embeds[left_idx:right_idx] = tmp_image_feature except RuntimeError as e: diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 89fdacfa2..5035810f8 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -39,7 +39,6 @@ suites = { "test_triton_attention_kernels.py", "test_triton_attention_backend.py", "test_update_weights_from_disk.py", - "test_vision_chunked_prefill.py", "test_vision_openai_server.py", "test_session_control.py", ], diff --git a/test/srt/test_vision_chunked_prefill.py b/test/srt/test_vision_chunked_prefill.py deleted file mode 100644 index f7725f17b..000000000 --- a/test/srt/test_vision_chunked_prefill.py +++ /dev/null @@ -1,173 +0,0 @@ -""" -Usage: -python3 -m unittest test_vision_chunked_prefill.TestVisionChunkedPrefill.test_chunked_prefill -""" - -import base64 -import io -import os -import unittest -from concurrent.futures import ThreadPoolExecutor -from typing import Union - -import numpy as np -import requests -from decord import VideoReader, cpu -from PIL import Image - -from sglang.srt.utils import kill_process_tree -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - popen_launch_server, -) - - -class TestVisionChunkedPrefill(unittest.TestCase): - def prepare_video_messages(self, video_path, max_frames_num=8): - vr = VideoReader(video_path, ctx=cpu(0)) - total_frame_num = len(vr) - uniform_sampled_frames = np.linspace( - 0, total_frame_num - 1, max_frames_num, dtype=int - ) - frame_idx = uniform_sampled_frames.tolist() - frames = vr.get_batch(frame_idx).asnumpy() - - base64_frames = [] - for frame in frames: - pil_img = Image.fromarray(frame) - buff = io.BytesIO() - pil_img.save(buff, format="JPEG") - base64_str = base64.b64encode(buff.getvalue()).decode("utf-8") - base64_frames.append(base64_str) - - messages = [{"role": "user", "content": []}] - frame_format = { - "type": "image_url", - "image_url": {"url": "data:image/jpeg;base64,{}"}, - "modalities": "video", - } - - for base64_frame in base64_frames: - frame_format["image_url"]["url"] = "data:image/jpeg;base64,{}".format( - base64_frame - ) - messages[0]["content"].append(frame_format.copy()) - - prompt = {"type": "text", "text": "Please describe the video briefly."} - messages[0]["content"].append(prompt) - - return messages - - def get_prompt_from_messages(self, messages): - text = ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - "<|im_start|>user\n" - ) - image_data = [] - for content in messages[0]["content"]: - if content["type"] == "image_url": - text += "\n" - image_data.append(content["image_url"]["url"]) - text += "Please describe the video briefly.<|im_end|>\n<|im_start|>assistant\n" - return text, image_data - - def generate(self, text, image_data): - response = requests.post( - self.base_url + "/generate", - json={ - "text": text, - "image_data": image_data, - "sampling_params": { - "temperature": 0, - "max_new_tokens": 32, - "no_stop_trim": True, - "skip_special_tokens": False, - }, - "modalities": ["multi-images"], - }, - ).json() - return response["text"] - - def generate_for_video(self, batch, num_frame) -> Union[str, list[str]]: - # prepare the video input about Steven introducing ipod nano - url = "https://raw.githubusercontent.com/evolvinglmms-lab/sglang/dev/onevision_local/assets/jobs.mp4" - cache_dir = os.path.expanduser("~/.cache") - file_path = os.path.join(cache_dir, "jobs.mp4") - os.makedirs(cache_dir, exist_ok=True) - if not os.path.exists(file_path): - response = requests.get(url) - response.raise_for_status() - with open(file_path, "wb") as f: - f.write(response.content) - - if not batch: - assert isinstance(num_frame, int) - messages = self.prepare_video_messages(file_path, max_frames_num=num_frame) - text, image_data = self.get_prompt_from_messages(messages) - return self.generate(text, image_data) - else: - assert isinstance(num_frame, list) - func_args = [] - for max_frames_num in num_frame: - messages = self.prepare_video_messages( - file_path, - max_frames_num=max_frames_num, - ) - text, image_data = self.get_prompt_from_messages(messages) - func_args.append((text, image_data)) - - with ThreadPoolExecutor(max_workers=10) as executor: - responses = list(executor.map(lambda p: self.generate(*p), func_args)) - - return responses - - def run_generate(self, chunked_prefill_size, batch, num_frame): - # launch server - model = "lmms-lab/llava-onevision-qwen2-7b-ov" - # model = "meta-llama/Llama-3.2-11B-Vision-Instruct" - self.base_url = DEFAULT_URL_FOR_TEST - process = popen_launch_server( - model, - self.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--chunked-prefill-size", - f"{chunked_prefill_size}", - ], - ) - try: - return self.generate_for_video(batch, num_frame) - finally: - kill_process_tree(process.pid) - - def test_chunked_prefill(self): - output_chunked = self.run_generate( - chunked_prefill_size=1024, batch=False, num_frame=1 - ) - output_no_chunked = self.run_generate( - chunked_prefill_size=-1, batch=False, num_frame=1 - ) - - print("output with chunked prefill:") - print(output_chunked) - print("output without chunked prefill:") - print(output_no_chunked) - assert output_chunked == output_no_chunked - - output_chunked = self.run_generate( - chunked_prefill_size=1024, batch=True, num_frame=[2, 6, 8, 10] - ) - output_no_chunked = self.run_generate( - chunked_prefill_size=-1, batch=True, num_frame=[2, 6, 8, 10] - ) - - print("output with chunked prefill:") - print(output_chunked) - print("output without chunked prefill:") - print(output_no_chunked) - assert output_chunked == output_no_chunked - - -if __name__ == "__main__": - unittest.main()