[NPU] Fix Z-Image negative-branch rotary embeddings for CFG (#23538)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
GXIN
2026-05-03 16:18:26 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent 44ca2d01fc
commit e37f46fcf7
2 changed files with 51 additions and 1 deletions
@@ -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,
@@ -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()