Address review comments (benchmark logs, ids, and design comments)

Signed-off-by: Yufeng Zhu <yzhu802@gatech.edu>
Signed-off-by: princepride <wangzhipeng628@gmail.com>

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",

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,

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

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

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,