diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 858313f0f..878159b81 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -1022,6 +1022,7 @@ def _save_generated_outputs( """Decode and save base64 images/videos from successful responses.""" os.makedirs(save_dir, exist_ok=True) saved = 0 + failed = 0 for idx, (req, out) in enumerate(zip(requests_list, outputs)): if not out.success or not out.response_body: @@ -1074,9 +1075,10 @@ def _save_generated_outputs( f.write(img_bytes) saved += 1 except Exception as e: - print(f"Warning: failed to save image for request {idx}: {e}") + failed += 1 + logger.warning(f"Failed to save image for request {idx}: {e}", exc_info=True) - print(f"Saved {saved} generated image(s) to {save_dir}") + logger.info(f"Saved {saved} generated image(s) to {save_dir}. Failed to save {failed} image(s).") def wait_for_service(base_url: str, timeout: int = 120) -> None: @@ -1347,13 +1349,18 @@ if __name__ == "__main__": default=1, help="Number of warmup requests to run before measurement.", ) + # NOTE Changed default from 1 to 2 because some models (e.g., Bagel) run + # `num_timesteps - 1` denoising iterations. A default of 1 results in 0 steps, + # which causes errors. + # TODO If this slightly longer warmup causes regression issues for other + # diffusion pipelines in the future, consider implementing model-specific + # overrides instead of a global default. parser.add_argument( "--warmup-num-inference-steps", type=int, default=2, - help="num_inference_steps used for warmup requests. " - "Must be >= 2 to ensure at least one denoising step is executed " - "(some models, e.g. Bagel, run num_timesteps-1 denoising iterations).", + help="Number of inference steps used for warmup requests. " + "Default is 2 to ensure at least one denoising step is executed.", ) parser.add_argument( "--warmup-concurrency", diff --git a/vllm_omni/diffusion/layers/mot/configs/README b/docs/user_guide/diffusion/mot_config.md similarity index 100% rename from vllm_omni/diffusion/layers/mot/configs/README rename to docs/user_guide/diffusion/mot_config.md diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index e48dde4d4..4f221e9e2 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/tests/diffusion/kernels/mot/test_mot_linear.py @@ -299,7 +299,7 @@ def _run_timing( @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM], - ids=lambda val: "", + ids=[f"img{num}_K3584_N4608" for num in _IMAGE_NUM], ) def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool): dcfg = _parse_dtype(dtype) @@ -384,7 +384,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM], - ids=lambda val: "", + ids=[f"img{num}_K3584_N3584" for num in _IMAGE_NUM], ) def test_mot_o_proj( image_num: int, diff --git a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py index 95df80831..7ce7bcadd 100644 --- a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py +++ b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py @@ -78,6 +78,13 @@ class MoTQKVParallelLinear(QKVParallelLinear): # ---- 2) Create vae weights (permanent submodule) ---- assert self.quant_method is not None + # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace + # container to hold the secondary VAE weights. + # This is a design choice for dynamically attaching `quant_method` to the bare module. + # This is slightly unconventional but necessary. It allows vLLM's framework + # (specifically `process_weights_after_loading`) to automatically discover + # and process these secondary weights without requiring a dedicated, + # boilerplate Module subclass just for the VAE pathway. self.gen_exp = torch.nn.Module() # Use the same weight_loader as text diff --git a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py index f69aab2d7..fcd57644f 100644 --- a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py +++ b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py @@ -71,6 +71,13 @@ class MoTRowParallelLinear(RowParallelLinear): # ---- Step 2: Create vae weights (permanent submodule) ---- assert self.quant_method is not None + # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace + # container to hold the secondary VAE weights. + # This is a design choice for dynamically attaching `quant_method` to the bare module. + # This is slightly unconventional but necessary. It allows vLLM's framework + # (specifically `process_weights_after_loading`) to automatically discover + # and process these secondary weights without requiring a dedicated, + # boilerplate Module subclass just for the VAE pathway. self.gen_exp = torch.nn.Module() # Select weight_loader consistent with text diff --git a/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py b/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py index ec5c8d7ac..f1bffbbac 100644 --- a/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py +++ b/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py @@ -152,6 +152,10 @@ def get_mot_default_config( } # Very small M (tail batches, final feature concat, etc.) + # NOTE If M is not a multiple of the block size (e.g., M=3 with BLOCK_SIZE_M=16), + # the kernel will allocate a full tile. Correctness is guaranteed by boundary masks. + # We tolerate this minor compute/SRAM padding to avoid generating excessive + # Triton compiled variants for extremely small sequence lengths. if M <= 64: return { "BLOCK_SIZE_M": 16 if M <= 16 else 32,