From 9bf93d293e0e5970014352c918eb6cf1c5cabfad Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Wed, 14 Oct 2026 08:52:24 +0000 Subject: [PATCH 1/3] [hooks] Support QwenImage2.1 in MagCache `QwenImage21TransformerBlock` was not registered in `TransformerBlockRegistry`, so enabling MagCache on `QwenImage21Transformer2DModel` raised `ValueError: Model class ... not registered`. Register it as a single-stream block (`return_hidden_states_index=0`, no encoder hidden states), matching how the model calls it. With the pipeline's KV cache, the hooked blocks see prefix + target tokens on the prefill step and only the target tokens afterwards. A step skipped right after the prefill therefore received a residual longer than its input and fell through to the "dimension mismatch" branch, returning the input without any residual; calibration crashed on the same shape change. Apply / compare the trailing (target) tokens of the residual instead, following the trailing-token alignment used for Flux2 in #14756. Tests: `MagCacheTesterMixin` for `QwenImage21Transformer2DModel` and `QwenImage21Pipeline`, plus KV-cache skip and calibration tests on the transformer. Refs: https://github.com/huggingface/diffusers/pull/14756 --- src/diffusers/hooks/_helpers.py | 10 +++ src/diffusers/hooks/mag_cache.py | 10 +++ .../test_models_transformer_qwenimage21.py | 86 +++++++++++++++++++ .../pipelines/qwenimage21/test_qwenimage21.py | 5 ++ 4 files changed, 111 insertions(+) diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 1c17d349562b..208fa810c61f 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -196,6 +196,7 @@ def _register_transformer_blocks_metadata(): MotifVideoTransformerBlock, ) from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock + from ..models.transformers.transformer_qwenimage21 import QwenImage21TransformerBlock from ..models.transformers.transformer_wan import WanTransformerBlock from ..models.transformers.transformer_z_image import ZImageTransformerBlock @@ -358,6 +359,15 @@ def _register_transformer_blocks_metadata(): ), ) + # QwenImage2.1 + TransformerBlockRegistry.register( + model_class=QwenImage21TransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=None, + ), + ) + # HunyuanImage2.1 TransformerBlockRegistry.register( model_class=HunyuanImageTransformerBlock, diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index eb76e343b30e..e897915ff813 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -249,6 +249,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): if diff > 0: output = output.clone() output[:, diff:, :] = output[:, diff:, :] + res + elif diff < 0: + # The residual was computed on a longer sequence (e.g. the QwenImage2.1 KV-cache prefill step + # covers prefix + target tokens, later steps only the target); the target tokens sit at the end. + output = output + res[:, -output.shape[1] :, :] else: logger.warning( f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. " @@ -371,6 +375,12 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc # norm(dim=-1) gives magnitude of each token vector curr_norm = torch.linalg.norm(current_residual.float(), dim=-1) prev_norm = torch.linalg.norm(state.previous_residual.float(), dim=-1) + if curr_norm.ndim == 2 and prev_norm.ndim == 2 and curr_norm.shape[1] != prev_norm.shape[1]: + # Sequence length changed between steps (e.g. KV-cache prefill vs. cached steps); compare the + # trailing tokens both steps share. + num_tokens = min(curr_norm.shape[1], prev_norm.shape[1]) + curr_norm = curr_norm[:, -num_tokens:] + prev_norm = prev_norm[:, -num_tokens:] # Avoid division by zero ratio = (curr_norm / (prev_norm + 1e-8)).mean().item() diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index f6d0a0004d0c..57c1b746b80f 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -25,6 +25,7 @@ from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, + MagCacheTesterMixin, MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, @@ -364,6 +365,91 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): assert not torch.allclose(with_kv_cache[3], uncached[3], atol=1e-5) +class TestQwenImage21TransformerMagCache(QwenImage21TransformerTesterConfig, MagCacheTesterMixin): + """MagCache tests for QwenImage 2.1 Transformer.""" + + def _run_denoising_loop(self, model, step_inputs, use_kv_cache): + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. + kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None + outputs = [] + with torch.no_grad(), model.cache_context("cond"): + for index, inputs in enumerate(step_inputs): + kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None + output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] + outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) + return outputs + + def _get_step_inputs(self, num_steps): + base_inputs = self.get_dummy_inputs() + generator = torch.Generator("cpu").manual_seed(0) + return [ + dict( + base_inputs, + hidden_states=randn_tensor( + base_inputs["hidden_states"].shape, generator=generator, device=torch_device + ), + timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), + ) + for index in range(num_steps) + ] + + def test_mag_cache_with_kv_cache(self): + """ + With the KV cache, the hooked blocks see prefix + target tokens on the prefill step and target tokens only + afterwards. A step skipped right after the prefill must reuse the target part of the prefill residual, giving + the same target tokens as a loop that recomputes the full sequence every step. + """ + from diffusers import MagCacheConfig + + num_steps = 4 + step_inputs = self._get_step_inputs(num_steps) + + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + uncached = self._run_denoising_loop(model, step_inputs, use_kv_cache=False) + + # Step 0 computes, step 1 is skipped, step 2 computes, step 3 is skipped. + model.enable_cache( + MagCacheConfig( + num_inference_steps=num_steps, + retention_ratio=0.0, + threshold=100.0, + max_skip_steps=1, + mag_ratios=[1.0] * num_steps, + ) + ) + with_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) + model._reset_stateful_cache() + without_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=False) + model.disable_cache() + + for step, (cached, reference) in enumerate(zip(with_kv_cache, without_kv_cache)): + assert cached.shape == reference.shape + torch.testing.assert_close(cached, reference, atol=2e-5, rtol=2e-5, msg=f"step {step} diverged") + + # The skipped steps must actually be approximated, otherwise the comparison above proves nothing. + assert not torch.allclose(with_kv_cache[1], uncached[1], atol=1e-5) + + def test_mag_cache_calibration_with_kv_cache(self): + from diffusers import MagCacheConfig + + num_steps = 4 + step_inputs = self._get_step_inputs(num_steps) + + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.enable_cache(MagCacheConfig(num_inference_steps=num_steps, calibrate=True)) + # Calibration never skips, so the outputs must match the uncached loop exactly. + calibrated = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) + model.disable_cache() + reference = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) + + for step, (cached, expected) in enumerate(zip(calibrated, reference)): + torch.testing.assert_close(cached, expected, msg=f"step {step} diverged") + + class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin): @property def ckpt_path(self): diff --git a/tests/pipelines/qwenimage21/test_qwenimage21.py b/tests/pipelines/qwenimage21/test_qwenimage21.py index 72b960c33936..21f07e2ece6c 100644 --- a/tests/pipelines/qwenimage21/test_qwenimage21.py +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -38,6 +38,7 @@ BasePipelineTesterConfig, LoraMemoryTesterMixin, LoraTesterMixin, + MagCacheTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -347,3 +348,7 @@ def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_outp class TestQwenImage21PipelineLoRAMemory(QwenImage21PipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests for the Qwen-Image 2.1 pipeline.""" + + +class TestQwenImage21PipelineMagCache(QwenImage21PipelineTesterConfig, MagCacheTesterMixin): + """MagCache tests for the Qwen-Image 2.1 pipeline.""" From d52f18d4621d05b6205df4c654bf09a02770c78a Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Wed, 14 Oct 2026 10:01:40 +0000 Subject: [PATCH 2/3] [tests] Share the QwenImage2.1 KV-cache denoising loop between cache tests Move `_run_denoising_loop` and the per-step input builder onto `QwenImage21TransformerTesterConfig` so the TaylorSeer and MagCache KV-cache tests reuse them instead of carrying copies, and fold the MagCache calibration check into the KV-cache skip test. --- .../test_models_transformer_qwenimage21.py | 112 +++++++----------- 1 file changed, 40 insertions(+), 72 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index 57c1b746b80f..8943c1ffcea3 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -97,6 +97,33 @@ def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str "img_mask": img_mask, } + def _get_step_inputs(self, num_steps): + base_inputs = self.get_dummy_inputs() + generator = torch.Generator("cpu").manual_seed(0) + return [ + dict( + base_inputs, + hidden_states=randn_tensor( + base_inputs["hidden_states"].shape, generator=generator, device=torch_device + ), + timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), + ) + for index in range(num_steps) + ] + + def _run_denoising_loop(self, model, step_inputs, use_kv_cache): + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. + kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None + outputs = [] + with torch.no_grad(), model.cache_context("cond"): + for index, inputs in enumerate(step_inputs): + kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None + output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] + outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) + return outputs + class TestQwenImage21Transformer(QwenImage21TransformerTesterConfig, ModelTesterMixin): @pytest.mark.skip( @@ -302,19 +329,6 @@ class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, At class TestQwenImage21TransformerTaylorSeerCache(QwenImage21TransformerTesterConfig, TaylorSeerCacheTesterMixin): """TaylorSeerCache tests for QwenImage 2.1 Transformer.""" - def _run_denoising_loop(self, model, step_inputs, use_kv_cache): - from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache - - # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. - kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None - outputs = [] - with torch.no_grad(), model.cache_context("cond"): - for index, inputs in enumerate(step_inputs): - kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None - output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] - outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) - return outputs - @pytest.mark.parametrize("use_lite_mode", [False, True]) def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): """ @@ -324,19 +338,7 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): """ from diffusers import TaylorSeerCacheConfig - num_steps = 8 - base_inputs = self.get_dummy_inputs() - generator = torch.Generator("cpu").manual_seed(0) - step_inputs = [ - dict( - base_inputs, - hidden_states=randn_tensor( - base_inputs["hidden_states"].shape, generator=generator, device=torch_device - ), - timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), - ) - for index in range(num_steps) - ] + step_inputs = self._get_step_inputs(num_steps=8) torch.manual_seed(0) model = self.model_class(**self.get_init_dict()).to(torch_device).eval() @@ -368,38 +370,12 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): class TestQwenImage21TransformerMagCache(QwenImage21TransformerTesterConfig, MagCacheTesterMixin): """MagCache tests for QwenImage 2.1 Transformer.""" - def _run_denoising_loop(self, model, step_inputs, use_kv_cache): - from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache - - # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. - kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None - outputs = [] - with torch.no_grad(), model.cache_context("cond"): - for index, inputs in enumerate(step_inputs): - kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None - output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] - outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) - return outputs - - def _get_step_inputs(self, num_steps): - base_inputs = self.get_dummy_inputs() - generator = torch.Generator("cpu").manual_seed(0) - return [ - dict( - base_inputs, - hidden_states=randn_tensor( - base_inputs["hidden_states"].shape, generator=generator, device=torch_device - ), - timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), - ) - for index in range(num_steps) - ] - def test_mag_cache_with_kv_cache(self): """ With the KV cache, the hooked blocks see prefix + target tokens on the prefill step and target tokens only - afterwards. A step skipped right after the prefill must reuse the target part of the prefill residual, giving - the same target tokens as a loop that recomputes the full sequence every step. + afterwards. Calibration must compare residuals across that shape change, and a step skipped right after the + prefill must reuse the target part of the prefill residual, giving the same target tokens as a loop that + recomputes the full sequence every step. """ from diffusers import MagCacheConfig @@ -410,6 +386,15 @@ def test_mag_cache_with_kv_cache(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() uncached = self._run_denoising_loop(model, step_inputs, use_kv_cache=False) + # Calibration never skips, so it must reproduce the uncached outputs. + model.enable_cache(MagCacheConfig(num_inference_steps=num_steps, calibrate=True)) + calibrated = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) + model.disable_cache() + for step, (cached, reference) in enumerate(zip(calibrated, uncached)): + torch.testing.assert_close( + cached, reference, atol=2e-5, rtol=2e-5, msg=f"calibration step {step} diverged" + ) + # Step 0 computes, step 1 is skipped, step 2 computes, step 3 is skipped. model.enable_cache( MagCacheConfig( @@ -432,23 +417,6 @@ def test_mag_cache_with_kv_cache(self): # The skipped steps must actually be approximated, otherwise the comparison above proves nothing. assert not torch.allclose(with_kv_cache[1], uncached[1], atol=1e-5) - def test_mag_cache_calibration_with_kv_cache(self): - from diffusers import MagCacheConfig - - num_steps = 4 - step_inputs = self._get_step_inputs(num_steps) - - torch.manual_seed(0) - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - model.enable_cache(MagCacheConfig(num_inference_steps=num_steps, calibrate=True)) - # Calibration never skips, so the outputs must match the uncached loop exactly. - calibrated = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) - model.disable_cache() - reference = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) - - for step, (cached, expected) in enumerate(zip(calibrated, reference)): - torch.testing.assert_close(cached, expected, msg=f"step {step} diverged") - class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin): @property From 86960edbfda2a74d9764de73067b03701fea525e Mon Sep 17 00:00:00 2001 From: Pham Hong Vinh Date: Wed, 14 Oct 2026 18:11:12 +0000 Subject: [PATCH 3/3] [tests] Use only MagCacheTesterMixin for the QwenImage2.1 transformer Match the Flux2 MagCache tests (#14756): drop the bespoke KV-cache test and keep the shared denoising helpers in the TaylorSeer test as on main. --- .../test_models_transformer_qwenimage21.py | 101 +++++------------- 1 file changed, 26 insertions(+), 75 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index 8943c1ffcea3..0308b39b6605 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -97,33 +97,6 @@ def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str "img_mask": img_mask, } - def _get_step_inputs(self, num_steps): - base_inputs = self.get_dummy_inputs() - generator = torch.Generator("cpu").manual_seed(0) - return [ - dict( - base_inputs, - hidden_states=randn_tensor( - base_inputs["hidden_states"].shape, generator=generator, device=torch_device - ), - timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), - ) - for index in range(num_steps) - ] - - def _run_denoising_loop(self, model, step_inputs, use_kv_cache): - from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache - - # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. - kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None - outputs = [] - with torch.no_grad(), model.cache_context("cond"): - for index, inputs in enumerate(step_inputs): - kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None - output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] - outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) - return outputs - class TestQwenImage21Transformer(QwenImage21TransformerTesterConfig, ModelTesterMixin): @pytest.mark.skip( @@ -329,6 +302,19 @@ class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, At class TestQwenImage21TransformerTaylorSeerCache(QwenImage21TransformerTesterConfig, TaylorSeerCacheTesterMixin): """TaylorSeerCache tests for QwenImage 2.1 Transformer.""" + def _run_denoising_loop(self, model, step_inputs, use_kv_cache): + from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache + + # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target. + kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None + outputs = [] + with torch.no_grad(), model.cache_context("cond"): + for index, inputs in enumerate(step_inputs): + kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None + output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0] + outputs.append(output[:, -inputs["hidden_states"].shape[1] :]) + return outputs + @pytest.mark.parametrize("use_lite_mode", [False, True]) def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): """ @@ -338,7 +324,19 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): """ from diffusers import TaylorSeerCacheConfig - step_inputs = self._get_step_inputs(num_steps=8) + num_steps = 8 + base_inputs = self.get_dummy_inputs() + generator = torch.Generator("cpu").manual_seed(0) + step_inputs = [ + dict( + base_inputs, + hidden_states=randn_tensor( + base_inputs["hidden_states"].shape, generator=generator, device=torch_device + ), + timestep=torch.tensor([1.0 - index / num_steps], device=torch_device), + ) + for index in range(num_steps) + ] torch.manual_seed(0) model = self.model_class(**self.get_init_dict()).to(torch_device).eval() @@ -370,53 +368,6 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): class TestQwenImage21TransformerMagCache(QwenImage21TransformerTesterConfig, MagCacheTesterMixin): """MagCache tests for QwenImage 2.1 Transformer.""" - def test_mag_cache_with_kv_cache(self): - """ - With the KV cache, the hooked blocks see prefix + target tokens on the prefill step and target tokens only - afterwards. Calibration must compare residuals across that shape change, and a step skipped right after the - prefill must reuse the target part of the prefill residual, giving the same target tokens as a loop that - recomputes the full sequence every step. - """ - from diffusers import MagCacheConfig - - num_steps = 4 - step_inputs = self._get_step_inputs(num_steps) - - torch.manual_seed(0) - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - uncached = self._run_denoising_loop(model, step_inputs, use_kv_cache=False) - - # Calibration never skips, so it must reproduce the uncached outputs. - model.enable_cache(MagCacheConfig(num_inference_steps=num_steps, calibrate=True)) - calibrated = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) - model.disable_cache() - for step, (cached, reference) in enumerate(zip(calibrated, uncached)): - torch.testing.assert_close( - cached, reference, atol=2e-5, rtol=2e-5, msg=f"calibration step {step} diverged" - ) - - # Step 0 computes, step 1 is skipped, step 2 computes, step 3 is skipped. - model.enable_cache( - MagCacheConfig( - num_inference_steps=num_steps, - retention_ratio=0.0, - threshold=100.0, - max_skip_steps=1, - mag_ratios=[1.0] * num_steps, - ) - ) - with_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=True) - model._reset_stateful_cache() - without_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=False) - model.disable_cache() - - for step, (cached, reference) in enumerate(zip(with_kv_cache, without_kv_cache)): - assert cached.shape == reference.shape - torch.testing.assert_close(cached, reference, atol=2e-5, rtol=2e-5, msg=f"step {step} diverged") - - # The skipped steps must actually be approximated, otherwise the comparison above proves nothing. - assert not torch.allclose(with_kv_cache[1], uncached[1], atol=1e-5) - class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin): @property