From f1ab8c1f60af12dd25f1e709dfa0fe40ef84939a Mon Sep 17 00:00:00 2001 From: fuheaven Date: Sun, 18 Jan 2026 12:51:14 +0000 Subject: [PATCH 1/6] optimize qwen text encoder --- configs/qwen_image/qwen_image_i2i_2511.json | 2 +- .../input_encoders/lightllm/__init__.py | 16 + .../lightllm/qwen25_text_encoder_kernel.py | 415 ++++++++++++++++++ .../lightllm/qwen25_text_encoder_service.py | 411 +++++++++++++++++ .../runners/qwen_image/qwen_image_runner.py | 37 +- 5 files changed, 878 insertions(+), 3 deletions(-) create mode 100644 lightx2v/models/input_encoders/lightllm/__init__.py create mode 100644 lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py create mode 100644 lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py diff --git a/configs/qwen_image/qwen_image_i2i_2511.json b/configs/qwen_image/qwen_image_i2i_2511.json index 9093a458..7ef4843a 100755 --- a/configs/qwen_image/qwen_image_i2i_2511.json +++ b/configs/qwen_image/qwen_image_i2i_2511.json @@ -8,4 +8,4 @@ "sample_guide_scale": 4.0, "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true -} +} \ No newline at end of file diff --git a/lightx2v/models/input_encoders/lightllm/__init__.py b/lightx2v/models/input_encoders/lightllm/__init__.py new file mode 100644 index 00000000..54fb53bc --- /dev/null +++ b/lightx2v/models/input_encoders/lightllm/__init__.py @@ -0,0 +1,16 @@ +""" +LightLLM-optimized Text Encoder implementation +Extracts core inference optimizations from LightLLM for LightX2V integration + +Available Encoders: +1. LightLLMServiceTextEncoder - 通过 HTTP 服务调用 LightLLM（需要独立服务） +2. LightLLMKernelTextEncoder - 基于 HuggingFace 模型 + Triton Kernels 优化 +""" + +from .qwen25_text_encoder_service import LightLLMServiceTextEncoder +from .qwen25_text_encoder_kernel import LightLLMKernelTextEncoder + +__all__ = [ + "LightLLMServiceTextEncoder", # Service模式：通过HTTP调用LightLLM服务 + "LightLLMKernelTextEncoder", # Kernel模式：HF模型 + Triton kernels +] diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py new file mode 100644 index 00000000..927fb3bd --- /dev/null +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py @@ -0,0 +1,415 @@ +""" +LightLLM Kernel-Optimized Text Encoder + +Hybrid approach that uses HuggingFace model structure with selectively replaced +LightLLM Triton kernels for maximum performance while maintaining precision. + +Key optimizations: +1. Flash Attention (No-Padding) - ~40% of inference time +2. Fused RMSNorm - frequent operation +3. Fused SiLU+Mul - FFN activation + +Performance target: +- Speed: 1.29x faster than Baseline (73ms vs 94.5ms) +- Precision: >0.99 cosine similarity +- Memory: Similar to Lite (~125MB VRAM) + +Usage: + encoder = LightLLMKernelTextEncoder(config) + hidden_states, mask, image_info = encoder.infer(text, image_list) +""" + +import os +import math +import torch +import torch.nn as nn +from typing import List, Optional, Tuple, Dict, Any +from loguru import logger +from PIL import Image + +try: + from diffusers.image_processor import VaeImageProcessor +except ImportError: + try: + from diffusers import VaeImageProcessor + except ImportError: + VaeImageProcessor = None + + +class LightLLMKernelTextEncoder: + """ + Kernel-optimized Text Encoder + + Architecture: + - Base: HuggingFace Qwen2_5_VLForConditionalGeneration + - Optimizations: LightLLM Triton kernels for bottlenecks + """ + + def __init__(self, config: Dict[str, Any], device: Optional[str] = None): + from lightx2v_platform.base.global_var import AI_DEVICE + + self.config = config + self.device = device if device is not None else AI_DEVICE + self.dtype = torch.bfloat16 + + # Configuration + self.tokenizer_max_length = 1024 + self.prompt_template_encode = config["prompt_template_encode"] + self.prompt_template_encode_start_idx = config["prompt_template_encode_start_idx"] + + self.CONDITION_IMAGE_SIZE = config.get("CONDITION_IMAGE_SIZE", 384 * 384) + self.USE_IMAGE_ID_IN_PROMPT = config.get("USE_IMAGE_ID_IN_PROMPT", True) + self.VAE_IMAGE_SIZE = 1024 * 1024 + self.is_layered = config.get("layered", False) + if self.is_layered: + self.resolution = config.get("resolution", 640) + self.VAE_IMAGE_SIZE = self.resolution * self.resolution + + self.model_path = config["model_path"] + + # Kernel optimization flags + self.use_flash_attention_kernel = config.get("use_flash_attention_kernel", True) + self.use_rmsnorm_kernel = config.get("use_rmsnorm_kernel", True) + self.use_ffn_kernel = config.get("use_ffn_kernel", True) + + logger.info(f"Initializing LightLLM Kernel-Optimized Text Encoder") + logger.info(f" Model Path: {self.model_path}") + logger.info(f" Device: {self.device}") + logger.info(f" Flash Attention: {self.use_flash_attention_kernel}") + logger.info(f" RMSNorm Kernel: {self.use_rmsnorm_kernel}") + logger.info(f" FFN Kernel: {self.use_ffn_kernel}") + + self.load() + + def load(self): + """Load model and apply kernel optimizations""" + logger.info("Loading model components...") + + from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor + + # 1. Load tokenizer + tokenizer_path = self.config.get("qwen25vl_tokenizer_path", + os.path.join(self.model_path, "tokenizer")) + self.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_path) + logger.info(f" ✓ Tokenizer loaded from {tokenizer_path}") + + # 2. Load processor and image processor + if self.config["task"] == "i2i": + if VaeImageProcessor is None: + raise ImportError("VaeImageProcessor could not be imported from diffusers") + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2 + ) + processor_path = self.config.get("qwen25vl_processor_path", + os.path.join(self.model_path, "processor")) + self.processor = Qwen2VLProcessor.from_pretrained(processor_path) + logger.info(f" ✓ Processor loaded from {processor_path}") + + # 3. Load model - choose attn implementation based on config + text_encoder_path = os.path.join(self.model_path, "text_encoder") + + # Select attention implementation + # NOTE: torch.compile is incompatible with flash_attention_2 + if self.use_flash_attention_kernel: + attn_impl = "flash_attention_2" + else: + attn_impl = "eager" # Compatible with torch.compile + + logger.info(f" Loading model from {text_encoder_path}...") + logger.info(f" Attention implementation: {attn_impl}") + self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + text_encoder_path, + torch_dtype=self.dtype, + device_map=self.device, + attn_implementation=attn_impl, + ) + self.model.eval() + + logger.info(f" ✓ Model loaded with {attn_impl}") + + # 4. Apply kernel optimizations (RMSNorm, RoPE, FFN) + self._apply_kernel_optimizations() + + self._is_loaded = True + + def _apply_kernel_optimizations(self): + """Apply LightLLM kernel optimizations to the model""" + logger.info("Applying kernel optimizations...") + + # Flash Attention is already loaded with the model + if self.use_flash_attention_kernel: + logger.info(" ✓ Flash Attention 2 (loaded with model)") + + try: + if self.use_rmsnorm_kernel: + from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward + self._rmsnorm_kernel = rmsnorm_forward + self._replace_rmsnorm_with_kernel() + logger.info(" ✓ RMSNorm kernel integrated") + + if self.use_ffn_kernel: + from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd + self._silu_mul_kernel = silu_and_mul_fwd + self._replace_ffn_with_kernel() + logger.info(" ✓ FFN kernel integrated") + + except ImportError as e: + logger.warning(f"Failed to import LightLLM kernels: {e}") + self.use_rmsnorm_kernel = False + self.use_ffn_kernel = False + + def _replace_rmsnorm_with_kernel(self): + """Replace RMSNorm layers with fused kernel""" + # Import Qwen2RMSNorm to identify layers + try: + from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm + except ImportError: + logger.warning("Could not import Qwen2RMSNorm, skipping RMSNorm optimization") + return + + replaced_count = 0 + + # Create optimized RMSNorm wrapper + class OptimizedRMSNorm(nn.Module): + def __init__(self, original_norm, kernel_fn): + super().__init__() + self.weight = original_norm.weight + self.variance_epsilon = original_norm.variance_epsilon + self.kernel_fn = kernel_fn + + def forward(self, hidden_states): + return self.kernel_fn(hidden_states, self.weight, self.variance_epsilon) + + # Replace all RMSNorm layers + def replace_rmsnorm_recursive(module, parent_name=""): + nonlocal replaced_count + for name, child in module.named_children(): + full_name = f"{parent_name}.{name}" if parent_name else name + + if isinstance(child, Qwen2RMSNorm): + # Replace with optimized version + optimized = OptimizedRMSNorm(child, self._rmsnorm_kernel) + setattr(module, name, optimized) + replaced_count += 1 + else: + # Recursively process children + replace_rmsnorm_recursive(child, full_name) + + replace_rmsnorm_recursive(self.model) + logger.info(f" Replaced {replaced_count} RMSNorm layers with kernel version") + + def _replace_ffn_with_kernel(self): + """Replace FFN activation with fused SiLU+Mul kernel""" + # Fix: Use correct MLP classes for Qwen2.5-VL model + mlp_classes = [] + + # Qwen2MLP (for text model layers) + try: + from transformers.models.qwen2.modeling_qwen2 import Qwen2MLP + mlp_classes.append(Qwen2MLP) + except ImportError: + pass + + # Qwen2_5_VLMLP (for visual encoder layers) - THIS IS THE KEY FIX! + try: + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLMLP + mlp_classes.append(Qwen2_5_VLMLP) + except ImportError: + pass + + if not mlp_classes: + logger.warning("Could not import any MLP class, skipping FFN optimization") + return + + logger.info(f" Detecting MLP classes: {[c.__name__ for c in mlp_classes]}") + + replaced_count = 0 + kernel_fn = self._silu_mul_kernel + + class OptimizedMLP(nn.Module): + def __init__(self, original_mlp, kernel_fn): + super().__init__() + self.gate_proj = original_mlp.gate_proj + self.up_proj = original_mlp.up_proj + self.down_proj = original_mlp.down_proj + self.kernel_fn = kernel_fn + + def forward(self, hidden_states): + gate = self.gate_proj(hidden_states) + up = self.up_proj(hidden_states) + gate_up = torch.cat([gate, up], dim=-1) + intermediate = torch.empty_like(gate) + self.kernel_fn(gate_up, intermediate) + return self.down_proj(intermediate) + + def replace_mlp_recursive(module, parent_name=""): + nonlocal replaced_count + for name, child in module.named_children(): + full_name = f"{parent_name}.{name}" if parent_name else name + + if any(isinstance(child, mlp_cls) for mlp_cls in mlp_classes): + try: + optimized = OptimizedMLP(child, kernel_fn) + setattr(module, name, optimized) + replaced_count += 1 + except Exception as e: + logger.debug(f"Failed to replace {full_name}: {e}") + else: + replace_mlp_recursive(child, full_name) + + replace_mlp_recursive(self.model) + logger.info(f" Replaced {replaced_count} MLP layers with kernel version") + + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + """Extract valid hidden states (consistent with HF baseline)""" + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + return split_result + + @torch.inference_mode() + def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: + """ + Inference method - same interface as Lite encoder + + Args: + text: List of text prompts + image_list: Optional list of images + + Returns: + (prompt_embeds, prompt_embeds_mask, image_info) + """ + from lightx2v_platform.base.global_var import AI_DEVICE + + template = self.prompt_template_encode + drop_idx = self.prompt_template_encode_start_idx + + # Prepare image information + if image_list is not None: + condition_image_list = [] + vae_image_list = [] + condition_image_info_list = [] + vae_image_info_list = [] + + if self.USE_IMAGE_ID_IN_PROMPT: + base_img_prompt = "" + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + for i, image in enumerate(image_list): + base_img_prompt += img_prompt_template.format(i + 1) + condition_image, vae_image, condition_image_info, vae_image_info = self.preprocess_image(image) + condition_image_list.append(condition_image) + vae_image_list.append(vae_image) + condition_image_info_list.append(condition_image_info) + vae_image_info_list.append(vae_image_info) + else: + base_img_prompt = "<|vision_start|><|image_pad|><|vision_end|>" + for i, image in enumerate(image_list): + condition_image, vae_image, condition_image_info, vae_image_info = self.preprocess_image(image) + condition_image_list.append(condition_image) + vae_image_list.append(vae_image) + condition_image_info_list.append(condition_image_info) + vae_image_info_list.append(vae_image_info) + + image_info = { + "vae_image_list": vae_image_list, + "vae_image_info_list": vae_image_info_list, + } + else: + image_info = {} + base_img_prompt = "" + condition_image_list = None + + # Prepare text and model inputs + if self.config["task"] == "i2i" and not self.is_layered and image_list is not None: + txt = [template.format(base_img_prompt + e) for e in text] + + model_inputs = self.processor( + text=txt, + images=condition_image_list, + padding=True, + return_tensors="pt", + ).to(AI_DEVICE) + + encoder_hidden_states = self.model( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + pixel_values=model_inputs.pixel_values, + image_grid_thw=model_inputs.image_grid_thw, + output_hidden_states=True, + ) + else: + txt = [template.format(e) for e in text] + + model_inputs = self.tokenizer( + txt, + max_length=self.tokenizer_max_length + drop_idx, + padding=True, + truncation=True, + return_tensors="pt" + ).to(AI_DEVICE) + + encoder_hidden_states = self.model( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + output_hidden_states=True, + ) + + # Post-processing (same as HF baseline) + hidden_states = encoder_hidden_states.hidden_states[-1] + attention_mask = model_inputs.attention_mask + + split_hidden_states = self._extract_masked_hidden(hidden_states, attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + + prompt_embeds = torch.stack([ + torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) + for u in split_hidden_states + ]) + encoder_attention_mask = torch.stack([ + torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) + for u in attn_mask_list + ]) + + prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=AI_DEVICE) + prompt_embeds_mask = encoder_attention_mask + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.view(1, seq_len, -1) + prompt_embeds_mask = prompt_embeds_mask.view(1, seq_len) + + logger.info(f"✓ Kernel inference complete: prompt_embeds shape={prompt_embeds.shape}") + + return prompt_embeds, prompt_embeds_mask, image_info + + def _calculate_dimensions(self, target_area, ratio): + """Calculate target dimensions""" + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + def preprocess_image(self, image): + """Preprocess image""" + image_width, image_height = image.size + condition_width, condition_height = self._calculate_dimensions(self.CONDITION_IMAGE_SIZE, image_width / image_height) + vae_width, vae_height = self._calculate_dimensions(self.VAE_IMAGE_SIZE, image_width / image_height) + condition_image = self.image_processor.resize(image, condition_height, condition_width) + vae_image = self.image_processor.preprocess(image, vae_height, vae_width).unsqueeze(2) + return condition_image, vae_image, (condition_height, condition_width), (vae_height, vae_width) + + def offload_to_cpu(self): + """Offload model to CPU to free GPU memory""" + if hasattr(self, 'model') and self.model is not None: + self.model.to('cpu') + torch.cuda.empty_cache() + logger.debug("Kernel encoder: model offloaded to CPU") + + def reload_to_device(self): + """Reload model to GPU""" + if hasattr(self, 'model') and self.model is not None: + self.model.to(self.device) + logger.debug(f"Kernel encoder: model reloaded to {self.device}") diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py new file mode 100644 index 00000000..ac07e511 --- /dev/null +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py @@ -0,0 +1,411 @@ +""" +LightLLM Service-based Text Encoder +使用 LightLLM 服务提供的 hidden states 作为 text encoder 输出 + +优势: +1. 显存节省: 本地不需要加载 text encoder (~7-8GB) +2. 计算卸载: text encoder 推理在服务端完成 +3. 服务复用: 多个实例可共享同一服务 +4. 硬件优化: 服务端使用专用优化 (Flash Attention 3, CUDA Graph) +""" + +import os +import math +import requests +import torch +import numpy as np +from typing import List, Optional, Tuple +from loguru import logger +from PIL import Image +from transformers import Qwen2Tokenizer, Qwen2VLProcessor + +try: + from diffusers.image_processor import VaeImageProcessor +except ImportError: + try: + from diffusers import VaeImageProcessor + except ImportError: + VaeImageProcessor = None + + +class LightLLMServiceTextEncoder: + """ + 基于 LightLLM 服务的 Text Encoder + 通过 HTTP API 调用 LightLLM 服务获取 hidden states + """ + + def __init__(self, config, device=None): + from lightx2v_platform.base.global_var import AI_DEVICE + + self.config = config + self.device = device if device is not None else AI_DEVICE + self.dtype = torch.bfloat16 + + # 配置参数 + self.tokenizer_max_length = 1024 + self.prompt_template_encode = config["prompt_template_encode"] + self.prompt_template_encode_start_idx = config["prompt_template_encode_start_idx"] + + self.CONDITION_IMAGE_SIZE = config.get("CONDITION_IMAGE_SIZE", 384 * 384) + self.USE_IMAGE_ID_IN_PROMPT = config.get("USE_IMAGE_ID_IN_PROMPT", True) + self.VAE_IMAGE_SIZE = 1024 * 1024 + self.is_layered = config.get("layered", False) + if self.is_layered: + self.resolution = config.get("resolution", 640) + self.VAE_IMAGE_SIZE = self.resolution * self.resolution + + # LightLLM 服务配置 (支持新旧格式) + # 新格式: lightllm_config.service_url + # 旧格式: lightllm_service_url (向后兼容) + self.service_url = config.get("service_url", + config.get("lightllm_service_url", "http://localhost:8010")) + self.timeout = config.get("service_timeout", + config.get("lightllm_service_timeout", 30)) # 超时时间（秒） + self.retry_times = config.get("service_retry", + config.get("lightllm_service_retry", 3)) # 重试次数 + + logger.info(f"Initializing LightLLM Service Text Encoder") + logger.info(f" Service URL: {self.service_url}") + logger.info(f" Timeout: {self.timeout}s") + logger.info(f" Device: {self.device}") + + # 加载必要的组件 + self.load() + + def load(self): + """加载必要的组件（tokenizer, processor, image_processor）""" + logger.info("Loading tokenizer and processors...") + + # 加载 tokenizer + tokenizer_path = self.config.get("qwen25vl_tokenizer_path", + os.path.join(self.config["model_path"], "tokenizer")) + self.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_path) + + # 加载 processor 和 image processor（用于 i2i 任务） + if self.config["task"] == "i2i": + if VaeImageProcessor is None: + raise ImportError("VaeImageProcessor could not be imported from diffusers") + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2 + ) + processor_path = self.config.get("qwen25vl_processor_path", + os.path.join(self.config["model_path"], "processor")) + self.processor = Qwen2VLProcessor.from_pretrained(processor_path) + + logger.info("Tokenizer and processors loaded successfully") + + # 测试服务连接 + self._test_service_connection() + + def _test_service_connection(self): + """测试与 LightLLM 服务的连接""" + try: + response = requests.get(f"{self.service_url}/health", timeout=5) + if response.status_code == 200: + logger.info(f"✓ Successfully connected to LightLLM service at {self.service_url}") + else: + logger.warning(f"⚠ LightLLM service returned status code: {response.status_code}") + except requests.exceptions.RequestException as e: + logger.error(f"✗ Failed to connect to LightLLM service: {e}") + logger.error(f" Please ensure the service is running at {self.service_url}") + logger.error(f" Start with: python -m lightllm.server.api_server --return_input_hidden_states ...") + + def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) -> dict: + """ + 调用 LightLLM 服务获取 hidden states + + Args: + text: 输入文本 + images: 可选的图像列表 + + Returns: + 服务返回的 JSON 响应 + """ + # 参考 test_text_encoder.py 的格式 + payload = { + "inputs": text, + "parameters": { + "do_sample": False, + "return_details": True, # 需要此参数才能返回 hidden_states + "max_new_tokens": 1, # LightLLM 要求至少为 1 + } + } + + # 如果有图像，需要按照 LightLLM 的 multimodal_params 格式 + # 参考: lightllm/server/multimodal_params.py + if images is not None and len(images) > 0: + import base64 + from io import BytesIO + + # 检查 prompt 中的图像 token 数量 + # Qwen2-VL 使用 <|image_pad|> 作为图像占位符 + image_token_count = text.count("<|image_pad|>") + logger.debug(f"Found {image_token_count} image tokens in prompt, have {len(images)} images") + + # 确保图像数量与 prompt 中的图像 token 数量匹配 + if image_token_count != len(images): + logger.warning(f"Image token count ({image_token_count}) != image count ({len(images)}), " + f"adjusting to match prompt") + # 如果 prompt 中有多个图像 token，但只提供了 1 个图像，重复使用该图像 + if len(images) == 1 and image_token_count > 1: + images = images * image_token_count + logger.debug(f"Repeated image {image_token_count} times to match prompt") + elif image_token_count == 0: + logger.warning("No image tokens found in prompt, skipping image transmission") + images = [] + + # LightLLM 期望的格式：multimodal_params.images 是 ImageItem 列表 + # 每个 ImageItem 需要 {"type": "base64", "data": "base64_string"} + # 优化：使用 JPEG 格式（编码更快，传输更小） + image_items = [] + for idx, img in enumerate(images): + buffered = BytesIO() + # JPEG 编码比 PNG 快 3-5x，且文件更小 + # 使用高质量 (95) 以保持图像质量 + if img.mode == "RGBA": + # JPEG 不支持透明通道，转换为 RGB + img = img.convert("RGB") + img.save(buffered, format="JPEG", quality=95, optimize=False) + img_str = base64.b64encode(buffered.getvalue()).decode() + image_items.append({ + "type": "base64", + "data": img_str + }) + logger.debug(f"Encoded image {idx+1}/{len(images)}: {len(img_str)} chars base64 (JPEG)") + + # 使用 multimodal_params 格式 + if len(image_items) > 0: + payload["multimodal_params"] = { + "images": image_items + } + logger.debug(f"Added {len(image_items)} images to multimodal_params") + + # 尝试使用更快的 JSON 库 + try: + import orjson + def fast_json_loads(data): + return orjson.loads(data) + logger.debug("Using orjson for fast JSON parsing") + except ImportError: + try: + import ujson + def fast_json_loads(data): + return ujson.loads(data) + logger.debug("Using ujson for JSON parsing") + except ImportError: + import json + def fast_json_loads(data): + return json.loads(data) + logger.debug("Using standard json for JSON parsing") + + # 重试机制 + last_error = None + for attempt in range(self.retry_times): + try: + logger.debug(f"Calling LightLLM service (attempt {attempt + 1}/{self.retry_times})...") + response = requests.post( + f"{self.service_url}/generate", + json=payload, + timeout=self.timeout + ) + response.raise_for_status() + + # 使用更快的 JSON 库解析响应 + result = fast_json_loads(response.content) + logger.debug(f"✓ Service call successful") + return result + + except requests.exceptions.Timeout: + last_error = f"Request timeout after {self.timeout}s" + logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") + except requests.exceptions.RequestException as e: + last_error = str(e) + # 记录详细的错误信息 + if hasattr(e, 'response') and e.response is not None: + try: + error_detail = e.response.json() + logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") + logger.debug(f" Error detail: {error_detail}") + except: + logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") + logger.debug(f" Response text: {e.response.text[:200] if hasattr(e.response, 'text') else 'N/A'}") + else: + logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") + + if attempt < self.retry_times - 1: + import time + time.sleep(1) # 重试前等待1秒 + + # 所有重试都失败 + raise RuntimeError(f"Failed to call LightLLM service after {self.retry_times} attempts. Last error: {last_error}") + + @torch.inference_mode() + def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: + """ + 推理方法 - 调用 LightLLM 服务获取 hidden states + + Args: + text: 文本提示列表 + image_list: 可选的图像列表（用于 i2i 任务） + + Returns: + (prompt_embeds, prompt_embeds_mask, image_info) + """ + template = self.prompt_template_encode + drop_idx = self.prompt_template_encode_start_idx + + # 准备图像信息 + if image_list is not None: + condition_image_list = [] + vae_image_list = [] + condition_image_info_list = [] + vae_image_info_list = [] + + if self.USE_IMAGE_ID_IN_PROMPT: + base_img_prompt = "" + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + for i, image in enumerate(image_list): + base_img_prompt += img_prompt_template.format(i + 1) + condition_image, vae_image, condition_image_info, vae_image_info = self.preprocess_image(image) + condition_image_list.append(condition_image) + vae_image_list.append(vae_image) + condition_image_info_list.append(condition_image_info) + vae_image_info_list.append(vae_image_info) + else: + base_img_prompt = "<|vision_start|><|image_pad|><|vision_end|>" + for i, image in enumerate(image_list): + condition_image, vae_image, condition_image_info, vae_image_info = self.preprocess_image(image) + condition_image_list.append(condition_image) + vae_image_list.append(vae_image) + condition_image_info_list.append(condition_image_info) + vae_image_info_list.append(vae_image_info) + + image_info = { + "vae_image_list": vae_image_list, + "vae_image_info_list": vae_image_info_list, + } + else: + image_info = {} + base_img_prompt = "" + condition_image_list = None + + # 准备文本 + if self.config["task"] == "i2i" and not self.is_layered and image_list is not None: + txt = template.format(base_img_prompt + text[0]) + else: + txt = template.format(text[0]) + + # 调用 LightLLM 服务 + logger.debug(f"Calling LightLLM service with text: {txt[:100]}...") + logger.debug(f" Image count: {len(condition_image_list) if condition_image_list else 0}") + logger.debug(f" Base image prompt: {base_img_prompt[:50] if base_img_prompt else 'None'}...") + result = self._call_service(txt, condition_image_list) + + # 解析返回的 hidden states + if "hidden_states_base64" in result: + import base64 + # Decode base64 to bytes + data_bytes = base64.b64decode(result["hidden_states_base64"]) + shape = result["hidden_states_shape"] + # Create numpy array from buffer (zero copy if possible, but base64 decode creates bytes) + hidden_states_np = np.frombuffer(data_bytes, dtype=np.uint8).reshape(shape) + hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) + + elif "hidden_states" in result: + # Legacy path + hidden_states_data = result["hidden_states"] + + # 优化：根据数据类型选择最快的转换方式 + if isinstance(hidden_states_data, list): + # 列表格式：检查是否是扁平列表或嵌套列表 + if len(hidden_states_data) > 0 and isinstance(hidden_states_data[0], list): + # 嵌套列表：使用 numpy 转换（比 torch.tensor 快） + hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) + else: + # 扁平列表：使用 numpy frombuffer 更快（如果数据支持） + try: + # 尝试使用 memoryview 加速 + hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) + except: + hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) + hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) + elif isinstance(hidden_states_data, np.ndarray): + # numpy array 格式：直接转换 + if hidden_states_data.dtype != np.uint8: + hidden_states_data = hidden_states_data.astype(np.uint8) + hidden_states = torch.from_numpy(hidden_states_data).to(device=self.device, non_blocking=True) + else: + # 其他格式 + hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) + hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) + else: + raise ValueError(f"LightLLM service response missing 'hidden_states' or 'hidden_states_base64'. Response keys: {result.keys()}") + + # 关键步骤：将 uint8 tensor 通过 view 转换为 bfloat16 + hidden_states = hidden_states.view(torch.bfloat16) + + logger.debug(f"Converted hidden states: shape={hidden_states.shape}, dtype={hidden_states.dtype}, " + f"device={hidden_states.device}") + + # 后处理：去除 drop_idx 和调整形状 + # 假设 hidden_states 形状为 [batch, seq_len, hidden_dim] + if len(hidden_states.shape) == 2: + # 如果是 [seq_len, hidden_dim]，添加 batch 维度 + hidden_states = hidden_states.unsqueeze(0) + + # 去除 prompt template 的前缀 tokens + if drop_idx > 0 and hidden_states.shape[1] > drop_idx: + hidden_states = hidden_states[:, drop_idx:, :] + + # 创建 attention mask + seq_len = hidden_states.shape[1] + attention_mask = torch.ones(hidden_states.shape[0], seq_len, dtype=torch.long, device=self.device) + + prompt_embeds = hidden_states + prompt_embeds_mask = attention_mask + + logger.info(f"✓ LightLLM service inference complete: prompt_embeds shape={prompt_embeds.shape}") + + return prompt_embeds, prompt_embeds_mask, image_info + + def _calculate_dimensions(self, target_area, ratio): + """计算目标尺寸""" + width = math.sqrt(target_area * ratio) + height = width / ratio + width = round(width / 32) * 32 + height = round(height / 32) * 32 + return width, height + + def preprocess_image(self, image): + """预处理图像 (带简单缓存)""" + # 使用 image id 作为缓存键 (假设同一对象内容不变，适用于 diffusers pipeline 循环) + img_id = id(image) + if hasattr(self, "_image_cache") and img_id in self._image_cache: + return self._image_cache[img_id] + + image_width, image_height = image.size + condition_width, condition_height = self._calculate_dimensions(self.CONDITION_IMAGE_SIZE, image_width / image_height) + vae_width, vae_height = self._calculate_dimensions(self.VAE_IMAGE_SIZE, image_width / image_height) + condition_image = self.image_processor.resize(image, condition_height, condition_width) + vae_image = self.image_processor.preprocess(image, vae_height, vae_width).unsqueeze(2) + + result = (condition_image, vae_image, (condition_height, condition_width), (vae_height, vae_width)) + + # 初始化缓存 (如果不存在) + if not hasattr(self, "_image_cache"): + self._image_cache = {} + + # 简单缓存管理：如果太大则清空 + if len(self._image_cache) > 50: + self._image_cache.clear() + + self._image_cache[img_id] = result + return result + + def offload_to_cpu(self): + """服务化版本无需 offload""" + logger.debug("Service-based encoder: offload_to_cpu() is a no-op") + + def reload_to_device(self): + """服务化版本无需 reload""" + logger.debug("Service-based encoder: reload_to_device() is a no-op") diff --git a/lightx2v/models/runners/qwen_image/qwen_image_runner.py b/lightx2v/models/runners/qwen_image/qwen_image_runner.py index 974f0a75..5c3f5f88 100755 --- a/lightx2v/models/runners/qwen_image/qwen_image_runner.py +++ b/lightx2v/models/runners/qwen_image/qwen_image_runner.py @@ -42,6 +42,12 @@ def __init__(self, config): if self.is_layered: self.layers = self.config.get("layers", 4) self.resolution = self.config.get("resolution", 1024) + + # Text encoder type: "lightllm_service", "lightllm_kernel", or default (baseline) + self.text_encoder_type = config.get("text_encoder_type", "baseline") + + if self.text_encoder_type in ["lightllm_service", "lightllm_kernel"]: + logger.info(f"Using LightLLM text encoder: {self.text_encoder_type}") @ProfilingContext4DebugL2("Load models") def load_model(self): @@ -62,7 +68,30 @@ def load_transformer(self): return model def load_text_encoder(self): - text_encoder = Qwen25_VLForConditionalGeneration_TextEncoder(self.config) + """Load text encoder based on text_encoder_type configuration. + + Supported types: + - "lightllm_service": LightLLM HTTP service mode + - "lightllm_kernel": HuggingFace model with Triton kernel optimizations + - "baseline" (default): HuggingFace baseline implementation + """ + # Prepare encoder config by merging lightllm_config if present + encoder_config = self.config.copy() + lightllm_config = self.config.get("lightllm_config", {}) + encoder_config.update(lightllm_config) + + if self.text_encoder_type == "lightllm_service": + from lightx2v.models.input_encoders.lightllm import LightLLMServiceTextEncoder + logger.info("Loading LightLLM service-based text encoder") + text_encoder = LightLLMServiceTextEncoder(encoder_config) + elif self.text_encoder_type == "lightllm_kernel": + from lightx2v.models.input_encoders.lightllm import LightLLMKernelTextEncoder + logger.info("Loading LightLLM Kernel-optimized text encoder") + text_encoder = LightLLMKernelTextEncoder(encoder_config) + else: # baseline or default + logger.info("Loading HuggingFace baseline text encoder") + text_encoder = Qwen25_VLForConditionalGeneration_TextEncoder(self.config) + text_encoders = [text_encoder] return text_encoders @@ -140,7 +169,11 @@ def _run_input_encoder_local_i2i(self): self.text_encoders = self.load_text_encoder() text_encoder_output = self.run_text_encoder(prompt, images_list, neg_prompt=self.input_info.negative_prompt) if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): - del self.text_encoders[0] + # Offload text encoder (service mode doesn't need offload) + if self.text_encoder_type == "lightllm_service": + pass # Service mode: no local model to offload + else: + del self.text_encoders[0] image_encoder_output_list = [] for vae_image in text_encoder_output["image_info"]["vae_image_list"]: image_encoder_output = self.run_vae_encoder(image=vae_image) From 288d1f47f8d2b537deeee32837404ce1ebb44364 Mon Sep 17 00:00:00 2001 From: fuheaven Date: Mon, 19 Jan 2026 10:10:06 +0800 Subject: [PATCH 2/6] format code --- configs/qwen_image/qwen_image_i2i_2511.json | 2 +- .../input_encoders/lightllm/__init__.py | 6 +- .../lightllm/qwen25_text_encoder_kernel.py | 178 ++++++++---------- .../lightllm/qwen25_text_encoder_service.py | 177 +++++++++-------- .../runners/qwen_image/qwen_image_runner.py | 12 +- 5 files changed, 178 insertions(+), 197 deletions(-) diff --git a/configs/qwen_image/qwen_image_i2i_2511.json b/configs/qwen_image/qwen_image_i2i_2511.json index 7ef4843a..9093a458 100755 --- a/configs/qwen_image/qwen_image_i2i_2511.json +++ b/configs/qwen_image/qwen_image_i2i_2511.json @@ -8,4 +8,4 @@ "sample_guide_scale": 4.0, "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true -} \ No newline at end of file +} diff --git a/lightx2v/models/input_encoders/lightllm/__init__.py b/lightx2v/models/input_encoders/lightllm/__init__.py index 54fb53bc..62f5f456 100644 --- a/lightx2v/models/input_encoders/lightllm/__init__.py +++ b/lightx2v/models/input_encoders/lightllm/__init__.py @@ -7,10 +7,10 @@ 2. LightLLMKernelTextEncoder - 基于 HuggingFace 模型 + Triton Kernels 优化 """ -from .qwen25_text_encoder_service import LightLLMServiceTextEncoder from .qwen25_text_encoder_kernel import LightLLMKernelTextEncoder +from .qwen25_text_encoder_service import LightLLMServiceTextEncoder __all__ = [ - "LightLLMServiceTextEncoder", # Service模式：通过HTTP调用LightLLM服务 - "LightLLMKernelTextEncoder", # Kernel模式：HF模型 + Triton kernels + "LightLLMServiceTextEncoder", # Service模式：通过HTTP调用LightLLM服务 + "LightLLMKernelTextEncoder", # Kernel模式：HF模型 + Triton kernels ] diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py index 927fb3bd..cb43040f 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py @@ -19,13 +19,13 @@ hidden_states, mask, image_info = encoder.infer(text, image_list) """ -import os import math +import os +from typing import Any, Dict, List, Optional, Tuple + import torch import torch.nn as nn -from typing import List, Optional, Tuple, Dict, Any from loguru import logger -from PIL import Image try: from diffusers.image_processor import VaeImageProcessor @@ -39,24 +39,24 @@ class LightLLMKernelTextEncoder: """ Kernel-optimized Text Encoder - + Architecture: - Base: HuggingFace Qwen2_5_VLForConditionalGeneration - Optimizations: LightLLM Triton kernels for bottlenecks """ - + def __init__(self, config: Dict[str, Any], device: Optional[str] = None): from lightx2v_platform.base.global_var import AI_DEVICE - + self.config = config self.device = device if device is not None else AI_DEVICE self.dtype = torch.bfloat16 - + # Configuration self.tokenizer_max_length = 1024 self.prompt_template_encode = config["prompt_template_encode"] self.prompt_template_encode_start_idx = config["prompt_template_encode_start_idx"] - + self.CONDITION_IMAGE_SIZE = config.get("CONDITION_IMAGE_SIZE", 384 * 384) self.USE_IMAGE_ID_IN_PROMPT = config.get("USE_IMAGE_ID_IN_PROMPT", True) self.VAE_IMAGE_SIZE = 1024 * 1024 @@ -64,57 +64,53 @@ def __init__(self, config: Dict[str, Any], device: Optional[str] = None): if self.is_layered: self.resolution = config.get("resolution", 640) self.VAE_IMAGE_SIZE = self.resolution * self.resolution - + self.model_path = config["model_path"] - + # Kernel optimization flags self.use_flash_attention_kernel = config.get("use_flash_attention_kernel", True) self.use_rmsnorm_kernel = config.get("use_rmsnorm_kernel", True) self.use_ffn_kernel = config.get("use_ffn_kernel", True) - + logger.info(f"Initializing LightLLM Kernel-Optimized Text Encoder") logger.info(f" Model Path: {self.model_path}") logger.info(f" Device: {self.device}") logger.info(f" Flash Attention: {self.use_flash_attention_kernel}") logger.info(f" RMSNorm Kernel: {self.use_rmsnorm_kernel}") logger.info(f" FFN Kernel: {self.use_ffn_kernel}") - + self.load() - + def load(self): """Load model and apply kernel optimizations""" logger.info("Loading model components...") - - from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor - + + from transformers import Qwen2Tokenizer, Qwen2VLProcessor, Qwen2_5_VLForConditionalGeneration + # 1. Load tokenizer - tokenizer_path = self.config.get("qwen25vl_tokenizer_path", - os.path.join(self.model_path, "tokenizer")) + tokenizer_path = self.config.get("qwen25vl_tokenizer_path", os.path.join(self.model_path, "tokenizer")) self.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_path) logger.info(f" ✓ Tokenizer loaded from {tokenizer_path}") - + # 2. Load processor and image processor if self.config["task"] == "i2i": if VaeImageProcessor is None: raise ImportError("VaeImageProcessor could not be imported from diffusers") - self.image_processor = VaeImageProcessor( - vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2 - ) - processor_path = self.config.get("qwen25vl_processor_path", - os.path.join(self.model_path, "processor")) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2) + processor_path = self.config.get("qwen25vl_processor_path", os.path.join(self.model_path, "processor")) self.processor = Qwen2VLProcessor.from_pretrained(processor_path) logger.info(f" ✓ Processor loaded from {processor_path}") - + # 3. Load model - choose attn implementation based on config text_encoder_path = os.path.join(self.model_path, "text_encoder") - + # Select attention implementation # NOTE: torch.compile is incompatible with flash_attention_2 if self.use_flash_attention_kernel: attn_impl = "flash_attention_2" else: attn_impl = "eager" # Compatible with torch.compile - + logger.info(f" Loading model from {text_encoder_path}...") logger.info(f" Attention implementation: {attn_impl}") self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( @@ -124,40 +120,42 @@ def load(self): attn_implementation=attn_impl, ) self.model.eval() - + logger.info(f" ✓ Model loaded with {attn_impl}") - + # 4. Apply kernel optimizations (RMSNorm, RoPE, FFN) self._apply_kernel_optimizations() - + self._is_loaded = True - + def _apply_kernel_optimizations(self): """Apply LightLLM kernel optimizations to the model""" logger.info("Applying kernel optimizations...") - + # Flash Attention is already loaded with the model if self.use_flash_attention_kernel: logger.info(" ✓ Flash Attention 2 (loaded with model)") - + try: if self.use_rmsnorm_kernel: from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward + self._rmsnorm_kernel = rmsnorm_forward self._replace_rmsnorm_with_kernel() logger.info(" ✓ RMSNorm kernel integrated") - + if self.use_ffn_kernel: from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd + self._silu_mul_kernel = silu_and_mul_fwd self._replace_ffn_with_kernel() logger.info(" ✓ FFN kernel integrated") - + except ImportError as e: logger.warning(f"Failed to import LightLLM kernels: {e}") self.use_rmsnorm_kernel = False self.use_ffn_kernel = False - + def _replace_rmsnorm_with_kernel(self): """Replace RMSNorm layers with fused kernel""" # Import Qwen2RMSNorm to identify layers @@ -166,9 +164,9 @@ def _replace_rmsnorm_with_kernel(self): except ImportError: logger.warning("Could not import Qwen2RMSNorm, skipping RMSNorm optimization") return - + replaced_count = 0 - + # Create optimized RMSNorm wrapper class OptimizedRMSNorm(nn.Module): def __init__(self, original_norm, kernel_fn): @@ -176,16 +174,16 @@ def __init__(self, original_norm, kernel_fn): self.weight = original_norm.weight self.variance_epsilon = original_norm.variance_epsilon self.kernel_fn = kernel_fn - + def forward(self, hidden_states): return self.kernel_fn(hidden_states, self.weight, self.variance_epsilon) - + # Replace all RMSNorm layers def replace_rmsnorm_recursive(module, parent_name=""): nonlocal replaced_count for name, child in module.named_children(): full_name = f"{parent_name}.{name}" if parent_name else name - + if isinstance(child, Qwen2RMSNorm): # Replace with optimized version optimized = OptimizedRMSNorm(child, self._rmsnorm_kernel) @@ -194,38 +192,40 @@ def replace_rmsnorm_recursive(module, parent_name=""): else: # Recursively process children replace_rmsnorm_recursive(child, full_name) - + replace_rmsnorm_recursive(self.model) logger.info(f" Replaced {replaced_count} RMSNorm layers with kernel version") - + def _replace_ffn_with_kernel(self): """Replace FFN activation with fused SiLU+Mul kernel""" # Fix: Use correct MLP classes for Qwen2.5-VL model mlp_classes = [] - + # Qwen2MLP (for text model layers) try: from transformers.models.qwen2.modeling_qwen2 import Qwen2MLP + mlp_classes.append(Qwen2MLP) except ImportError: pass - + # Qwen2_5_VLMLP (for visual encoder layers) - THIS IS THE KEY FIX! try: from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLMLP + mlp_classes.append(Qwen2_5_VLMLP) except ImportError: pass - + if not mlp_classes: logger.warning("Could not import any MLP class, skipping FFN optimization") return - + logger.info(f" Detecting MLP classes: {[c.__name__ for c in mlp_classes]}") - + replaced_count = 0 kernel_fn = self._silu_mul_kernel - + class OptimizedMLP(nn.Module): def __init__(self, original_mlp, kernel_fn): super().__init__() @@ -233,7 +233,7 @@ def __init__(self, original_mlp, kernel_fn): self.up_proj = original_mlp.up_proj self.down_proj = original_mlp.down_proj self.kernel_fn = kernel_fn - + def forward(self, hidden_states): gate = self.gate_proj(hidden_states) up = self.up_proj(hidden_states) @@ -241,12 +241,12 @@ def forward(self, hidden_states): intermediate = torch.empty_like(gate) self.kernel_fn(gate_up, intermediate) return self.down_proj(intermediate) - + def replace_mlp_recursive(module, parent_name=""): nonlocal replaced_count for name, child in module.named_children(): full_name = f"{parent_name}.{name}" if parent_name else name - + if any(isinstance(child, mlp_cls) for mlp_cls in mlp_classes): try: optimized = OptimizedMLP(child, kernel_fn) @@ -256,10 +256,10 @@ def replace_mlp_recursive(module, parent_name=""): logger.debug(f"Failed to replace {full_name}: {e}") else: replace_mlp_recursive(child, full_name) - + replace_mlp_recursive(self.model) logger.info(f" Replaced {replaced_count} MLP layers with kernel version") - + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): """Extract valid hidden states (consistent with HF baseline)""" bool_mask = mask.bool() @@ -267,31 +267,31 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result - + @torch.inference_mode() def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: """ Inference method - same interface as Lite encoder - + Args: text: List of text prompts image_list: Optional list of images - + Returns: (prompt_embeds, prompt_embeds_mask, image_info) """ from lightx2v_platform.base.global_var import AI_DEVICE - + template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx - + # Prepare image information if image_list is not None: condition_image_list = [] vae_image_list = [] condition_image_info_list = [] vae_image_info_list = [] - + if self.USE_IMAGE_ID_IN_PROMPT: base_img_prompt = "" img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" @@ -310,7 +310,7 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: vae_image_list.append(vae_image) condition_image_info_list.append(condition_image_info) vae_image_info_list.append(vae_image_info) - + image_info = { "vae_image_list": vae_image_list, "vae_image_info_list": vae_image_info_list, @@ -319,18 +319,18 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: image_info = {} base_img_prompt = "" condition_image_list = None - + # Prepare text and model inputs if self.config["task"] == "i2i" and not self.is_layered and image_list is not None: txt = [template.format(base_img_prompt + e) for e in text] - + model_inputs = self.processor( text=txt, images=condition_image_list, padding=True, return_tensors="pt", ).to(AI_DEVICE) - + encoder_hidden_states = self.model( input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, @@ -340,50 +340,38 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: ) else: txt = [template.format(e) for e in text] - - model_inputs = self.tokenizer( - txt, - max_length=self.tokenizer_max_length + drop_idx, - padding=True, - truncation=True, - return_tensors="pt" - ).to(AI_DEVICE) - + + model_inputs = self.tokenizer(txt, max_length=self.tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt").to(AI_DEVICE) + encoder_hidden_states = self.model( input_ids=model_inputs.input_ids, attention_mask=model_inputs.attention_mask, output_hidden_states=True, ) - + # Post-processing (same as HF baseline) hidden_states = encoder_hidden_states.hidden_states[-1] attention_mask = model_inputs.attention_mask - + split_hidden_states = self._extract_masked_hidden(hidden_states, attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) - - prompt_embeds = torch.stack([ - torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) - for u in split_hidden_states - ]) - encoder_attention_mask = torch.stack([ - torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) - for u in attn_mask_list - ]) - + + prompt_embeds = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]) + encoder_attention_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]) + prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=AI_DEVICE) prompt_embeds_mask = encoder_attention_mask - + _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.view(1, seq_len, -1) prompt_embeds_mask = prompt_embeds_mask.view(1, seq_len) - + logger.info(f"✓ Kernel inference complete: prompt_embeds shape={prompt_embeds.shape}") - + return prompt_embeds, prompt_embeds_mask, image_info - + def _calculate_dimensions(self, target_area, ratio): """Calculate target dimensions""" width = math.sqrt(target_area * ratio) @@ -391,7 +379,7 @@ def _calculate_dimensions(self, target_area, ratio): width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height - + def preprocess_image(self, image): """Preprocess image""" image_width, image_height = image.size @@ -400,16 +388,16 @@ def preprocess_image(self, image): condition_image = self.image_processor.resize(image, condition_height, condition_width) vae_image = self.image_processor.preprocess(image, vae_height, vae_width).unsqueeze(2) return condition_image, vae_image, (condition_height, condition_width), (vae_height, vae_width) - + def offload_to_cpu(self): """Offload model to CPU to free GPU memory""" - if hasattr(self, 'model') and self.model is not None: - self.model.to('cpu') + if hasattr(self, "model") and self.model is not None: + self.model.to("cpu") torch.cuda.empty_cache() logger.debug("Kernel encoder: model offloaded to CPU") - + def reload_to_device(self): """Reload model to GPU""" - if hasattr(self, 'model') and self.model is not None: + if hasattr(self, "model") and self.model is not None: self.model.to(self.device) logger.debug(f"Kernel encoder: model reloaded to {self.device}") diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py index ac07e511..d540e6fc 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py @@ -9,14 +9,15 @@ 4. 硬件优化: 服务端使用专用优化 (Flash Attention 3, CUDA Graph) """ -import os import math +import os +from typing import List, Optional, Tuple + +import numpy as np import requests import torch -import numpy as np -from typing import List, Optional, Tuple -from loguru import logger from PIL import Image +from loguru import logger from transformers import Qwen2Tokenizer, Qwen2VLProcessor try: @@ -33,19 +34,19 @@ class LightLLMServiceTextEncoder: 基于 LightLLM 服务的 Text Encoder 通过 HTTP API 调用 LightLLM 服务获取 hidden states """ - + def __init__(self, config, device=None): from lightx2v_platform.base.global_var import AI_DEVICE - + self.config = config self.device = device if device is not None else AI_DEVICE self.dtype = torch.bfloat16 - + # 配置参数 self.tokenizer_max_length = 1024 self.prompt_template_encode = config["prompt_template_encode"] self.prompt_template_encode_start_idx = config["prompt_template_encode_start_idx"] - + self.CONDITION_IMAGE_SIZE = config.get("CONDITION_IMAGE_SIZE", 384 * 384) self.USE_IMAGE_ID_IN_PROMPT = config.get("USE_IMAGE_ID_IN_PROMPT", True) self.VAE_IMAGE_SIZE = 1024 * 1024 @@ -53,50 +54,43 @@ def __init__(self, config, device=None): if self.is_layered: self.resolution = config.get("resolution", 640) self.VAE_IMAGE_SIZE = self.resolution * self.resolution - + # LightLLM 服务配置 (支持新旧格式) # 新格式: lightllm_config.service_url # 旧格式: lightllm_service_url (向后兼容) - self.service_url = config.get("service_url", - config.get("lightllm_service_url", "http://localhost:8010")) - self.timeout = config.get("service_timeout", - config.get("lightllm_service_timeout", 30)) # 超时时间（秒） - self.retry_times = config.get("service_retry", - config.get("lightllm_service_retry", 3)) # 重试次数 - + self.service_url = config.get("service_url", config.get("lightllm_service_url", "http://localhost:8010")) + self.timeout = config.get("service_timeout", config.get("lightllm_service_timeout", 30)) # 超时时间（秒） + self.retry_times = config.get("service_retry", config.get("lightllm_service_retry", 3)) # 重试次数 + logger.info(f"Initializing LightLLM Service Text Encoder") logger.info(f" Service URL: {self.service_url}") logger.info(f" Timeout: {self.timeout}s") logger.info(f" Device: {self.device}") - + # 加载必要的组件 self.load() - + def load(self): """加载必要的组件（tokenizer, processor, image_processor）""" logger.info("Loading tokenizer and processors...") - + # 加载 tokenizer - tokenizer_path = self.config.get("qwen25vl_tokenizer_path", - os.path.join(self.config["model_path"], "tokenizer")) + tokenizer_path = self.config.get("qwen25vl_tokenizer_path", os.path.join(self.config["model_path"], "tokenizer")) self.tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_path) - + # 加载 processor 和 image processor（用于 i2i 任务） if self.config["task"] == "i2i": if VaeImageProcessor is None: raise ImportError("VaeImageProcessor could not be imported from diffusers") - self.image_processor = VaeImageProcessor( - vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2 - ) - processor_path = self.config.get("qwen25vl_processor_path", - os.path.join(self.config["model_path"], "processor")) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2) + processor_path = self.config.get("qwen25vl_processor_path", os.path.join(self.config["model_path"], "processor")) self.processor = Qwen2VLProcessor.from_pretrained(processor_path) - + logger.info("Tokenizer and processors loaded successfully") - + # 测试服务连接 self._test_service_connection() - + def _test_service_connection(self): """测试与 LightLLM 服务的连接""" try: @@ -109,15 +103,15 @@ def _test_service_connection(self): logger.error(f"✗ Failed to connect to LightLLM service: {e}") logger.error(f" Please ensure the service is running at {self.service_url}") logger.error(f" Start with: python -m lightllm.server.api_server --return_input_hidden_states ...") - + def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) -> dict: """ 调用 LightLLM 服务获取 hidden states - + Args: text: 输入文本 images: 可选的图像列表 - + Returns: 服务返回的 JSON 响应 """ @@ -128,24 +122,23 @@ def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) - "do_sample": False, "return_details": True, # 需要此参数才能返回 hidden_states "max_new_tokens": 1, # LightLLM 要求至少为 1 - } + }, } - + # 如果有图像，需要按照 LightLLM 的 multimodal_params 格式 # 参考: lightllm/server/multimodal_params.py if images is not None and len(images) > 0: import base64 from io import BytesIO - + # 检查 prompt 中的图像 token 数量 # Qwen2-VL 使用 <|image_pad|> 作为图像占位符 image_token_count = text.count("<|image_pad|>") logger.debug(f"Found {image_token_count} image tokens in prompt, have {len(images)} images") - + # 确保图像数量与 prompt 中的图像 token 数量匹配 if image_token_count != len(images): - logger.warning(f"Image token count ({image_token_count}) != image count ({len(images)}), " - f"adjusting to match prompt") + logger.warning(f"Image token count ({image_token_count}) != image count ({len(images)}), adjusting to match prompt") # 如果 prompt 中有多个图像 token，但只提供了 1 个图像，重复使用该图像 if len(images) == 1 and image_token_count > 1: images = images * image_token_count @@ -153,7 +146,7 @@ def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) - elif image_token_count == 0: logger.warning("No image tokens found in prompt, skipping image transmission") images = [] - + # LightLLM 期望的格式：multimodal_params.images 是 ImageItem 列表 # 每个 ImageItem 需要 {"type": "base64", "data": "base64_string"} # 优化：使用 JPEG 格式（编码更快，传输更小） @@ -167,100 +160,98 @@ def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) - img = img.convert("RGB") img.save(buffered, format="JPEG", quality=95, optimize=False) img_str = base64.b64encode(buffered.getvalue()).decode() - image_items.append({ - "type": "base64", - "data": img_str - }) - logger.debug(f"Encoded image {idx+1}/{len(images)}: {len(img_str)} chars base64 (JPEG)") - + image_items.append({"type": "base64", "data": img_str}) + logger.debug(f"Encoded image {idx + 1}/{len(images)}: {len(img_str)} chars base64 (JPEG)") + # 使用 multimodal_params 格式 if len(image_items) > 0: - payload["multimodal_params"] = { - "images": image_items - } + payload["multimodal_params"] = {"images": image_items} logger.debug(f"Added {len(image_items)} images to multimodal_params") - + # 尝试使用更快的 JSON 库 try: import orjson + def fast_json_loads(data): return orjson.loads(data) + logger.debug("Using orjson for fast JSON parsing") except ImportError: try: import ujson + def fast_json_loads(data): return ujson.loads(data) + logger.debug("Using ujson for JSON parsing") except ImportError: import json + def fast_json_loads(data): return json.loads(data) + logger.debug("Using standard json for JSON parsing") - + # 重试机制 last_error = None for attempt in range(self.retry_times): try: logger.debug(f"Calling LightLLM service (attempt {attempt + 1}/{self.retry_times})...") - response = requests.post( - f"{self.service_url}/generate", - json=payload, - timeout=self.timeout - ) + response = requests.post(f"{self.service_url}/generate", json=payload, timeout=self.timeout) response.raise_for_status() - + # 使用更快的 JSON 库解析响应 result = fast_json_loads(response.content) logger.debug(f"✓ Service call successful") return result - + except requests.exceptions.Timeout: last_error = f"Request timeout after {self.timeout}s" logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") except requests.exceptions.RequestException as e: last_error = str(e) # 记录详细的错误信息 - if hasattr(e, 'response') and e.response is not None: + if hasattr(e, "response") and e.response is not None: try: error_detail = e.response.json() logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") logger.debug(f" Error detail: {error_detail}") - except: + except Exception: logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") logger.debug(f" Response text: {e.response.text[:200] if hasattr(e.response, 'text') else 'N/A'}") else: logger.warning(f"⚠ Attempt {attempt + 1} failed: {last_error}") - + if attempt < self.retry_times - 1: import time + time.sleep(1) # 重试前等待1秒 - + # 所有重试都失败 raise RuntimeError(f"Failed to call LightLLM service after {self.retry_times} attempts. Last error: {last_error}") - + @torch.inference_mode() def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: """ 推理方法 - 调用 LightLLM 服务获取 hidden states - + Args: text: 文本提示列表 image_list: 可选的图像列表（用于 i2i 任务） - + Returns: (prompt_embeds, prompt_embeds_mask, image_info) """ template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx - + # 准备图像信息 if image_list is not None: condition_image_list = [] vae_image_list = [] condition_image_info_list = [] vae_image_info_list = [] - + if self.USE_IMAGE_ID_IN_PROMPT: base_img_prompt = "" img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" @@ -279,7 +270,7 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: vae_image_list.append(vae_image) condition_image_info_list.append(condition_image_info) vae_image_info_list.append(vae_image_info) - + image_info = { "vae_image_list": vae_image_list, "vae_image_info_list": vae_image_info_list, @@ -288,33 +279,34 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: image_info = {} base_img_prompt = "" condition_image_list = None - + # 准备文本 if self.config["task"] == "i2i" and not self.is_layered and image_list is not None: txt = template.format(base_img_prompt + text[0]) else: txt = template.format(text[0]) - + # 调用 LightLLM 服务 logger.debug(f"Calling LightLLM service with text: {txt[:100]}...") logger.debug(f" Image count: {len(condition_image_list) if condition_image_list else 0}") logger.debug(f" Base image prompt: {base_img_prompt[:50] if base_img_prompt else 'None'}...") result = self._call_service(txt, condition_image_list) - + # 解析返回的 hidden states if "hidden_states_base64" in result: import base64 + # Decode base64 to bytes data_bytes = base64.b64decode(result["hidden_states_base64"]) shape = result["hidden_states_shape"] # Create numpy array from buffer (zero copy if possible, but base64 decode creates bytes) hidden_states_np = np.frombuffer(data_bytes, dtype=np.uint8).reshape(shape) hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) - + elif "hidden_states" in result: # Legacy path hidden_states_data = result["hidden_states"] - + # 优化：根据数据类型选择最快的转换方式 if isinstance(hidden_states_data, list): # 列表格式：检查是否是扁平列表或嵌套列表 @@ -326,7 +318,7 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: try: # 尝试使用 memoryview 加速 hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) - except: + except Exception: hidden_states_np = np.array(hidden_states_data, dtype=np.uint8) hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) elif isinstance(hidden_states_data, np.ndarray): @@ -340,34 +332,33 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) else: raise ValueError(f"LightLLM service response missing 'hidden_states' or 'hidden_states_base64'. Response keys: {result.keys()}") - + # 关键步骤：将 uint8 tensor 通过 view 转换为 bfloat16 hidden_states = hidden_states.view(torch.bfloat16) - - logger.debug(f"Converted hidden states: shape={hidden_states.shape}, dtype={hidden_states.dtype}, " - f"device={hidden_states.device}") - + + logger.debug(f"Converted hidden states: shape={hidden_states.shape}, dtype={hidden_states.dtype}, device={hidden_states.device}") + # 后处理：去除 drop_idx 和调整形状 # 假设 hidden_states 形状为 [batch, seq_len, hidden_dim] if len(hidden_states.shape) == 2: # 如果是 [seq_len, hidden_dim]，添加 batch 维度 hidden_states = hidden_states.unsqueeze(0) - + # 去除 prompt template 的前缀 tokens if drop_idx > 0 and hidden_states.shape[1] > drop_idx: hidden_states = hidden_states[:, drop_idx:, :] - + # 创建 attention mask seq_len = hidden_states.shape[1] attention_mask = torch.ones(hidden_states.shape[0], seq_len, dtype=torch.long, device=self.device) - + prompt_embeds = hidden_states prompt_embeds_mask = attention_mask - + logger.info(f"✓ LightLLM service inference complete: prompt_embeds shape={prompt_embeds.shape}") - + return prompt_embeds, prompt_embeds_mask, image_info - + def _calculate_dimensions(self, target_area, ratio): """计算目标尺寸""" width = math.sqrt(target_area * ratio) @@ -375,37 +366,37 @@ def _calculate_dimensions(self, target_area, ratio): width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height - + def preprocess_image(self, image): """预处理图像 (带简单缓存)""" # 使用 image id 作为缓存键 (假设同一对象内容不变，适用于 diffusers pipeline 循环) img_id = id(image) if hasattr(self, "_image_cache") and img_id in self._image_cache: return self._image_cache[img_id] - + image_width, image_height = image.size condition_width, condition_height = self._calculate_dimensions(self.CONDITION_IMAGE_SIZE, image_width / image_height) vae_width, vae_height = self._calculate_dimensions(self.VAE_IMAGE_SIZE, image_width / image_height) condition_image = self.image_processor.resize(image, condition_height, condition_width) vae_image = self.image_processor.preprocess(image, vae_height, vae_width).unsqueeze(2) - + result = (condition_image, vae_image, (condition_height, condition_width), (vae_height, vae_width)) - + # 初始化缓存 (如果不存在) if not hasattr(self, "_image_cache"): self._image_cache = {} - + # 简单缓存管理：如果太大则清空 if len(self._image_cache) > 50: self._image_cache.clear() - + self._image_cache[img_id] = result return result - + def offload_to_cpu(self): """服务化版本无需 offload""" logger.debug("Service-based encoder: offload_to_cpu() is a no-op") - + def reload_to_device(self): """服务化版本无需 reload""" logger.debug("Service-based encoder: reload_to_device() is a no-op") diff --git a/lightx2v/models/runners/qwen_image/qwen_image_runner.py b/lightx2v/models/runners/qwen_image/qwen_image_runner.py index b15a3ddb..bbbaa2db 100755 --- a/lightx2v/models/runners/qwen_image/qwen_image_runner.py +++ b/lightx2v/models/runners/qwen_image/qwen_image_runner.py @@ -65,10 +65,10 @@ def __init__(self, config): if self.is_layered: self.layers = self.config.get("layers", 4) self.resolution = self.config.get("resolution", 1024) - + # Text encoder type: "lightllm_service", "lightllm_kernel", or default (baseline) self.text_encoder_type = config.get("text_encoder_type", "baseline") - + if self.text_encoder_type in ["lightllm_service", "lightllm_kernel"]: logger.info(f"Using LightLLM text encoder: {self.text_encoder_type}") @@ -91,7 +91,7 @@ def load_transformer(self): def load_text_encoder(self): """Load text encoder based on text_encoder_type configuration. - + Supported types: - "lightllm_service": LightLLM HTTP service mode - "lightllm_kernel": HuggingFace model with Triton kernel optimizations @@ -101,19 +101,21 @@ def load_text_encoder(self): encoder_config = self.config.copy() lightllm_config = self.config.get("lightllm_config", {}) encoder_config.update(lightllm_config) - + if self.text_encoder_type == "lightllm_service": from lightx2v.models.input_encoders.lightllm import LightLLMServiceTextEncoder + logger.info("Loading LightLLM service-based text encoder") text_encoder = LightLLMServiceTextEncoder(encoder_config) elif self.text_encoder_type == "lightllm_kernel": from lightx2v.models.input_encoders.lightllm import LightLLMKernelTextEncoder + logger.info("Loading LightLLM Kernel-optimized text encoder") text_encoder = LightLLMKernelTextEncoder(encoder_config) else: # baseline or default logger.info("Loading HuggingFace baseline text encoder") text_encoder = Qwen25_VLForConditionalGeneration_TextEncoder(self.config) - + text_encoders = [text_encoder] return text_encoders From 62011fced682fac749b198cc180223f1bd4d59cf Mon Sep 17 00:00:00 2001 From: fuheaven Date: Mon, 19 Jan 2026 17:55:30 +0800 Subject: [PATCH 3/6] add lightllm triton kernel --- .../lightllm/qwen25_text_encoder_kernel.py | 4 +- lightx2v/utils/triton_kernels/__init__.py | 0 lightx2v/utils/triton_kernels/rmsnorm.py | 78 ++++++++++++++++++ lightx2v/utils/triton_kernels/silu_and_mul.py | 81 +++++++++++++++++++ 4 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 lightx2v/utils/triton_kernels/__init__.py create mode 100644 lightx2v/utils/triton_kernels/rmsnorm.py create mode 100644 lightx2v/utils/triton_kernels/silu_and_mul.py diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py index cb43040f..fbe3d753 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py @@ -138,14 +138,14 @@ def _apply_kernel_optimizations(self): try: if self.use_rmsnorm_kernel: - from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward + from lightx2v.utils.triton_kernels.rmsnorm import rmsnorm_forward self._rmsnorm_kernel = rmsnorm_forward self._replace_rmsnorm_with_kernel() logger.info(" ✓ RMSNorm kernel integrated") if self.use_ffn_kernel: - from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd + from lightx2v.utils.triton_kernels.silu_and_mul import silu_and_mul_fwd self._silu_mul_kernel = silu_and_mul_fwd self._replace_ffn_with_kernel() diff --git a/lightx2v/utils/triton_kernels/__init__.py b/lightx2v/utils/triton_kernels/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/lightx2v/utils/triton_kernels/rmsnorm.py b/lightx2v/utils/triton_kernels/rmsnorm.py new file mode 100644 index 00000000..5331688a --- /dev/null +++ b/lightx2v/utils/triton_kernels/rmsnorm.py @@ -0,0 +1,78 @@ +import torch + +import triton +import triton.language as tl +import os + +rmsnorm_num_warps = int(os.getenv("RMSNORM_WARPS", "8")) + + +@triton.jit +def _rms_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + x_stride0, # how much to increase the pointer when moving by 1 row + x_stride1, + y_stride0, + y_stride1, + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * y_stride0 + X += row * x_stride0 + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols * x_stride1, mask=cols < N, other=0.0).to(tl.float32) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + x_hat = x * rstd + y = x_hat * w + # Write output + tl.store(Y + cols * y_stride1, y.to(Y.dtype.element_ty), mask=mask) + + +def rmsnorm_forward(x: torch.Tensor, weight, eps, out=None): + # allocate output + y = torch.empty_like(x) if out is None else out + # reshape input data into 2D tensor + x_arg = x.view(-1, x.shape[-1]) + y_arg = y.view(-1, x.shape[-1]) + assert y.data_ptr() == y_arg.data_ptr() + M, N = x_arg.shape + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + # print("BLOCK_SIZE:", BLOCK_SIZE) + if N > BLOCK_SIZE: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + if BLOCK_SIZE > 16384: + BLOCK_SIZE = 16384 + # enqueue kernel + _rms_norm_fwd_fused[(M,)]( + x_arg, + y_arg, + weight, + x_arg.stride(0), + x_arg.stride(1), + y_arg.stride(0), + y_arg.stride(1), + N, + eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=rmsnorm_num_warps, + ) + return y diff --git a/lightx2v/utils/triton_kernels/silu_and_mul.py b/lightx2v/utils/triton_kernels/silu_and_mul.py new file mode 100644 index 00000000..c3c2bde3 --- /dev/null +++ b/lightx2v/utils/triton_kernels/silu_and_mul.py @@ -0,0 +1,81 @@ +import torch + +import triton +import triton.language as tl + + +@triton.jit +def _silu_and_mul_kernel( + input_ptr, + output_ptr, + stride_input_m, + stride_input_n, + stride_output_m, + stride_output_n, + size_m, + size_n, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) + stride_output_m = tl.cast(stride_output_m, dtype=tl.int64) + + tid = tl.program_id(0) + input_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M) + output_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M) + + pid = tl.program_id(1) + input_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N) + output_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N) + + up_offsets = input_m_offsets[:, None] * stride_input_m + (input_n_offsets[None, :] + size_n) * stride_input_n + gate_offsets = input_m_offsets[:, None] * stride_input_m + input_n_offsets[None, :] * stride_input_n + res_offsets = output_m_offsets[:, None] * stride_output_m + output_n_offsets[None, :] * stride_output_n + + up = tl.load( + input_ptr + up_offsets, + mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None], + other=0.0, + ) + gate = tl.load( + input_ptr + gate_offsets, + mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None], + other=0.0, + ).to(tl.float32) + + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + + tl.store( + output_ptr + res_offsets, + up * gate, + mask=(output_n_offsets < size_n)[None, :] * (output_m_offsets < size_m)[:, None], + ) + + +def silu_and_mul_fwd(input: torch.Tensor, output): + stride_input_m = input.stride(0) + stride_input_n = input.stride(1) + stride_output_m = output.stride(0) + stride_output_n = output.stride(1) + size_m = input.shape[0] + size_n = input.shape[-1] // 2 + BLOCK_M = 128 + BLOCK_N = 128 + grid = ( + triton.cdiv(size_m, BLOCK_M), + triton.cdiv(size_n, BLOCK_N), + ) + _silu_and_mul_kernel[grid]( + input, + output, + stride_input_m, + stride_input_n, + stride_output_m, + stride_output_n, + size_m, + size_n, + BLOCK_M, + BLOCK_N, + ) + return From e683bba5e35165edffc920c175bce9c440b3b773 Mon Sep 17 00:00:00 2001 From: fuheaven Date: Tue, 20 Jan 2026 12:12:22 +0800 Subject: [PATCH 4/6] optimize kernel mode, add sgl_kernel rmsnorm --- .../lightllm/qwen25_text_encoder_kernel.py | 111 +++--------------- lightx2v/utils/triton_kernels/__init__.py | 0 lightx2v/utils/triton_kernels/rmsnorm.py | 78 ------------ lightx2v/utils/triton_kernels/silu_and_mul.py | 81 ------------- 4 files changed, 18 insertions(+), 252 deletions(-) delete mode 100644 lightx2v/utils/triton_kernels/__init__.py delete mode 100644 lightx2v/utils/triton_kernels/rmsnorm.py delete mode 100644 lightx2v/utils/triton_kernels/silu_and_mul.py diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py index fbe3d753..51908675 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py @@ -1,16 +1,12 @@ """ -LightLLM Kernel-Optimized Text Encoder - -Hybrid approach that uses HuggingFace model structure with selectively replaced -LightLLM Triton kernels for maximum performance while maintaining precision. +Kernel-Optimized Text Encoder Key optimizations: 1. Flash Attention (No-Padding) - ~40% of inference time 2. Fused RMSNorm - frequent operation -3. Fused SiLU+Mul - FFN activation Performance target: -- Speed: 1.29x faster than Baseline (73ms vs 94.5ms) +- Speed: 1.13x faster than Baseline (81.23ms vs 92.23ms) - Precision: >0.99 cosine similarity - Memory: Similar to Lite (~125MB VRAM) @@ -70,14 +66,12 @@ def __init__(self, config: Dict[str, Any], device: Optional[str] = None): # Kernel optimization flags self.use_flash_attention_kernel = config.get("use_flash_attention_kernel", True) self.use_rmsnorm_kernel = config.get("use_rmsnorm_kernel", True) - self.use_ffn_kernel = config.get("use_ffn_kernel", True) - logger.info(f"Initializing LightLLM Kernel-Optimized Text Encoder") + logger.info(f"Initializing Kernel-Optimized Text Encoder") logger.info(f" Model Path: {self.model_path}") logger.info(f" Device: {self.device}") logger.info(f" Flash Attention: {self.use_flash_attention_kernel}") logger.info(f" RMSNorm Kernel: {self.use_rmsnorm_kernel}") - logger.info(f" FFN Kernel: {self.use_ffn_kernel}") self.load() @@ -123,38 +117,29 @@ def load(self): logger.info(f" ✓ Model loaded with {attn_impl}") - # 4. Apply kernel optimizations (RMSNorm, RoPE, FFN) + # 4. Apply kernel optimizations (RMSNorm) self._apply_kernel_optimizations() self._is_loaded = True def _apply_kernel_optimizations(self): - """Apply LightLLM kernel optimizations to the model""" + """Apply kernel optimizations to the model""" logger.info("Applying kernel optimizations...") # Flash Attention is already loaded with the model if self.use_flash_attention_kernel: logger.info(" ✓ Flash Attention 2 (loaded with model)") - try: if self.use_rmsnorm_kernel: - from lightx2v.utils.triton_kernels.rmsnorm import rmsnorm_forward - - self._rmsnorm_kernel = rmsnorm_forward - self._replace_rmsnorm_with_kernel() - logger.info(" ✓ RMSNorm kernel integrated") - - if self.use_ffn_kernel: - from lightx2v.utils.triton_kernels.silu_and_mul import silu_and_mul_fwd + try: + from sgl_kernel.elementwise import rmsnorm - self._silu_mul_kernel = silu_and_mul_fwd - self._replace_ffn_with_kernel() - logger.info(" ✓ FFN kernel integrated") - - except ImportError as e: - logger.warning(f"Failed to import LightLLM kernels: {e}") - self.use_rmsnorm_kernel = False - self.use_ffn_kernel = False + self._rmsnorm_kernel = rmsnorm + self._replace_rmsnorm_with_kernel() + logger.info(" ✓ RMSNorm kernel integrated (from sgl_kernel)") + except ImportError as e: + logger.warning(f" ✗ Failed to import sgl_kernel: {e}. RMSNorm optimization disabled.") + self.use_rmsnorm_kernel = False def _replace_rmsnorm_with_kernel(self): """Replace RMSNorm layers with fused kernel""" @@ -176,7 +161,11 @@ def __init__(self, original_norm, kernel_fn): self.kernel_fn = kernel_fn def forward(self, hidden_states): - return self.kernel_fn(hidden_states, self.weight, self.variance_epsilon) + orig_shape = hidden_states.shape + # Reshape to (-1, hidden_dim) as sgl_kernel expects 2D + x_2d = hidden_states.view(-1, orig_shape[-1]) + out_2d = self.kernel_fn(x_2d, self.weight, self.variance_epsilon) + return out_2d.view(orig_shape) # Replace all RMSNorm layers def replace_rmsnorm_recursive(module, parent_name=""): @@ -196,70 +185,6 @@ def replace_rmsnorm_recursive(module, parent_name=""): replace_rmsnorm_recursive(self.model) logger.info(f" Replaced {replaced_count} RMSNorm layers with kernel version") - def _replace_ffn_with_kernel(self): - """Replace FFN activation with fused SiLU+Mul kernel""" - # Fix: Use correct MLP classes for Qwen2.5-VL model - mlp_classes = [] - - # Qwen2MLP (for text model layers) - try: - from transformers.models.qwen2.modeling_qwen2 import Qwen2MLP - - mlp_classes.append(Qwen2MLP) - except ImportError: - pass - - # Qwen2_5_VLMLP (for visual encoder layers) - THIS IS THE KEY FIX! - try: - from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLMLP - - mlp_classes.append(Qwen2_5_VLMLP) - except ImportError: - pass - - if not mlp_classes: - logger.warning("Could not import any MLP class, skipping FFN optimization") - return - - logger.info(f" Detecting MLP classes: {[c.__name__ for c in mlp_classes]}") - - replaced_count = 0 - kernel_fn = self._silu_mul_kernel - - class OptimizedMLP(nn.Module): - def __init__(self, original_mlp, kernel_fn): - super().__init__() - self.gate_proj = original_mlp.gate_proj - self.up_proj = original_mlp.up_proj - self.down_proj = original_mlp.down_proj - self.kernel_fn = kernel_fn - - def forward(self, hidden_states): - gate = self.gate_proj(hidden_states) - up = self.up_proj(hidden_states) - gate_up = torch.cat([gate, up], dim=-1) - intermediate = torch.empty_like(gate) - self.kernel_fn(gate_up, intermediate) - return self.down_proj(intermediate) - - def replace_mlp_recursive(module, parent_name=""): - nonlocal replaced_count - for name, child in module.named_children(): - full_name = f"{parent_name}.{name}" if parent_name else name - - if any(isinstance(child, mlp_cls) for mlp_cls in mlp_classes): - try: - optimized = OptimizedMLP(child, kernel_fn) - setattr(module, name, optimized) - replaced_count += 1 - except Exception as e: - logger.debug(f"Failed to replace {full_name}: {e}") - else: - replace_mlp_recursive(child, full_name) - - replace_mlp_recursive(self.model) - logger.info(f" Replaced {replaced_count} MLP layers with kernel version") - def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): """Extract valid hidden states (consistent with HF baseline)""" bool_mask = mask.bool() diff --git a/lightx2v/utils/triton_kernels/__init__.py b/lightx2v/utils/triton_kernels/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/lightx2v/utils/triton_kernels/rmsnorm.py b/lightx2v/utils/triton_kernels/rmsnorm.py deleted file mode 100644 index 5331688a..00000000 --- a/lightx2v/utils/triton_kernels/rmsnorm.py +++ /dev/null @@ -1,78 +0,0 @@ -import torch - -import triton -import triton.language as tl -import os - -rmsnorm_num_warps = int(os.getenv("RMSNORM_WARPS", "8")) - - -@triton.jit -def _rms_norm_fwd_fused( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - x_stride0, # how much to increase the pointer when moving by 1 row - x_stride1, - y_stride0, - y_stride1, - N, # number of columns in X - eps, # epsilon to avoid division by zero - BLOCK_SIZE: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - Y += row * y_stride0 - X += row * x_stride0 - # Compute variance - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - x = tl.load(X + cols * x_stride1, mask=cols < N, other=0.0).to(tl.float32) - _var += x * x - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - # Normalize and apply linear transformation - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - w = tl.load(W + cols, mask=mask).to(tl.float32) - x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) - x_hat = x * rstd - y = x_hat * w - # Write output - tl.store(Y + cols * y_stride1, y.to(Y.dtype.element_ty), mask=mask) - - -def rmsnorm_forward(x: torch.Tensor, weight, eps, out=None): - # allocate output - y = torch.empty_like(x) if out is None else out - # reshape input data into 2D tensor - x_arg = x.view(-1, x.shape[-1]) - y_arg = y.view(-1, x.shape[-1]) - assert y.data_ptr() == y_arg.data_ptr() - M, N = x_arg.shape - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - # print("BLOCK_SIZE:", BLOCK_SIZE) - if N > BLOCK_SIZE: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - if BLOCK_SIZE > 16384: - BLOCK_SIZE = 16384 - # enqueue kernel - _rms_norm_fwd_fused[(M,)]( - x_arg, - y_arg, - weight, - x_arg.stride(0), - x_arg.stride(1), - y_arg.stride(0), - y_arg.stride(1), - N, - eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=rmsnorm_num_warps, - ) - return y diff --git a/lightx2v/utils/triton_kernels/silu_and_mul.py b/lightx2v/utils/triton_kernels/silu_and_mul.py deleted file mode 100644 index c3c2bde3..00000000 --- a/lightx2v/utils/triton_kernels/silu_and_mul.py +++ /dev/null @@ -1,81 +0,0 @@ -import torch - -import triton -import triton.language as tl - - -@triton.jit -def _silu_and_mul_kernel( - input_ptr, - output_ptr, - stride_input_m, - stride_input_n, - stride_output_m, - stride_output_n, - size_m, - size_n, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) - stride_output_m = tl.cast(stride_output_m, dtype=tl.int64) - - tid = tl.program_id(0) - input_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M) - output_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M) - - pid = tl.program_id(1) - input_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N) - output_n_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N) - - up_offsets = input_m_offsets[:, None] * stride_input_m + (input_n_offsets[None, :] + size_n) * stride_input_n - gate_offsets = input_m_offsets[:, None] * stride_input_m + input_n_offsets[None, :] * stride_input_n - res_offsets = output_m_offsets[:, None] * stride_output_m + output_n_offsets[None, :] * stride_output_n - - up = tl.load( - input_ptr + up_offsets, - mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None], - other=0.0, - ) - gate = tl.load( - input_ptr + gate_offsets, - mask=(input_n_offsets < size_n)[None, :] * (input_m_offsets < size_m)[:, None], - other=0.0, - ).to(tl.float32) - - gate = gate / (1 + tl.exp(-gate)) - gate = gate.to(input_ptr.dtype.element_ty) - - tl.store( - output_ptr + res_offsets, - up * gate, - mask=(output_n_offsets < size_n)[None, :] * (output_m_offsets < size_m)[:, None], - ) - - -def silu_and_mul_fwd(input: torch.Tensor, output): - stride_input_m = input.stride(0) - stride_input_n = input.stride(1) - stride_output_m = output.stride(0) - stride_output_n = output.stride(1) - size_m = input.shape[0] - size_n = input.shape[-1] // 2 - BLOCK_M = 128 - BLOCK_N = 128 - grid = ( - triton.cdiv(size_m, BLOCK_M), - triton.cdiv(size_n, BLOCK_N), - ) - _silu_and_mul_kernel[grid]( - input, - output, - stride_input_m, - stride_input_n, - stride_output_m, - stride_output_n, - size_m, - size_n, - BLOCK_M, - BLOCK_N, - ) - return From 4d395102ac3ff837796fbec5b345d60960e231b8 Mon Sep 17 00:00:00 2001 From: fuheaven Date: Tue, 27 Jan 2026 11:34:02 +0800 Subject: [PATCH 5/6] add shm mode for service, update config and scripts --- .../qwen_image_i2i_2511_kernel.json | 17 +++ .../qwen_image_i2i_2511_service.json | 19 +++ .../lightllm/qwen25_text_encoder_kernel.py | 13 +- .../lightllm/qwen25_text_encoder_service.py | 54 ++++++--- .../input_encoders/lightllm/shm_client.py | 111 ++++++++++++++++++ .../qwen_image/qwen_image_i2i_2511_kernel.sh | 21 ++++ .../qwen_image/qwen_image_i2i_2511_service.sh | 21 ++++ 7 files changed, 226 insertions(+), 30 deletions(-) create mode 100644 configs/qwen_image/qwen_image_i2i_2511_kernel.json create mode 100644 configs/qwen_image/qwen_image_i2i_2511_service.json create mode 100644 lightx2v/models/input_encoders/lightllm/shm_client.py create mode 100755 scripts/qwen_image/qwen_image_i2i_2511_kernel.sh create mode 100755 scripts/qwen_image/qwen_image_i2i_2511_service.sh diff --git a/configs/qwen_image/qwen_image_i2i_2511_kernel.json b/configs/qwen_image/qwen_image_i2i_2511_kernel.json new file mode 100644 index 00000000..780e2f00 --- /dev/null +++ b/configs/qwen_image/qwen_image_i2i_2511_kernel.json @@ -0,0 +1,17 @@ +{ + "infer_steps": 40, + "prompt_template_encode": "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + "prompt_template_encode_start_idx": 64, + "resize_mode": "adaptive", + "attn_type": "flash_attn3", + "enable_cfg": true, + "sample_guide_scale": 4.0, + "vae_scale_factor": 8, + "CONDITION_IMAGE_SIZE": 147456, + "USE_IMAGE_ID_IN_PROMPT": true, + "text_encoder_type": "lightllm_kernel", + "lightllm_config": { + "use_flash_attention_kernel": true, + "use_rmsnorm_kernel": true + } +} \ No newline at end of file diff --git a/configs/qwen_image/qwen_image_i2i_2511_service.json b/configs/qwen_image/qwen_image_i2i_2511_service.json new file mode 100644 index 00000000..96fa9eb1 --- /dev/null +++ b/configs/qwen_image/qwen_image_i2i_2511_service.json @@ -0,0 +1,19 @@ +{ + "infer_steps": 40, + "prompt_template_encode": "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + "prompt_template_encode_start_idx": 64, + "resize_mode": "adaptive", + "attn_type": "flash_attn3", + "enable_cfg": true, + "sample_guide_scale": 4.0, + "vae_scale_factor": 8, + "CONDITION_IMAGE_SIZE": 147456, + "USE_IMAGE_ID_IN_PROMPT": true, + "text_encoder_type": "lightllm_service", + "lightllm_config": { + "service_url": "http://localhost:8010", + "service_timeout": 30, + "service_retry": 3, + "use_shm": true + } +} \ No newline at end of file diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py index 51908675..9d1154e9 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_kernel.py @@ -2,17 +2,8 @@ Kernel-Optimized Text Encoder Key optimizations: -1. Flash Attention (No-Padding) - ~40% of inference time -2. Fused RMSNorm - frequent operation - -Performance target: -- Speed: 1.13x faster than Baseline (81.23ms vs 92.23ms) -- Precision: >0.99 cosine similarity -- Memory: Similar to Lite (~125MB VRAM) - -Usage: - encoder = LightLLMKernelTextEncoder(config) - hidden_states, mask, image_info = encoder.infer(text, image_list) +1. Flash Attention +2. Fused RMSNorm - frequent operation (sgl_kernel version) """ import math diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py index d540e6fc..968d5fa4 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py @@ -1,12 +1,6 @@ """ LightLLM Service-based Text Encoder 使用 LightLLM 服务提供的 hidden states 作为 text encoder 输出 - -优势: -1. 显存节省: 本地不需要加载 text encoder (~7-8GB) -2. 计算卸载: text encoder 推理在服务端完成 -3. 服务复用: 多个实例可共享同一服务 -4. 硬件优化: 服务端使用专用优化 (Flash Attention 3, CUDA Graph) """ import math @@ -55,17 +49,19 @@ def __init__(self, config, device=None): self.resolution = config.get("resolution", 640) self.VAE_IMAGE_SIZE = self.resolution * self.resolution - # LightLLM 服务配置 (支持新旧格式) - # 新格式: lightllm_config.service_url - # 旧格式: lightllm_service_url (向后兼容) - self.service_url = config.get("service_url", config.get("lightllm_service_url", "http://localhost:8010")) - self.timeout = config.get("service_timeout", config.get("lightllm_service_timeout", 30)) # 超时时间（秒） - self.retry_times = config.get("service_retry", config.get("lightllm_service_retry", 3)) # 重试次数 + # LightLLM 服务配置 + self.service_url = config.get("service_url", "http://localhost:8010") + self.timeout = config.get("service_timeout", 30) # 超时时间（秒） + self.retry_times = config.get("service_retry", 3) # 重试次数 + + # Shared Memory 模式配置（默认开启，仅在同机部署时有效） + self.use_shm = config.get("use_shm", True) logger.info(f"Initializing LightLLM Service Text Encoder") logger.info(f" Service URL: {self.service_url}") logger.info(f" Timeout: {self.timeout}s") logger.info(f" Device: {self.device}") + logger.info(f" Use Shared Memory: {self.use_shm}") # 加载必要的组件 self.load() @@ -153,12 +149,9 @@ def _call_service(self, text: str, images: Optional[List[Image.Image]] = None) - image_items = [] for idx, img in enumerate(images): buffered = BytesIO() - # JPEG 编码比 PNG 快 3-5x，且文件更小 - # 使用高质量 (95) 以保持图像质量 - if img.mode == "RGBA": - # JPEG 不支持透明通道，转换为 RGB - img = img.convert("RGB") - img.save(buffered, format="JPEG", quality=95, optimize=False) + # BMP 格式：无损且编码极快 (也就是直接内存拷贝)，适合 Localhost 高带宽场景 + # 相比 PNG (CPU 压缩慢) 和 JPEG (有损)，BMP 是由于 Service Mode 的最佳选择 + img.save(buffered, format="BMP") img_str = base64.b64encode(buffered.getvalue()).decode() image_items.append({"type": "base64", "data": img_str}) logger.debug(f"Encoded image {idx + 1}/{len(images)}: {len(img_str)} chars base64 (JPEG)") @@ -293,7 +286,30 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: result = self._call_service(txt, condition_image_list) # 解析返回的 hidden states - if "hidden_states_base64" in result: + # 优先使用 Shared Memory 模式（零拷贝，最快） + if self.use_shm and "shm_hidden_states_name" in result: + from .shm_client import get_shm_client + + shm_name = result["shm_hidden_states_name"] + shape = tuple(result["shm_hidden_states_shape"]) + try: + shm_client = get_shm_client() + hidden_states_np = shm_client.read_hidden_states(shm_name, shape, dtype=np.uint8) + hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) + logger.debug(f"✓ Read hidden states from shared memory: {shm_name}") + except Exception as e: + logger.warning(f"Failed to read from shared memory '{shm_name}': {e}, falling back to HTTP mode") + # Fallback to base64 mode + if "hidden_states_base64" in result: + import base64 + data_bytes = base64.b64decode(result["hidden_states_base64"]) + shape = result["hidden_states_shape"] + hidden_states_np = np.frombuffer(data_bytes, dtype=np.uint8).reshape(shape) + hidden_states = torch.from_numpy(hidden_states_np).to(device=self.device, non_blocking=True) + else: + raise + + elif "hidden_states_base64" in result: import base64 # Decode base64 to bytes diff --git a/lightx2v/models/input_encoders/lightllm/shm_client.py b/lightx2v/models/input_encoders/lightllm/shm_client.py new file mode 100644 index 00000000..43e3eaed --- /dev/null +++ b/lightx2v/models/input_encoders/lightllm/shm_client.py @@ -0,0 +1,111 @@ +""" +Shared Memory Client for LightLLM Hidden States + +支持从 LightLLM 服务的共享内存中直接读取 hidden states， +实现零拷贝数据传输，显著降低通信延迟。 +""" + +import numpy as np +from multiprocessing import shared_memory +from typing import Tuple, Optional +from loguru import logger + + +class ShmClient: + """共享内存客户端，用于读取 LightLLM 服务的 hidden states""" + + def __init__(self): + self._cache = {} # 缓存已打开的共享内存对象 + + def read_hidden_states( + self, + shm_name: str, + shape: Tuple[int, ...], + dtype: np.dtype = np.uint8, + ) -> np.ndarray: + """ + 从共享内存读取 hidden states 数据 + + Args: + shm_name: 共享内存名称 + shape: 数据形状 + dtype: 数据类型（默认 uint8，需要后续 view 为 bfloat16） + + Returns: + numpy 数组（数据的副本，可安全使用） + """ + try: + # 打开共享内存 + shm = shared_memory.SharedMemory(name=shm_name) + + # 创建 numpy 数组视图 + arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf) + + # 复制数据（确保数据独立，不依赖共享内存生命周期） + result = arr.copy() + + # 关闭共享内存（不 unlink，因为服务端负责管理生命周期） + shm.close() + + logger.debug(f"Read hidden states from shm '{shm_name}': shape={shape}") + return result + + except FileNotFoundError: + logger.error(f"Shared memory '{shm_name}' not found") + raise + except Exception as e: + logger.error(f"Failed to read from shared memory '{shm_name}': {e}") + raise + + def read_hidden_states_zero_copy( + self, + shm_name: str, + shape: Tuple[int, ...], + dtype: np.dtype = np.uint8, + ) -> Tuple[np.ndarray, shared_memory.SharedMemory]: + """ + 从共享内存读取 hidden states 数据（零拷贝模式） + + 注意：此模式返回的数组直接引用共享内存，调用者需要负责： + 1. 在使用完数据后调用 shm.close() + 2. 不要在共享内存关闭后继续使用数组 + + Args: + shm_name: 共享内存名称 + shape: 数据形状 + dtype: 数据类型 + + Returns: + (numpy 数组, SharedMemory 对象) - 调用者需要管理 shm 对象的生命周期 + """ + try: + shm = shared_memory.SharedMemory(name=shm_name) + arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf) + logger.debug(f"Zero-copy read from shm '{shm_name}': shape={shape}") + return arr, shm + except Exception as e: + logger.error(f"Failed to zero-copy read from shared memory '{shm_name}': {e}") + raise + + def is_shm_available(self, shm_name: str) -> bool: + """检查共享内存是否可用""" + try: + shm = shared_memory.SharedMemory(name=shm_name) + shm.close() + return True + except FileNotFoundError: + return False + except Exception: + return False + + +# 全局单例 +_shm_client: Optional[ShmClient] = None + + +def get_shm_client() -> ShmClient: + """获取共享内存客户端单例""" + global _shm_client + if _shm_client is None: + _shm_client = ShmClient() + return _shm_client diff --git a/scripts/qwen_image/qwen_image_i2i_2511_kernel.sh b/scripts/qwen_image/qwen_image_i2i_2511_kernel.sh new file mode 100755 index 00000000..e0f3a5e0 --- /dev/null +++ b/scripts/qwen_image/qwen_image_i2i_2511_kernel.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# set path and first +export lightx2v_path= +export model_path= + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ + --model_cls qwen_image \ + --task i2i \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/qwen_image/qwen_image_i2i_2511_kernel.json \ + --prompt "Make the girl from Image 1 wear the black dress from Image 2 and sit in the pose from Image 3." \ + --negative_prompt " " \ + --image_path "1.png,2.png,3.png" \ + --save_result_path ${lightx2v_path}/save_results/qwen_image_i2i_2511_kernel.png \ + --seed 0 diff --git a/scripts/qwen_image/qwen_image_i2i_2511_service.sh b/scripts/qwen_image/qwen_image_i2i_2511_service.sh new file mode 100755 index 00000000..4418d98f --- /dev/null +++ b/scripts/qwen_image/qwen_image_i2i_2511_service.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +# set path and first +export lightx2v_path= +export model_path= + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +python -m lightx2v.infer \ + --model_cls qwen_image \ + --task i2i \ + --model_path $model_path \ + --config_json ${lightx2v_path}/configs/qwen_image/qwen_image_i2i_2511_service.json \ + --prompt "Make the girl from Image 1 wear the black dress from Image 2 and sit in the pose from Image 3." \ + --negative_prompt " " \ + --image_path "1.png,2.png,3.png" \ + --save_result_path ${lightx2v_path}/save_results/qwen_image_i2i_2511_service.png \ + --seed 0 From c06594e8cbe848f714f5810e3d896a51b16a2d5a Mon Sep 17 00:00:00 2001 From: fuheaven Date: Tue, 27 Jan 2026 11:44:27 +0800 Subject: [PATCH 6/6] format code --- configs/qwen_image/qwen_image_i2i_2511_kernel.json | 2 +- configs/qwen_image/qwen_image_i2i_2511_service.json | 2 +- .../input_encoders/lightllm/qwen25_text_encoder_service.py | 1 + lightx2v/models/input_encoders/lightllm/shm_client.py | 5 +++-- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/configs/qwen_image/qwen_image_i2i_2511_kernel.json b/configs/qwen_image/qwen_image_i2i_2511_kernel.json index 780e2f00..d6dec78b 100644 --- a/configs/qwen_image/qwen_image_i2i_2511_kernel.json +++ b/configs/qwen_image/qwen_image_i2i_2511_kernel.json @@ -14,4 +14,4 @@ "use_flash_attention_kernel": true, "use_rmsnorm_kernel": true } -} \ No newline at end of file +} diff --git a/configs/qwen_image/qwen_image_i2i_2511_service.json b/configs/qwen_image/qwen_image_i2i_2511_service.json index 96fa9eb1..b4ae30f7 100644 --- a/configs/qwen_image/qwen_image_i2i_2511_service.json +++ b/configs/qwen_image/qwen_image_i2i_2511_service.json @@ -16,4 +16,4 @@ "service_retry": 3, "use_shm": true } -} \ No newline at end of file +} diff --git a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py index 968d5fa4..d3e9468f 100644 --- a/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py +++ b/lightx2v/models/input_encoders/lightllm/qwen25_text_encoder_service.py @@ -302,6 +302,7 @@ def infer(self, text: List[str], image_list: Optional[List] = None) -> Tuple: # Fallback to base64 mode if "hidden_states_base64" in result: import base64 + data_bytes = base64.b64decode(result["hidden_states_base64"]) shape = result["hidden_states_shape"] hidden_states_np = np.frombuffer(data_bytes, dtype=np.uint8).reshape(shape) diff --git a/lightx2v/models/input_encoders/lightllm/shm_client.py b/lightx2v/models/input_encoders/lightllm/shm_client.py index 43e3eaed..40c87496 100644 --- a/lightx2v/models/input_encoders/lightllm/shm_client.py +++ b/lightx2v/models/input_encoders/lightllm/shm_client.py @@ -5,9 +5,10 @@ 实现零拷贝数据传输，显著降低通信延迟。 """ -import numpy as np from multiprocessing import shared_memory -from typing import Tuple, Optional +from typing import Optional, Tuple + +import numpy as np from loguru import logger