[Feature] gen-only FP8 Quantization Support for SenseNova-U1 (#3943)

Signed-off-by: leohuang257 <masugchds@gmail.com>

tests/diffusion/quantization/test_quantization_fp8.py

@@ -6,11 +6,13 @@ End-to-end tests for the unified quantization framework (PR #1764).
 Validates FP8 quantization works correctly for all supported model types:
   - Single-stage diffusion models (FLUX.1-dev, Qwen-Image, Z-Image-Turbo)
   - Multi-stage models (BAGEL: LLM + Diffusion)
+  - MoT gen-only models (SenseNova-U1: gen-path FP8, und-path BF16)
 
 Tests verify:
   1. FP8 quantization produces valid images
   2. Memory usage is lower than BF16 baseline
   3. Multi-stage models only quantize the diffusion stage (not the LLM stage)
+  4. MoT models only quantize gen-path layers (mot_gen)
 
 Usage:
     # Run all FP8 quantization tests
@@ -436,3 +438,92 @@ def test_single_stage_quantization_config_key():
         quantization_config="fp8",
     )
     assert len(images) >= 1
+
+
+# ─── SenseNova-U1 gen-only FP8 tests ─────────────────────────────────────────
+
+_SENSENOVA_MODEL = "SenseNova/SenseNova-U1-8B-MoT"
+
+
+def _generate_sensenova_u1_image(
+    quantization_config: str | None = None,
+    num_inference_steps: int = 20,
+) -> tuple[Any, float]:
+    """Generate an image with SenseNova-U1 (single-stage MoT, gen-only FP8).
+
+    Returns (generated_image, peak_memory_gib).
+    """
+    omni_kwargs: dict[str, Any] = {}
+    if quantization_config:
+        omni_kwargs["quantization_config"] = quantization_config
+
+    with OmniRunner(_SENSENOVA_MODEL, **omni_kwargs) as runner:
+        torch.accelerator.reset_peak_memory_stats()
+
+        sampling_params = OmniDiffusionSamplingParams(
+            height=1024,
+            width=1024,
+            seed=42,
+            num_inference_steps=num_inference_steps,
+            extra_args={
+                "cfg_scale": 4.0,
+                "cfg_norm": "none",
+                "timestep_shift": 3.0,
+            },
+        )
+        omni_outputs = list(
+            runner.omni.generate(
+                prompts={"prompt": "A cat sitting on a windowsill", "modalities": ["image"]},
+                sampling_params_list=sampling_params,
+            )
+        )
+
+        peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+
+        generated_image = None
+        for req_output in omni_outputs:
+            if images := getattr(req_output, "images", None):
+                generated_image = images[0]
+                break
+
+        assert generated_image is not None, "No images generated from SenseNova-U1"
+        assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
+
+        return generated_image, peak_mem
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_fp8_generates_image():
+    """SenseNova-U1 with gen-only FP8 generates a valid image.
+
+    FP8 is applied only to gen-path (mot_gen) layers; understanding-path
+    layers stay in BF16.
+    """
+    image, _ = _generate_sensenova_u1_image(quantization_config="fp8")
+    image.save("test_sensenova_u1_fp8.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_bf16_generates_image():
+    """SenseNova-U1 without quantization generates a valid image (baseline)."""
+    image, _ = _generate_sensenova_u1_image(quantization_config=None)
+    image.save("test_sensenova_u1_bf16.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_fp8_uses_less_memory():
+    """Gen-only FP8 should use less peak memory than BF16 for SenseNova-U1."""
+    _, mem_bf16 = _generate_sensenova_u1_image(
+        quantization_config=None,
+        num_inference_steps=10,
+    )
+    torch.accelerator.empty_cache()
+
+    _, mem_fp8 = _generate_sensenova_u1_image(
+        quantization_config="fp8",
+        num_inference_steps=10,
+    )
+
+    print(f"SenseNova-U1 BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"SenseNova-U1 FP8  peak memory: {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"