From e37f46fcf7bcdbe459105e97f10d099c3643d398 Mon Sep 17 00:00:00 2001
From: GXIN <37653830+gxxx-hum@users.noreply.github.com>
Date: Sun, 3 May 2026 21:18:26 +0800
Subject: [PATCH] [NPU] Fix Z-Image negative-branch rotary embeddings for CFG
 (#23538)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
---
 .../configs/pipeline_configs/zimage.py        |  7 ++-
 .../test/unit/test_zimage_pipeline_config.py  | 45 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py

diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
index dada61710..2bfc46d14 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py
@@ -361,9 +361,14 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
         }
 
     def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        prompt_embeds = (
+            batch.negative_prompt_embeds[0]
+            if batch.negative_prompt_embeds is not None
+            else batch.prompt_embeds[0]
+        )
         return {
             "freqs_cis": self.get_freqs_cis(
-                batch.prompt_embeds[0],
+                prompt_embeds,
                 batch.width,
                 batch.height,
                 device,
diff --git a/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py
new file mode 100644
index 000000000..aac9b99ef
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_zimage_pipeline_config.py
@@ -0,0 +1,45 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
+
+
+class TestZImagePipelineConfig(unittest.TestCase):
+    @patch("sglang.multimodal_gen.configs.pipeline_configs.zimage.get_sp_world_size")
+    def test_zimage_negative_prompt_rotary_embeddings_use_negative_prompt_len(
+        self, mock_get_sp_world_size
+    ) -> None:
+        """Negative CFG branch should build RoPE positions from negative prompt embeds."""
+        mock_get_sp_world_size.return_value = 1
+
+        config = ZImagePipelineConfig()
+        pos_seq_len = 19
+        neg_seq_len = 45
+        batch = SimpleNamespace(
+            prompt_embeds=[torch.ones(pos_seq_len, 2560)],
+            negative_prompt_embeds=[torch.ones(neg_seq_len, 2560)],
+            height=16,
+            width=16,
+        )
+
+        def rotary_emb(pos_ids):
+            return pos_ids
+
+        neg_kwargs = config.prepare_neg_cond_kwargs(
+            batch=batch,
+            device=torch.device("cpu"),
+            rotary_emb=rotary_emb,
+            dtype=torch.float32,
+        )
+
+        cap_pos_ids, image_pos_ids = neg_kwargs["freqs_cis"]
+        neg_cap_padded_len = 64
+        self.assertEqual(cap_pos_ids.shape, (neg_cap_padded_len, 3))
+        self.assertEqual(image_pos_ids[0].tolist(), [neg_cap_padded_len + 1, 0, 0])
+
+
+if __name__ == "__main__":
+    unittest.main()