# Copyright 2024 NVIDIA CORPORATION & AFFILIATES # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # SPDX-License-Identifier: Apache-2.0 """Self-forcing flow Euler samplers for chunk-causal autoregressive video. This module provides the streaming Sana-WM inference samplers: * ``SelfForcingFlowEuler`` is the base chunk-causal autoregressive sampler. It walks ``base_chunk_frames``-sized chunks left-to-right, denoising each chunk against a KV cache accumulated from previously generated chunks. * ``SelfForcingFlowEulerCamCtrl`` extends the base sampler with the camera conditioning extras (``camera_conditions``, ``chunk_plucker``, etc.), first-frame conditioning, and the 10-slot dual-mode (state / concat) KV cache layout used by the camctrl ``forward_long`` path. This is the sampler used by the end-to-end streaming Sana-WM + LTX-2 refiner. """ from __future__ import annotations import importlib import os import sys import torch # Diffusers ships with a hard ``import flash_attn`` in some attention backends # that raises before ``flash_attn_interface`` (FA4) is considered. We # temporarily hide the installed ``flash_attn`` module so diffusers takes the # FA-not-installed branch, then restore it so downstream code can still use FA. _fa_spec = importlib.util.find_spec("flash_attn") _has_fa = _fa_spec is not None _real_fa_module = None if _has_fa: _real_fa_module = sys.modules.get("flash_attn") sys.modules["flash_attn"] = None try: from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps finally: if _has_fa: if _real_fa_module is not None: sys.modules["flash_attn"] = _real_fa_module else: del sys.modules["flash_attn"] from tqdm import tqdm from diffusion.model.nets.basic_modules import CachedGLUMBConvTemp from diffusion.model.nets.sana_blocks import CachedCausalAttention # --------------------------------------------------------------------------- # Shared helpers # --------------------------------------------------------------------------- def _inject_sliced_extras( extra: dict[str, object], kwargs: dict, num_chunk_frames: int, end_f: int, ) -> None: """Inject ``extra`` kwargs into ``kwargs``, slicing temporal dims. Tensors whose temporal axis is longer than ``num_chunk_frames`` are sliced to ``[end_f - num_chunk_frames, end_f)``. Layouts handled: * ``(B, C, T, H, W)`` — e.g. ``chunk_plucker``; sliced on dim 2. * ``(B, T, ...)`` — e.g. ``camera_conditions``; sliced on dim 1. Any key already present in ``kwargs`` is left untouched. """ begin_f = end_f - num_chunk_frames for k, v in extra.items(): if k in kwargs: continue if isinstance(v, torch.Tensor): if v.ndim == 5: kwargs[k] = v[:, :, begin_f:end_f] if v.shape[2] > num_chunk_frames else v elif v.ndim >= 3 and v.shape[1] > num_chunk_frames: kwargs[k] = v[:, begin_f:end_f] else: kwargs[k] = v else: kwargs[k] = v def _pop_extra_model_kwargs(model_kwargs: dict) -> dict: """Pop all keys from ``model_kwargs`` except ``mask`` and ``data_info``. The popped entries are the "extras" (camera tensors, RoPE caches, etc.) that need per-chunk temporal slicing before being forwarded to the model. """ extra: dict = {} for key in list(model_kwargs): if key not in ("mask", "data_info"): extra[key] = model_kwargs.pop(key) return extra # --------------------------------------------------------------------------- # Cache-slot index constants # --------------------------------------------------------------------------- # # The camctrl ``forward_long`` path uses a 10-slot KV cache per block. Two # layouts share the slot table, distinguished by slot 6 (``_SLOT_TYPE_FLAG``): # # Old layout (ChunkCausalGDN / ChunkCausalSoftmaxAttn): # 0: k, 1: v, 2: beta, 3: decay, 4: shortconv, 5: tconv, 6-9: None # # New layout (CachedChunkCausalGDN / CachedChunkCausalSoftmaxAttn): # GDN: 0: S_kv state, 1: S_z state, 2: cam_S_kv state, 3: cam K-conv state # Softmax: 0: k post-RoPE, 1: v, 2: cam_k post-UCPE, 3: cam_v post-UCPE # Both: 4: shortconv, 5: tconv, 6: type flag (1.0=state, 0.0=concat), # 9: FFN tconv state written by ``CachedGLUMBConvTemp``. _NUM_CACHE_SLOTS = 10 # 7 active + 3 reserved _SLOT_K = 0 _SLOT_V = 1 _SLOT_BETA = 2 _SLOT_DECAY = 3 _SLOT_SHORTCONV = 4 _SLOT_TCONV = 5 _SLOT_TYPE_FLAG = 6 # Old-layout concat slots (when type flag is absent). _CONCAT_SLOTS = (_SLOT_K, _SLOT_V, _SLOT_BETA, _SLOT_DECAY) # New-layout softmax concat slots (when type flag == 0.0). _SOFTMAX_CONCAT_SLOTS = (0, 1, 2, 3) _LAST_CHUNK_SLOTS = (_SLOT_SHORTCONV, _SLOT_TCONV) # --------------------------------------------------------------------------- # Base self-forcing sampler # --------------------------------------------------------------------------- class SelfForcingFlowEuler: """Chunk-causal autoregressive flow-Euler sampler with KV cache support. Walks ``base_chunk_frames``-sized chunks left-to-right. For each chunk the sampler runs the diffusion schedule, then performs one extra ``t = 0`` forward with ``save_kv_cache=True`` to write the KV cache that subsequent chunks consume. This implements the "self-forcing" recipe where every chunk is conditioned on the model's own previously-generated context. """ def __init__( self, model_fn: object, condition: torch.Tensor, uncondition: torch.Tensor, cfg_scale: float, flow_shift: float = 3.0, model_kwargs: dict | None = None, base_chunk_frames: int = 10, num_cached_blocks: int = -1, **kwargs: object, ) -> None: self.model = model_fn self.condition = condition self.uncondition = uncondition self.cfg_scale = cfg_scale self.model_kwargs = model_kwargs or {} self.mask = self.model_kwargs.pop("mask", None) self.flow_shift = flow_shift self.base_chunk_frames = base_chunk_frames self.rank = os.environ.get("RANK", 0) self.cached_modules = None # Populate ``self.cached_modules`` and ``self.num_model_blocks``. self.get_cached_modules_by_block() self.num_cached_blocks = num_cached_blocks self.use_softmax_attention = kwargs.get("use_softmax_attention", False) self.sink_token = kwargs.get("sink_token", False) def create_autoregressive_segments(self, total_frames: int) -> list[int]: """Build chunk boundaries for an autoregressive sweep. Returns a list of frame indices ``[0, c1, c2, ..., total_frames]`` of length ``num_chunks + 1`` such that chunk ``i`` covers frames ``[chunk_indices[i], chunk_indices[i + 1])``. The first chunk absorbs any remainder so subsequent chunks all have exactly ``base_chunk_frames`` frames. """ remained_frames = total_frames % self.base_chunk_frames num_chunks = total_frames // self.base_chunk_frames chunk_indices = [0] for i in range(num_chunks): cur_idx = chunk_indices[-1] + self.base_chunk_frames if i == 0: cur_idx += remained_frames chunk_indices.append(cur_idx) return chunk_indices def get_cached_modules_by_block(self) -> list[list[torch.nn.Module]]: """Locate ``CachedCausalAttention`` and ``CachedGLUMBConvTemp`` modules. The result is a list (one entry per transformer block) of the cached modules inside that block. ``self.num_model_blocks`` is set as a side effect. """ if self.cached_modules is not None: return self.cached_modules # Unwrap DDP if present. model = self.model.module if hasattr(self.model, "module") else self.model cached_modules: list[list[torch.nn.Module]] = [] def collect_from_block(block: torch.nn.Module, block_idx: int) -> list[torch.nn.Module]: attention_modules: list[torch.nn.Module] = [] conv_modules: list[torch.nn.Module] = [] def collect_recursive(module: torch.nn.Module) -> None: if isinstance(module, CachedCausalAttention): attention_modules.append(module) elif isinstance(module, CachedGLUMBConvTemp): conv_modules.append(module) for child in module.children(): collect_recursive(child) collect_recursive(block) return attention_modules + conv_modules if hasattr(model, "blocks"): blocks = model.blocks elif hasattr(model, "transformer_blocks"): blocks = model.transformer_blocks elif hasattr(model, "layers"): blocks = model.layers else: raise ValueError("Model does not have any blocks") self.num_model_blocks = len(blocks) for block_idx, block in enumerate(blocks): block_modules = collect_from_block(block, block_idx) cached_modules.append(block_modules) self.cached_modules = cached_modules return cached_modules # NOTE: SelfForcingFlowEulerCamCtrl overrides ``_initialize_kv_cache``, # ``_accumulate_softmax_kv_cache``, ``accumulate_kv_cache`` and ``sample`` # with the 10-slot dual-mode (state / concat) cache layout used by the # camctrl ``forward_long`` path. The base class keeps ``__init__``, # ``create_autoregressive_segments`` and ``get_cached_modules_by_block`` # only — the inherited entry points for CamCtrl. The non-camctrl # streaming path (6-slot softmax + 3-slot linear caches) is intentionally # not shipped in this repo. # --------------------------------------------------------------------------- # CamCtrl self-forcing sampler # --------------------------------------------------------------------------- class SelfForcingFlowEulerCamCtrl(SelfForcingFlowEuler): """SelfForcingFlowEuler with camera conditioning and first-frame anchoring. Wraps ``SelfForcingFlowEuler`` to support the camctrl ``forward_long`` API used by the streaming Sana-WM pipeline: * Camera tensors (``camera_conditions``, ``chunk_plucker``, etc.) are popped from ``model_kwargs`` at init and injected into each model call sliced to the current temporal window. * The KV cache uses the 10-slot layout with a dual-mode (state / concat) type flag at slot 6, and supports a "sink" chunk anchored at the start of the sequence when the sliding window has scrolled past chunk 0. * ``condition_frame_info`` in ``data_info`` (e.g. ``{0: 0.0}``) marks frames that should be treated as fully clean and restored after every denoising step so the per-token scheduler cannot corrupt them. """ def __init__( self, model_fn: object, condition: torch.Tensor, uncondition: torch.Tensor, cfg_scale: float, model_kwargs: dict | None = None, **kw: object, ) -> None: model_kwargs = model_kwargs or {} self._extra_model_kwargs = _pop_extra_model_kwargs(model_kwargs) super().__init__( model_fn, condition, uncondition, cfg_scale, model_kwargs=model_kwargs, **kw, ) self._patch_model() # ------------------------------------------------------------------ # Camera tensor slicing # ------------------------------------------------------------------ def _patch_model(self) -> None: """Monkey-patch ``model.forward_long`` to inject sliced camera tensors.""" extra = self._extra_model_kwargs orig_forward_long = self.model.forward_long # Keys that ``forward_long`` recomputes or doesn't need: # - ``pos_embeds``: RoPE is recomputed from (start_f, end_f) or frame_index. # - ``cam_pos_embeds``: dict of full-sequence tensors; forward_long # recomputes from sliced camera_conditions instead. # - ``chunk_index``: not used in KV-cache mode (blocks check kv_cache). # - ``frame_index``: forwarded explicitly (not via _inject_sliced_extras). _SKIP_FOR_FORWARD_LONG = frozenset({"pos_embeds", "chunk_index", "cam_pos_embeds", "frame_index"}) def _forward_long_with_extras( x: torch.Tensor, timestep: torch.Tensor, y: torch.Tensor, mask: torch.Tensor | None = None, **kwargs: object, ) -> object: end_f = kwargs.get("end_f", x.shape[2]) num_chunk_frames = x.shape[2] filtered_extra = {k: v for k, v in extra.items() if k not in _SKIP_FOR_FORWARD_LONG} _inject_sliced_extras(filtered_extra, kwargs, num_chunk_frames, end_f) return orig_forward_long(x, timestep, y, mask=mask, **kwargs) self.model.forward_long = _forward_long_with_extras # ------------------------------------------------------------------ # Sample (override for first-frame conditioning + distilled schedules) # ------------------------------------------------------------------ @torch.no_grad() def sample( self, latents: torch.Tensor, steps: int = 50, generator: torch.Generator | None = None, *, denoising_step_list: list[int] | None = None, **kwargs: object, ) -> torch.Tensor: """Sample with first-frame conditioning support. Reads ``condition_frame_info`` from ``data_info`` (e.g. ``{0: 0.0}`` means frame 0 is fully clean). For each conditioned frame that falls inside the current chunk, its latent is restored after every denoising step so the scheduler cannot corrupt it. Args: latents: Initial latent tensor of shape ``(B, C, T, H, W)``. steps: Number of denoising steps per chunk. Ignored when ``denoising_step_list`` is supplied. generator: Optional torch generator (currently unused; the scheduler is deterministic given the noise input). denoising_step_list: Optional explicit student timestep schedule (e.g. ``[1000, 967, 908, 764, 0]``). MUST end with 0. When provided, ``steps`` is ignored and the ``FlowMatchEulerDiscreteScheduler`` is set up with these exact sigmas (no shift re-applied — the schedule is taken verbatim, so it should already incorporate the teacher's ``flow_shift``). Use this for distilled students that were trained on a fixed subsampled subset of teacher timesteps. Returns: Denoised latent tensor with the same shape as ``latents``. """ for _ in self.sample_chunks( latents, steps=steps, generator=generator, denoising_step_list=denoising_step_list, **kwargs, ): pass return latents @torch.no_grad() def sample_chunks( self, latents: torch.Tensor, steps: int = 50, generator: torch.Generator | None = None, *, denoising_step_list: list[int] | None = None, **kwargs: object, ): """Streaming variant of :meth:`sample` — yields one chunk at a time. After each AR chunk completes (denoising + KV-cache save pass), yields a tuple ``(chunk_idx, latent_chunk_view, start_f, end_f)`` where ``latent_chunk_view`` is a *view* into the in-place-mutated ``latents`` tensor for the just-finished chunk. The view stays valid for the remainder of inference (subsequent chunks never overwrite earlier frames), so the orchestrator may launch downstream work on a separate CUDA stream and continue pulling chunks without copying. ``sample(latents, ...)`` is implemented as ``for _ in sample_chunks(...)`` and returns ``latents`` after exhaustion, so the legacy whole-volume API is preserved. """ # Resolve scheduler factory once (a fresh instance is built per chunk). if denoising_step_list is not None: if len(denoising_step_list) < 2 or denoising_step_list[-1] != 0: raise ValueError( "denoising_step_list must have >=2 entries and end with 0; " f"got {denoising_step_list}" ) # Drop trailing 0; FlowMatchEulerDiscreteScheduler auto-appends sigma=0. # ``shift=1.0`` keeps our explicit sigmas verbatim (no second shift). _explicit_sigmas = [float(t) / 1000.0 for t in denoising_step_list[:-1]] else: _explicit_sigmas = None device = self.condition.device do_classifier_free_guidance = self.cfg_scale > 1 batch_size, num_latent_channels, total_frames, height, width = latents.shape self.total_frames = total_frames if total_frames <= self.base_chunk_frames: raise ValueError("Please use FlowEuler for short videos") chunk_indices = self.create_autoregressive_segments(total_frames) self._chunk_indices = chunk_indices num_chunks = len(chunk_indices) - 1 kv_cache = self._initialize_kv_cache(num_chunks) kv_save_stride = int(os.environ.get("SANA_WM_STAGE1_KV_SAVE_STRIDE", "1")) if kv_save_stride < 0: raise ValueError("SANA_WM_STAGE1_KV_SAVE_STRIDE must be >= 0.") assert self.condition.shape[0] == batch_size or self.condition.shape[0] == num_chunks if self.condition.shape[0] == batch_size: self.condition = self.condition.repeat_interleave(num_chunks, dim=0) self.mask = self.mask[None].repeat_interleave(num_chunks, dim=0) if self.mask is not None else None # -- First-frame conditioning -- data_info = self.model_kwargs.pop("data_info", {}) condition_frame_info = data_info.pop("condition_frame_info", {}) # Save a clean copy of conditioned frame latents. init_latents = latents.clone() image_vae_embeds = data_info.get("image_vae_embeds", None) # Build the scheduler once (sigmas / shift don't change per chunk). if _explicit_sigmas is not None: _shared_scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0) _shared_scheduler.set_timesteps(sigmas=_explicit_sigmas, device=device) _shared_timesteps = _shared_scheduler.timesteps _shared_num_steps = len(_shared_timesteps) else: _shared_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.flow_shift) _shared_timesteps, _shared_num_steps = retrieve_timesteps(_shared_scheduler, steps, device, None) for chunk_idx in range(num_chunks): ( chunk_kv_cache, num_chunks_to_accumulate, sink_num, num_cached_frames, ) = self.accumulate_kv_cache(kv_cache, chunk_idx) prompt_embeds = self.condition[chunk_idx].unsqueeze(0) if do_classifier_free_guidance: prompt_embeds = torch.cat([self.uncondition, prompt_embeds], dim=0) mask = self.mask[chunk_idx] if self.mask is not None else None # Reuse the scheduler built outside the chunk loop; re-set # timesteps to reset its internal ``_step_index`` to zero. self.scheduler = _shared_scheduler if _explicit_sigmas is not None: self.scheduler.set_timesteps(sigmas=_explicit_sigmas, device=device) timesteps = self.scheduler.timesteps num_inference_steps = _shared_num_steps else: timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, steps, device, None) start_f = chunk_indices[chunk_idx] end_f = chunk_indices[chunk_idx + 1] end_f - start_f max(chunk_idx - self.num_cached_blocks, 0) if self.num_cached_blocks > 0 else 0 # Build frame_index for the current chunk's RoPE positions. # # In camctrl, ``CachedChunkCausalSoftmaxAttn`` caches POST-RoPE # K/V (each chunk keeps its own absolute positions). When attending # over [sink + window + current], cached tokens already have the # correct rope baked in — only the CURRENT chunk's Q/K need # rotation, which uses positions ``[start_f, end_f)``. So # frame_index here describes only the current chunk; sink_token # affects cache contents but not the per-chunk rope shape. # # We still pass frame_index (instead of just relying on # start_f/end_f) so future experiments can override per-frame # positions for the current chunk without changing this # scaffolding. frame_index: torch.Tensor | None = None rope_start_f = start_f rope_end_f = end_f if sink_num > 0: frame_index = torch.arange(start_f, end_f, device=device, dtype=torch.long) # Shallow copy — we only need to override image_vae_embeds per chunk. local_data_info = dict(data_info) if image_vae_embeds is not None: local_data_info["image_vae_embeds"] = image_vae_embeds[:, :, start_f:end_f] # Identify conditioned frames inside this chunk (local indices). chunk_frames = end_f - start_f cond_local_indices: list[int] = [] for frame_idx in condition_frame_info: if start_f <= frame_idx < end_f: cond_local_indices.append(frame_idx - start_f) # Build a per-frame mask instead of a full (B,C,F,H,W) tensor. # The model consumes frame-level timesteps and the scheduler uses # the same frame value broadcast over spatial tokens. condition_frame_mask = None if cond_local_indices: condition_frame_mask = torch.zeros( batch_size, chunk_frames, device=device, dtype=torch.float32, ) for loc in cond_local_indices: condition_frame_mask[:, loc] = 1.0 spatial_tokens = height * width for i, t in tqdm( list(enumerate(timesteps)), disable=os.getenv("DPM_TQDM", "False") == "True", desc=f"Processing chunk {chunk_idx}", ): latent_model_input = ( torch.cat([latents[:, :, start_f:end_f]] * 2) if do_classifier_free_guidance else latents[:, :, start_f:end_f] ) # Keep the timestep on device without `.item()` syncs, but # avoid materialising a full channel x spatial mask. t_dev = t.to(device=device, dtype=torch.float32).reshape(1) if condition_frame_mask is None: timestep_frames = t_dev.expand(batch_size, chunk_frames) per_token_timesteps = t_dev.expand(batch_size, chunk_frames * spatial_tokens) else: timestep_frames = (1.0 - condition_frame_mask) * t_dev per_token_timesteps = ( timestep_frames[:, :, None] .expand( batch_size, chunk_frames, spatial_tokens, ) .reshape(batch_size, -1) ) timestep_tensor_model = timestep_frames[:, None, :] if do_classifier_free_guidance: timestep_tensor_model = torch.cat([timestep_tensor_model, timestep_tensor_model], dim=0) noise_pred, _ = self.model( latent_model_input, timestep_tensor_model, prompt_embeds, start_f=rope_start_f, end_f=rope_end_f, frame_index=frame_index, save_kv_cache=False, kv_cache=chunk_kv_cache, mask=mask, data_info=local_data_info, **self.model_kwargs, ) if isinstance(noise_pred, Transformer2DModelOutput): noise_pred = noise_pred[0] if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.cfg_scale * (noise_pred_text - noise_pred_uncond) # Per-token scheduler step with ``per_token_timesteps`` (same # reshape convention as ChunkFlowEuler). latents_dtype = latents.dtype chunk_latents_cur = latents[:, :, start_f:end_f] chunk_shape = chunk_latents_cur.shape denoised = self.scheduler.step( -noise_pred.reshape( batch_size, num_latent_channels, -1, ).transpose(1, 2), t, chunk_latents_cur.reshape( batch_size, num_latent_channels, -1, ).transpose(1, 2), per_token_timesteps=per_token_timesteps, return_dict=False, )[0] denoised = denoised.transpose(1, 2).reshape(chunk_shape) latents[:, :, start_f:end_f] = denoised # Safety: explicitly restore conditioned frames in case of # numerical drift from the per-token scheduler step. for loc in cond_local_indices: latents[:, :, start_f + loc] = init_latents[:, :, start_f + loc] if latents.dtype != latents_dtype: latents = latents.to(latents_dtype) # KV cache save pass — populates the chunk's clean-sigma K/V for # future chunks' self-attention. A stride >1 is an experimental # Stage-1-only approximation; stage 2 still refines every chunk. do_kv_save = kv_save_stride == 1 or (kv_save_stride > 1 and chunk_idx % kv_save_stride == 0) if kv_save_stride == 0: do_kv_save = bool(self.sink_token and chunk_idx == 0) if do_kv_save: latent_model_input = ( torch.cat([latents[:, :, start_f:end_f]] * 2) if do_classifier_free_guidance else latents[:, :, start_f:end_f] ) timestep = torch.zeros(latent_model_input.shape[0], device=device) noise_pred, updated_kv_cache = self.model( latent_model_input, timestep, prompt_embeds, start_f=rope_start_f, end_f=rope_end_f, frame_index=frame_index, save_kv_cache=True, kv_cache=chunk_kv_cache, mask=mask, data_info=local_data_info, **self.model_kwargs, ) kv_cache[chunk_idx] = updated_kv_cache else: kv_cache[chunk_idx] = [[None] * _NUM_CACHE_SLOTS for _ in range(self.num_model_blocks)] yield chunk_idx, latents[:, :, start_f:end_f], start_f, end_f # ------------------------------------------------------------------ # KV cache management # ------------------------------------------------------------------ def _initialize_kv_cache(self, num_chunks: int) -> list[list[list[torch.Tensor | None]]]: """Create empty 10-slot cache: ``kv_cache[chunk][block] = [None]*10``.""" return [[[None] * _NUM_CACHE_SLOTS for _ in range(self.num_model_blocks)] for _ in range(num_chunks)] def accumulate_kv_cache(self, kv_cache: list, chunk_idx: int): """Override parent dispatcher to always use 10-slot cache logic.""" return self._accumulate_softmax_kv_cache(kv_cache, chunk_idx) def _accumulate_softmax_kv_cache( self, kv_cache: list, chunk_idx: int, ) -> tuple[list, int, int, int]: """Accumulate KV cache for chunk ``chunk_idx``. Two cache layouts are supported, distinguished by slot 6 (type flag): * **State-based** (type flag == 1.0, GDN blocks): slots 0-3 hold recurrent states from the last chunk — no concatenation needed. ``sink_token`` has no effect here: state already encodes full history. * **Concat-based** (type flag == 0.0, Softmax blocks): slots 0-3 hold K/V tensors concatenated across cached chunks. When ``self.sink_token`` is set and the sliding window has scrolled past chunk 0, chunk 0 is always retained at the front (sink anchor). * **Legacy** (type flag absent): falls back to the old concat logic using beta/decay detection. Sink behavior mirrors the softmax path. Slot 4 (shortconv) always comes from the preceding chunk. Slot 9 (``kv_cache[-1]``) holds the FFN tconv state written by ``CachedGLUMBConvTemp`` and comes from the preceding chunk. Returns: ``(cur_kv_cache, num_chunks_accumulated, sink_num, num_cached_frames)``. """ if chunk_idx == 0: return kv_cache[0], 0, 0, 0 cur_kv_cache = kv_cache[chunk_idx] # Clamp to >= 0: when ``chunk_idx < num_cached_blocks`` the window has # not yet slid past chunk 0, so the effective start is 0. start_chunk_idx = max(chunk_idx - self.num_cached_blocks, 0) if self.num_cached_blocks > 0 else 0 # Sink-aware iteration order: when ``num_cached_blocks`` slid past # chunk 0, prepend chunk 0 as a permanent anchor. sink_num = 0 valid_cached_chunks = list(range(start_chunk_idx, chunk_idx)) if self.sink_token and self.num_cached_blocks > 0: s = max(chunk_idx - self.num_cached_blocks + 1, 0) if s > 0: valid_cached_chunks = [0] + list(range(s, chunk_idx)) sink_num = self._chunk_indices[1] - self._chunk_indices[0] valid_cached_chunks = [ i for i in valid_cached_chunks if kv_cache[i][0][_SLOT_K] is not None or kv_cache[i][0][_SLOT_TYPE_FLAG] is not None ] # Count cached frames in latent units (independent of patch_size). The # sampler builds frame_index in latent units so this stays consistent. num_cached_frames = sum(self._chunk_indices[i + 1] - self._chunk_indices[i] for i in valid_cached_chunks) prev_cache_idx = valid_cached_chunks[-1] if valid_cached_chunks else chunk_idx for block_id in range(self.num_model_blocks): prev_last = kv_cache[prev_cache_idx][block_id] # Detect cache layout from type flag (slot 6). type_flag = prev_last[_SLOT_TYPE_FLAG] if prev_last[_SLOT_TYPE_FLAG] is not None else None type_flag_value = None if type_flag is not None: type_flag_value = float(type_flag.item()) if isinstance(type_flag, torch.Tensor) else float(type_flag) if type_flag_value is not None and type_flag_value > 0.5: # --- State-based (GDN): last chunk's state is the full history --- # NOTE: ``CachedGLUMBConvTemp`` writes tconv state to # ``kv_cache[-1]`` (slot 9), not ``_SLOT_TCONV`` (slot 5). We # must read from [-1] and place into [-1] so the MLP finds it # on the next chunk. cur_kv_cache[block_id] = [ prev_last[0], # S_kv state (or accumulated softmax k) prev_last[1], # S_z state (or accumulated softmax v) prev_last[2], # cam_S_kv state prev_last[3], # camera K ShortConv state prev_last[_SLOT_SHORTCONV], # ShortConv state None, # (slot 5 unused) prev_last[_SLOT_TYPE_FLAG], # type flag None, None, prev_last[-1], # FFN tconv state (slot 9) ] elif type_flag_value is not None: # --- Concat-based (Softmax): concatenate K/V across chunks --- acc: list[torch.Tensor | None] = [None] * _NUM_CACHE_SLOTS for i in valid_cached_chunks: prev = kv_cache[i][block_id] if prev[0] is None: continue for s in _SOFTMAX_CONCAT_SLOTS: if prev[s] is None: continue # Softmax K/V are (B, H, N, D) — concat along dim 2. if acc[s] is None: acc[s] = prev[s] else: acc[s] = torch.cat([acc[s], prev[s]], dim=2) cur_kv_cache[block_id] = [ acc[0], # accumulated k acc[1], # accumulated v acc[2], # accumulated cam_k acc[3], # accumulated cam_v prev_last[_SLOT_SHORTCONV], # ShortConv state (last chunk) None, # (slot 5 unused) prev_last[_SLOT_TYPE_FLAG], # type flag None, None, prev_last[-1], # FFN tconv state (slot 9) ] else: # --- Legacy layout (no type flag): old concat logic --- acc_legacy: list[torch.Tensor | None] = [None] * _NUM_CACHE_SLOTS for i in valid_cached_chunks: prev = kv_cache[i][block_id] if prev[_SLOT_K] is None: continue is_gdn = prev[_SLOT_BETA] is not None kv_cat_dim = -1 if is_gdn else 2 for s in _CONCAT_SLOTS: if prev[s] is None: continue if s in (_SLOT_K, _SLOT_V): cat_dim = kv_cat_dim elif s == _SLOT_BETA: cat_dim = 2 else: cat_dim = -1 if acc_legacy[s] is None: acc_legacy[s] = prev[s] else: acc_legacy[s] = torch.cat([acc_legacy[s], prev[s]], dim=cat_dim) cur_kv_cache[block_id] = [ acc_legacy[_SLOT_K], acc_legacy[_SLOT_V], acc_legacy[_SLOT_BETA], acc_legacy[_SLOT_DECAY], prev_last[_SLOT_SHORTCONV], None, # (slot 5 unused) None, None, None, prev_last[-1], # FFN tconv state (slot 9) ] # Evict cached chunks outside the (possibly sink-augmented) window. if self.num_cached_blocks > 0: kept = set(valid_cached_chunks) for i in range(chunk_idx): if i not in kept: kv_cache[i][block_id] = [None] * _NUM_CACHE_SLOTS return ( cur_kv_cache, chunk_idx - start_chunk_idx, sink_num, num_cached_frames, )