diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 612a2843e..3414f00a2 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import functools +from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union import diffusers @@ -109,6 +110,60 @@ def _local_seq_len(seq_len: int, sp_world_size: int) -> int: _get_qkv_projections = get_qkv_projections +def _safe_tensor_version(tensor: torch.Tensor) -> Optional[int]: + """Read a tensor version counter without rejecting inference tensors.""" + return None if tensor.is_inference() else tensor._version + + +@dataclass(frozen=True, eq=False) +class _QwenModulationCacheKey: + timestep: torch.Tensor + timestep_version: Optional[int] + additional_t_cond: Optional[torch.Tensor] + additional_t_cond_version: Optional[int] + hidden_dtype: torch.dtype + hidden_device: torch.device + + def matches(self, other: "_QwenModulationCacheKey") -> bool: + return ( + self.timestep is other.timestep + and self.timestep_version == other.timestep_version + and self.additional_t_cond is other.additional_t_cond + and self.additional_t_cond_version == other.additional_t_cond_version + and self.hidden_dtype == other.hidden_dtype + and self.hidden_device == other.hidden_device + ) + + +def _qwen_modulation_cache_key( + timestep: Optional[torch.Tensor], + additional_t_cond: Optional[torch.Tensor], + hidden_states: torch.Tensor, +) -> Optional[_QwenModulationCacheKey]: + """Build the key shared by the two serial CFG forwards at one timestep.""" + if ( + not isinstance(timestep, torch.Tensor) + or torch.is_grad_enabled() + or torch.compiler.is_compiling() + or is_in_breakable_cuda_graph() + or (timestep.device.type == "cuda" and torch.cuda.is_current_stream_capturing()) + ): + return None + + return _QwenModulationCacheKey( + timestep=timestep, + timestep_version=_safe_tensor_version(timestep), + additional_t_cond=additional_t_cond, + additional_t_cond_version=( + _safe_tensor_version(additional_t_cond) + if isinstance(additional_t_cond, torch.Tensor) + else None + ), + hidden_dtype=hidden_states.dtype, + hidden_device=hidden_states.device, + ) + + class QwenTimestepProjEmbeddings(nn.Module): def __init__(self, embedding_dim, use_additional_t_cond=False): super().__init__() @@ -977,6 +1032,13 @@ class QwenImageTransformerBlock(nn.Module): self.fused_res_ln_ss_gate_select01 = ( FusedResidualLayerNormScaleShiftGateSelect01() ) + self._modulation_cache: Optional[ + Tuple[ + _QwenModulationCacheKey, + torch.Tensor, + torch.Tensor, + ] + ] = None nunchaku_enabled = ( quant_config is not None @@ -1058,6 +1120,30 @@ class QwenImageTransformerBlock(nn.Module): ) -> torch.Tensor: return self.fuse_mul_add(a, b, c, k) + def _get_modulation_params( + self, + temb_img_silu: torch.Tensor, + temb_txt_silu: torch.Tensor, + cache_key: Optional[_QwenModulationCacheKey], + ) -> Tuple[torch.Tensor, torch.Tensor]: + cached = self._modulation_cache + if ( + cache_key is not None + and cached is not None + and cached[0].matches(cache_key) + ): + self._modulation_cache = None + return cached[1], cached[2] + + img_mod_params, _ = self.img_mod[1](temb_img_silu) + txt_mod_params, _ = self.txt_mod[1](temb_txt_silu) + self._modulation_cache = ( + (cache_key, img_mod_params, txt_mod_params) + if cache_key is not None + else None + ) + return img_mod_params, txt_mod_params + def _modulate( self, x: torch.Tensor, @@ -1165,10 +1251,14 @@ class QwenImageTransformerBlock(nn.Module): image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, modulate_index: Optional[torch.Tensor] = None, + modulation_cache_key: Optional[_QwenModulationCacheKey] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # Get modulation parameters for both streams - img_mod_params, _ = self.img_mod[1](temb_img_silu) # [B, 6*dim] - txt_mod_params, _ = self.txt_mod[1](temb_txt_silu) # [B, 6*dim] + img_mod_params, txt_mod_params = self._get_modulation_params( + temb_img_silu, + temb_txt_silu, + modulation_cache_key, + ) if ( self.quant_config is not None @@ -1512,6 +1602,12 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): hidden_states, _ = self.img_in(hidden_states) + modulation_cache_key = _qwen_modulation_cache_key( + timestep, + additional_t_cond, + hidden_states, + ) + timestep = (timestep / 1000).to(hidden_states.dtype) if self.zero_cond_t: @@ -1611,6 +1707,7 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): image_rotary_emb=image_rotary_emb, joint_attention_kwargs=block_attention_kwargs, modulate_index=modulate_index, + modulation_cache_key=modulation_cache_key, ) # controlnet residual diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 79c691f61..9e5bd6f5c 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -95,6 +95,10 @@ from sglang.multimodal_gen.runtime.models.dits.longcat_image import ( _apply_longcat_qknorm_rope, ) from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate +from sglang.multimodal_gen.runtime.models.dits.qwen_image import ( + QwenImageTransformerBlock, + _qwen_modulation_cache_key, +) from sglang.multimodal_gen.runtime.models.dits.sana import ( _eager_ln_modulate as _sana_eager_ln_modulate, ) @@ -291,6 +295,117 @@ class TestFlux2EagerFusions(CustomTestCase): self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1) +# ------------------------------------------------------------------------- +# Qwen-Image -- reuse timestep-only modulation across serial CFG branches +# ------------------------------------------------------------------------- + + +class _CountingProjection(nn.Module): + def __init__(self, offset: float): + super().__init__() + self.offset = offset + self.calls = 0 + + def forward(self, x): + self.calls += 1 + return x + self.offset, None + + +class TestQwenImageModulationCache(CustomTestCase): + def _block(self): + block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock) + nn.Module.__init__(block) + block.img_mod = nn.ModuleList([nn.Identity(), _CountingProjection(1.0)]) + block.txt_mod = nn.ModuleList([nn.Identity(), _CountingProjection(2.0)]) + block._modulation_cache = None + return block + + def _key(self, timestep, hidden, additional_t_cond=None): + with torch.no_grad(): + return _qwen_modulation_cache_key( + timestep, + additional_t_cond, + hidden, + ) + + def test_matching_cfg_key_reuses_both_modulation_projections(self): + block = self._block() + timestep = torch.tensor([500.0], device="cuda") + hidden = torch.empty(1, 17, 32, device="cuda", dtype=torch.bfloat16) + img_temb = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16) + txt_temb = torch.randn_like(img_temb) + key = self._key(timestep, hidden) + + first = block._get_modulation_params(img_temb, txt_temb, key) + second = block._get_modulation_params(img_temb, txt_temb, key) + + self.assertIs(first[0], second[0]) + self.assertIs(first[1], second[1]) + self.assertIsNone(block._modulation_cache) + self.assertEqual(block.img_mod[1].calls, 1) + self.assertEqual(block.txt_mod[1].calls, 1) + + def test_tensor_identity_version_and_condition_invalidate_cache(self): + block = self._block() + timestep = torch.tensor([500.0], device="cuda") + hidden = torch.empty(1, 17, 32, device="cuda", dtype=torch.bfloat16) + temb = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16) + key = self._key(timestep, hidden) + block._get_modulation_params(temb, temb, key) + + with torch.no_grad(): + timestep.add_(1) + mutated = self._key(timestep, hidden) + block._get_modulation_params(temb, temb, mutated) + self.assertEqual(block.img_mod[1].calls, 2) + + same_value_new_tensor = self._key(timestep.clone(), hidden) + block._get_modulation_params(temb, temb, same_value_new_tensor) + self.assertEqual(block.img_mod[1].calls, 3) + + condition = torch.tensor([1], device="cuda") + conditioned = self._key(timestep, hidden, condition) + block._get_modulation_params(temb, temb, conditioned) + self.assertEqual(block.img_mod[1].calls, 4) + + def test_grad_enabled_path_disables_and_clears_cache(self): + block = self._block() + timestep = torch.tensor([500.0], device="cuda") + hidden = torch.empty(1, 17, 32, device="cuda", dtype=torch.bfloat16) + temb = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16) + key = self._key(timestep, hidden) + block._get_modulation_params(temb, temb, key) + + self.assertIsNone(_qwen_modulation_cache_key(timestep, None, hidden)) + block._get_modulation_params(temb, temb, None) + + self.assertIsNone(block._modulation_cache) + self.assertEqual(block.img_mod[1].calls, 2) + + def test_inference_tensors_cache_and_graph_path_falls_back(self): + block = self._block() + with torch.inference_mode(): + timestep = torch.tensor([500.0], device="cuda") + hidden = torch.empty(1, 17, 32, device="cuda", dtype=torch.bfloat16) + temb = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16) + key = _qwen_modulation_cache_key(timestep, None, hidden) + + first = block._get_modulation_params(temb, temb, key) + second = block._get_modulation_params(temb, temb, key) + + self.assertIs(first[0], second[0]) + self.assertEqual(block.img_mod[1].calls, 1) + + with ( + patch( + "sglang.multimodal_gen.runtime.models.dits.qwen_image.is_in_breakable_cuda_graph", + return_value=True, + ), + torch.no_grad(), + ): + self.assertIsNone(_qwen_modulation_cache_key(timestep, None, hidden)) + + # ------------------------------------------------------------------------- # GLM-Image -- LayerNorm + modulate and per-head qk LayerNorm # -------------------------------------------------------------------------