Elastic-forcing-code / 14B /code /pipeline /causal_inference.py
Add training and inference source code
2ea2835 verified 29 days ago
31.4 kB
from typing import List, Optional
import torch

from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper

from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation


def _wan_cache_geometry(generator: WanDiffusionWrapper):
    """Return the cache geometry exposed by the loaded Wan backbone."""
    model = generator.model
    try:
        num_transformer_blocks = len(model.blocks)
        num_heads = int(model.num_heads)
        dim = int(model.dim)
    except AttributeError as exc:
        raise AttributeError(
            "Wan generator.model must expose blocks, num_heads, and dim"
        ) from exc

    if num_transformer_blocks <= 0:
        raise ValueError("Wan generator.model.blocks must not be empty")
    if num_heads <= 0 or dim <= 0 or dim % num_heads != 0:
        raise ValueError(
            f"invalid Wan attention geometry: dim={dim}, num_heads={num_heads}"
        )
    return num_transformer_blocks, num_heads, dim // num_heads


class CausalInferencePipeline(torch.nn.Module):
    def __init__(
            self,
            args,
            device,
            generator=None,
            text_encoder=None,
            vae=None
    ):
        super().__init__()
        # Step 1: Initialize all models
        if generator is None:
            # Mirror ``model/base.py:_initialize_models`` so that ``attention_type``
            # and ``num_frame_per_block`` from the config flow into the wrapper
            # (necessary for the block-causal linear-attention student).
            gen_kwargs = dict(getattr(args, "model_kwargs", {}))
            attn_type = getattr(args, "attention_type", "softmax")
            # ``gdn`` is injected into a softmax backbone (see wan/modules/gdn.py).
            is_gdn = attn_type == "gdn"
            backbone_attn_type = "softmax" if is_gdn else attn_type
            gen_kwargs.setdefault("attention_type", backbone_attn_type)
            if backbone_attn_type in ("linear", "hybrid"):
                gen_kwargs.setdefault(
                    "num_frame_per_block", getattr(args, "num_frame_per_block", 1),
                )
            if backbone_attn_type == "hybrid":
                gen_kwargs.setdefault(
                    "hybrid_gate", bool(getattr(args, "hybrid_gate", False)),
                )
                gen_kwargs.setdefault(
                    "hybrid_gate_init_lambda", float(getattr(args, "hybrid_gate_init_lambda", 0.12)),
                )
                gen_kwargs.setdefault(
                    "hybrid_gate_type", str(getattr(args, "hybrid_gate_type", "none")),
                )
                gen_kwargs.setdefault(
                    "alpha_init_bias", float(getattr(args, "alpha_init_bias", -3.0)),
                )
                gen_kwargs.setdefault(
                    "alpha_detach_denom_delta", bool(getattr(args, "alpha_detach_denom_delta", True)),
                )
                gen_kwargs.setdefault(
                    "alpha_detach_state_summary", bool(getattr(args, "alpha_detach_state_summary", True)),
                )
                # Fast alpha-hybrid path knobs (see ``model/base.py``).
                gen_kwargs.setdefault(
                    "use_fast_alpha_hybrid", bool(getattr(args, "use_fast_alpha_hybrid", False)),
                )
                gen_kwargs.setdefault(
                    "hybrid_local_impl", str(getattr(args, "hybrid_local_impl", "torch")),
                )
                gen_kwargs.setdefault(
                    "hybrid_cross_impl", str(getattr(args, "hybrid_cross_impl", "torch")),
                )
                gen_kwargs.setdefault(
                    "alpha_combine_impl", str(getattr(args, "alpha_combine_impl", "torch")),
                )
                gen_kwargs.setdefault(
                    "fallback_to_torch_on_error", bool(getattr(args, "fallback_to_torch_on_error", True)),
                )
                gen_kwargs.setdefault(
                    "match_torch_centering", bool(getattr(args, "match_torch_centering", True)),
                )
                gen_kwargs.setdefault(
                    "fast_alpha_maxperf", bool(getattr(args, "fast_alpha_maxperf", False)),
                )
                # ---- mixed-layer attention ----
                sm_indices = getattr(args, "softmax_layer_indices", None)
                sm_interval = getattr(args, "softmax_layer_interval", None)
                if sm_indices is not None:
                    gen_kwargs.setdefault("softmax_layer_indices", list(sm_indices))
                if sm_interval is not None:
                    gen_kwargs.setdefault("softmax_layer_interval", int(sm_interval))
                gen_kwargs.setdefault(
                    "softmax_layer_local_attn_size",
                    int(getattr(args, "softmax_layer_local_attn_size", -1)),
                )
                gen_kwargs.setdefault(
                    "softmax_layer_sink_size",
                    int(getattr(args, "softmax_layer_sink_size", 0)),
                )
            self.generator = WanDiffusionWrapper(**gen_kwargs, is_causal=True)
            # ── GDN (ARL²) injection — happens BEFORE the checkpoint load in
            # inference.py so the injected keys line up.  load_stage2=False:
            # inference.py loads the full generator (Stage-2 or post-DMD) itself.
            # Keep the ordinary softmax inference path independent of the
            # optional GDN/FLA stack. Importing ``wan.modules.gdn`` eagerly
            # would require ``fla`` even when this config does not request any
            # GDN layers.
            gdn_configured = is_gdn or any(
                getattr(args, name, None) is not None
                for name in (
                    "gdn_layer_indices",
                    "gdn_layer_indices_json",
                    "arl2_hybrid_layers",
                    "arl2_hybrid_layers_json",
                )
            )
            if gdn_configured:
                from wan.modules.gdn import (
                    gdn_layers_from_config,
                    setup_gdn_generator,
                )
                gdn_layers = gdn_layers_from_config(args)
                setup_gdn_generator(self.generator, args, load_stage2=False)
                print(f"[causal_inference] GDN hybrid injected at layers {gdn_layers}")
        else:
            self.generator = generator
        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
        self.vae = WanVAEWrapper() if vae is None else vae

        # Step 2: Initialize all causal hyperparmeters
        self.scheduler = self.generator.get_scheduler()
        self.denoising_step_list = torch.tensor(
            args.denoising_step_list, dtype=torch.long)
        if args.warp_denoising_step:
            timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]

        # Evaluated 14B sampler: preserve the four trained non-zero
        # stages, then append the terminal t=0 cleanup after warping.
        self.inference_terminal_zero_cleanup = bool(
            getattr(args, "inference_terminal_zero_cleanup", False)
        )
        if (
            self.inference_terminal_zero_cleanup
            and float(self.denoising_step_list[-1]) != 0.0
        ):
            self.denoising_step_list = torch.cat(
                (self.denoising_step_list, self.denoising_step_list.new_zeros(1))
            )

        # MMD training called scheduler.add_noise(..., t=0) before each cache
        # write.  Because extra_one_step=True, that is sigma ~= 0.00498 rather
        # than exact x0.  Match that cache-input distribution without adding
        # noise to the decoded output.
        self.inference_train_matched_context_cache = bool(
            getattr(args, "inference_train_matched_context_cache", False)
        )
        self.inference_aux_noise_seed = int(
            getattr(args, "inference_aux_noise_seed", 0)
        )
        self._inference_call_index = 0

        (
            self.num_transformer_blocks,
            self.num_heads,
            self.head_dim,
        ) = _wan_cache_geometry(self.generator)
        self.frame_seq_length = 1560

        self.attention_type = getattr(self.generator, "attention_type", "softmax")
        self.kv_cache1 = None
        self.args = args
        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
        self.independent_first_frame = args.independent_first_frame
        # Linear-attention student does not have a (sliding) ``local_attn_size``;
        # treat it as unbounded so the recurrent state is allocated for the full
        # video.
        self.local_attn_size = getattr(self.generator.model, "local_attn_size", -1)

        print(
            f"KV inference with {self.num_frame_per_block} frames per block "
            f"(attention_type={self.attention_type})"
        )

        if self.num_frame_per_block > 1:
            self.generator.model.num_frame_per_block = self.num_frame_per_block

    def _set_advance_state(self, value: bool) -> None:
        """Toggle ``advance_state`` on every linear/hybrid kv-cache dict.
        See ``SelfForcingTrainingPipeline._set_advance_state``: the flag
        controls whether the linear / hybrid recurrent state is updated in
        place on a given forward.  Pipelines set ``False`` for every
        denoising-loop pass and ``True`` only on the clean re-run so each
        block contributes EXACTLY ONCE to the persistent state, matching
        the chunkwise training semantics and aligning train / inference
        cache distributions.  No-op for softmax caches.
        """
        if self.kv_cache1 is None:
            return
        for cache in self.kv_cache1:
            if "state" in cache:
                cache["advance_state"] = value
            # ARL2 HybridBlockAttention layers: clean-pass GDN update flag.
            if "gdn_state" in cache:
                cache["update_state"] = value

    def inference(
        self,
        noise: torch.Tensor,
        text_prompts: List[str],
        initial_latent: Optional[torch.Tensor] = None,
        return_latents: bool = False,
        profile: bool = False,
        low_memory: bool = False,
        decode: bool = True,
        standard_noise_override=None,
        aux_noise_override=None,
    ) -> torch.Tensor:
        """
        Perform inference on the given noise and text prompts.
        Inputs:
            noise (torch.Tensor): The input noise tensor of shape
                (batch_size, num_output_frames, num_channels, height, width).
            text_prompts (List[str]): The list of text prompts.
            initial_latent (torch.Tensor): The initial latent tensor of shape
                (batch_size, num_input_frames, num_channels, height, width).
                If num_input_frames is 1, perform image to video.
                If num_input_frames is greater than 1, perform video extension.
            return_latents (bool): Whether to return the latents.
        Outputs:
            video (torch.Tensor): The generated video tensor of shape
                (batch_size, num_output_frames, num_channels, height, width).
                It is normalized to be in the range [0, 1].
        """
        batch_size, num_frames, num_channels, height, width = noise.shape
        aux_generator = torch.Generator(device=noise.device)
        aux_generator.manual_seed(
            self.inference_aux_noise_seed + self._inference_call_index
        )
        self._inference_call_index += 1

        def aux_randn_like(tensor: torch.Tensor) -> torch.Tensor:
            if aux_noise_override is not None:
                value = next(aux_noise_override).to(tensor.device)
                return value
            return torch.randn(
                tensor.shape,
                dtype=tensor.dtype,
                device=tensor.device,
                generator=aux_generator,
            )
        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
            # If the first frame is independent and the first frame is provided, then the number of frames in the
            # noise should still be a multiple of num_frame_per_block
            assert num_frames % self.num_frame_per_block == 0
            num_blocks = num_frames // self.num_frame_per_block
        else:
            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
            assert (num_frames - 1) % self.num_frame_per_block == 0
            num_blocks = (num_frames - 1) // self.num_frame_per_block
        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
        conditional_dict = self.text_encoder(
            text_prompts=text_prompts
        )

        if low_memory:
            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
            move_model_to_device_with_memory_preservation(self.text_encoder, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)

        output = torch.zeros(
            [batch_size, num_output_frames, num_channels, height, width],
            device=noise.device,
            dtype=noise.dtype
        )

        # Set up profiling if requested
        if profile:
            init_start = torch.cuda.Event(enable_timing=True)
            init_end = torch.cuda.Event(enable_timing=True)
            diffusion_start = torch.cuda.Event(enable_timing=True)
            diffusion_end = torch.cuda.Event(enable_timing=True)
            vae_start = torch.cuda.Event(enable_timing=True)
            vae_end = torch.cuda.Event(enable_timing=True)
            block_times = []
            block_start = torch.cuda.Event(enable_timing=True)
            block_end = torch.cuda.Event(enable_timing=True)
            init_start.record()

        # Step 1: Initialize KV cache to all zeros
        if self.kv_cache1 is None:
            self._initialize_kv_cache(
                batch_size=batch_size,
                dtype=noise.dtype,
                device=noise.device
            )
            self._initialize_crossattn_cache(
                batch_size=batch_size,
                dtype=noise.dtype,
                device=noise.device
            )
        else:
            # reset cross attn cache
            for block_index in range(self.num_transformer_blocks):
                self.crossattn_cache[block_index]["is_init"] = False
            # reset kv cache
            for block_index in range(len(self.kv_cache1)):
                cache = self.kv_cache1[block_index]
                cache["global_end_index"] = torch.tensor(
                    [0], dtype=torch.long, device=noise.device)
                cache["local_end_index"] = torch.tensor(
                    [0], dtype=torch.long, device=noise.device)
                # Linear/hybrid layers carry a recurrent state; softmax
                # layers carry a K/V cache.  In a MIXED model each layer
                # has only one of the two -- zero whichever is present.
                if "state" in cache:
                    cache["state"].zero_()
                if "normalizer" in cache:
                    cache["normalizer"].zero_()
                if "k" in cache:
                    cache["k"].zero_()
                if "v" in cache:
                    cache["v"].zero_()
                # ARL2 GDN layers: drop the leaked recurrent state so the next
                # video starts from a clean (lazily re-allocated) state.
                if "gdn_state" in cache:
                    cache["gdn_state"] = None
                    cache["update_state"] = True

        # Step 2: Cache context feature
        current_start_frame = 0
        if initial_latent is not None:
            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
            # Initial-latent context forwards write clean K/V into the cache
            # exactly once per block -> advance_state=True.
            self._set_advance_state(True)
            if self.independent_first_frame:
                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
                assert (num_input_frames - 1) % self.num_frame_per_block == 0
                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
                output[:, :1] = initial_latent[:, :1]
                self.generator(
                    noisy_image_or_video=initial_latent[:, :1],
                    conditional_dict=conditional_dict,
                    timestep=timestep * 0,
                    kv_cache=self.kv_cache1,
                    crossattn_cache=self.crossattn_cache,
                    current_start=current_start_frame * self.frame_seq_length,
                )
                current_start_frame += 1
            else:
                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
                assert num_input_frames % self.num_frame_per_block == 0
                num_input_blocks = num_input_frames // self.num_frame_per_block

            for _ in range(num_input_blocks):
                current_ref_latents = \
                    initial_latent[:, current_start_frame:current_start_frame + self.num_frame_per_block]
                output[:, current_start_frame:current_start_frame + self.num_frame_per_block] = current_ref_latents
                self.generator(
                    noisy_image_or_video=current_ref_latents,
                    conditional_dict=conditional_dict,
                    timestep=timestep * 0,
                    kv_cache=self.kv_cache1,
                    crossattn_cache=self.crossattn_cache,
                    current_start=current_start_frame * self.frame_seq_length,
                )
                current_start_frame += self.num_frame_per_block

        if profile:
            init_end.record()
            torch.cuda.synchronize()
            diffusion_start.record()

        # Step 3: Temporal denoising loop
        all_num_frames = [self.num_frame_per_block] * num_blocks
        if self.independent_first_frame and initial_latent is None:
            all_num_frames = [1] + all_num_frames
        for current_num_frames in all_num_frames:
            if profile:
                block_start.record()

            noisy_input = noise[
                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]

            # All denoising-loop forwards on this block read the same cross-block
            # cache and do NOT advance it; only the clean re-run (Step 3.3)
            # advances the cache.  No-op for softmax caches.
            self._set_advance_state(False)

            # Step 3.1: Spatial denoising loop
            for index, current_timestep in enumerate(self.denoising_step_list):
                print(f"current_timestep: {current_timestep}")
                # set current timestep
                timestep = torch.ones(
                    [batch_size, current_num_frames],
                    device=noise.device,
                    dtype=torch.int64) * current_timestep

                if index < len(self.denoising_step_list) - 1:
                    _, denoised_pred = self.generator(
                        noisy_image_or_video=noisy_input,
                        conditional_dict=conditional_dict,
                        timestep=timestep,
                        kv_cache=self.kv_cache1,
                        crossattn_cache=self.crossattn_cache,
                        current_start=current_start_frame * self.frame_seq_length
                    )
                    next_timestep = self.denoising_step_list[index + 1]
                    if (
                        self.inference_terminal_zero_cleanup
                        and float(next_timestep) == 0.0
                    ):
                        transition_noise = aux_randn_like(
                            denoised_pred.flatten(0, 1)
                        )
                    else:
                        target = denoised_pred.flatten(0, 1)
                        if standard_noise_override is None:
                            transition_noise = torch.randn_like(target)
                        else:
                            transition_noise = next(standard_noise_override).to(target.device)
                    noisy_input = self.scheduler.add_noise(
                        denoised_pred.flatten(0, 1),
                        transition_noise,
                        next_timestep * torch.ones(
                            [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
                    ).unflatten(0, denoised_pred.shape[:2])
                else:
                    # for getting real output
                    _, denoised_pred = self.generator(
                        noisy_image_or_video=noisy_input,
                        conditional_dict=conditional_dict,
                        timestep=timestep,
                        kv_cache=self.kv_cache1,
                        crossattn_cache=self.crossattn_cache,
                        current_start=current_start_frame * self.frame_seq_length
                    )

            # Step 3.2: record the model's output
            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred

            # Step 3.3: update KV cache at t=0.  Hybrid mode noises only this
            # cache input to match training; the stored output remains clean.
            context_timestep = torch.ones_like(timestep) * self.args.context_noise
            cache_input = denoised_pred
            if self.inference_train_matched_context_cache:
                flat_cache_input = cache_input.flatten(0, 1)
                cache_input = self.scheduler.add_noise(
                    flat_cache_input,
                    aux_randn_like(flat_cache_input),
                    context_timestep.flatten(0, 1),
                ).unflatten(0, denoised_pred.shape[:2])
            # This is the single write the linear/hybrid cache gets for this
            # block -- restore the default ``advance_state=True`` before the
            # forward so the clean K/V actually lands in the cache.
            self._set_advance_state(True)
            self.generator(
                noisy_image_or_video=cache_input,
                conditional_dict=conditional_dict,
                timestep=context_timestep,
                kv_cache=self.kv_cache1,
                crossattn_cache=self.crossattn_cache,
                current_start=current_start_frame * self.frame_seq_length,
            )

            if profile:
                block_end.record()
                torch.cuda.synchronize()
                block_time = block_start.elapsed_time(block_end)
                block_times.append(block_time)

            # Step 3.4: update the start and end frame indices
            current_start_frame += current_num_frames

        if profile:
            # End diffusion timing and synchronize CUDA
            diffusion_end.record()
            torch.cuda.synchronize()
            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
            init_time = init_start.elapsed_time(init_end)
            vae_start.record()

        # Step 4: Decode the output (skippable for latent-only generation, e.g.
        # ARL² Stage-2 teacher-video latent synthesis where pixels are not needed)
        if not decode:
            return (None, output) if return_latents else output
        video = self.vae.decode_to_pixel(output, use_cache=False)
        video = (video * 0.5 + 0.5).clamp(0, 1)

        if profile:
            # End VAE timing and synchronize CUDA
            vae_end.record()
            torch.cuda.synchronize()
            vae_time = vae_start.elapsed_time(vae_end)
            total_time = init_time + diffusion_time + vae_time

            print("Profiling results:")
            print(f"  - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)")
            print(f"  - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)")
            for i, block_time in enumerate(block_times):
                print(f"    - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)")
            print(f"  - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)")
            print(f"  - Total time: {total_time:.2f} ms")

        if return_latents:
            return video, output
        else:
            return video

    def _initialize_kv_cache(self, batch_size, dtype, device):
        """
        Initialize a Per-GPU KV cache for the Wan model.
        Per-layer cache shape depends on that layer's attention type:
        * ``softmax`` layers carry a sliding-window K/V tensor
          ``[B, kv_cache_size, H, Dv]`` in the activation dtype (bf16).
          The cache size matches the all-softmax baseline.
        * ``linear`` / ``hybrid`` layers carry a recurrent state
          ``[B, H, M, Dv]`` plus denominator ``[B, H, M]`` in fp32 for
          numerical stability over the long causal rollout.
        For a MIXED model (some softmax, some linear/hybrid) we walk the
        ``per_layer_self_attn_types`` list and pick the right shape per
        layer.  Pure-linear / pure-hybrid models use the same code path
        with every layer set to ``"linear_or_hybrid"``.
        """
        kv_cache1 = []
        model = self.generator.model

        # ── ARL2 path: softmax CausalWanModel with some blocks swapped to
        # HybridBlockAttention (GDN).  GDN layers use a lazily-allocated
        # gdn_state (+ update_state flag); other layers keep K/V.
        def is_gdn_block(block):
            cls = type(block.self_attn)
            return (
                cls.__name__ == "HybridBlockAttention"
                and cls.__module__ == "wan.modules.hybrid_attention"
            )

        if any(is_gdn_block(b) for b in model.blocks):
            sm_size = (self.local_attn_size * self.frame_seq_length
                       if self.local_attn_size != -1 else 32760)
            for b in model.blocks:
                entry = {
                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
                }
                if is_gdn_block(b):
                    entry["gdn_state"] = None
                    entry["update_state"] = True
                else:
                    entry["k"] = torch.zeros([batch_size, sm_size, self.num_heads, self.head_dim], dtype=dtype, device=device)
                    entry["v"] = torch.zeros([batch_size, sm_size, self.num_heads, self.head_dim], dtype=dtype, device=device)
                kv_cache1.append(entry)
            self.kv_cache1 = kv_cache1
            return

        if self.attention_type in ("linear", "hybrid"):
            # Default: every layer is linear/hybrid.  A mixed model exposes
            # an explicit per-layer types list that overrides this.
            per_layer_types = getattr(
                model, "per_layer_self_attn_types",
                ["linear_or_hybrid"] * self.num_transformer_blocks,
            )
            # Softmax-layer KV cache size mirrors the all-softmax baseline.
            softmax_cache_size = (
                self.local_attn_size * self.frame_seq_length
                if self.local_attn_size != -1
                else 32760
            )
            for layer_idx in range(self.num_transformer_blocks):
                ltype = per_layer_types[layer_idx] if layer_idx < len(per_layer_types) else "linear_or_hybrid"
                if ltype == "softmax":
                    kv_cache1.append({
                        "k": torch.zeros(
                            [batch_size, softmax_cache_size, self.num_heads, self.head_dim],
                            dtype=dtype, device=device,
                        ),
                        "v": torch.zeros(
                            [batch_size, softmax_cache_size, self.num_heads, self.head_dim],
                            dtype=dtype, device=device,
                        ),
                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
                    })
                else:
                    # linear / hybrid recurrent state.  Kernel dim defaults
                    # to head_dim but may be overridden via num_kernel_features.
                    kernel_dim = self.head_dim
                    try:
                        first_block = model.blocks[layer_idx]
                        kernel_dim = first_block.self_attn.kernel_dim
                    except (AttributeError, IndexError):
                        pass
                    kv_cache1.append({
                        "state": torch.zeros(
                            [batch_size, self.num_heads, kernel_dim, self.head_dim],
                            dtype=torch.float32, device=device,
                        ),
                        "normalizer": torch.zeros(
                            [batch_size, self.num_heads, kernel_dim],
                            dtype=torch.float32, device=device,
                        ),
                        "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
                        "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
                    })
        else:
            # All-softmax model (CausalWanModel path; no mixing).
            if self.local_attn_size != -1:
                kv_cache_size = self.local_attn_size * self.frame_seq_length
            else:
                kv_cache_size = 32760
            for _ in range(self.num_transformer_blocks):
                kv_cache1.append({
                    "k": torch.zeros([batch_size, kv_cache_size, self.num_heads, self.head_dim], dtype=dtype, device=device),
                    "v": torch.zeros([batch_size, kv_cache_size, self.num_heads, self.head_dim], dtype=dtype, device=device),
                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
                })

        self.kv_cache1 = kv_cache1  # always store the clean cache

    def _initialize_crossattn_cache(self, batch_size, dtype, device):
        """
        Initialize a Per-GPU cross-attention cache for the Wan model.
        """
        crossattn_cache = []

        for _ in range(self.num_transformer_blocks):
            crossattn_cache.append({
                "k": torch.zeros([batch_size, 512, self.num_heads, self.head_dim], dtype=dtype, device=device),
                "v": torch.zeros([batch_size, 512, self.num_heads, self.head_dim], dtype=dtype, device=device),
                "is_init": False
            })
        self.crossattn_cache = crossattn_cache