[Diffusion] Cache Qwen-Image modulation across serial CFG branches (#37090)

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]
+        img_mod_params, txt_mod_params = self._get_modulation_params(
-        txt_mod_params, _ = self.txt_mod[1](temb_txt_silu)  # [B, 6*dim]
+            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

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
 # -------------------------------------------------------------------------