minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-31 01:40:29 +08:00
committed by GitHub
parent 8c28cdd116
commit 5ab97c4f44
2 changed files with 214 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+99 -2
View File
@@ -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
test/registered/kernels/ops/diffusion/test_model_fast_paths.py
+115
View File
@@ -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
# -------------------------------------------------------------------------