# 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 """Triton-kernel dispatch wrappers for streaming chunk-causal inference. These are pure helpers that take an attention-block ``layer`` instance and the current chunk's tensors and run the fused kernels in :mod:`diffusion.model.ops.fused_gdn`, :mod:`diffusion.model.ops.fused_gdn_chunkwise`, and :mod:`diffusion.model.ops.fused_cam_gdn`. The ``nn.Module`` wrappers live in :mod:`diffusion.model.nets.sana_gdn_blocks` (``CachedChunkCausalGDN`` / ``CachedChunkCausalSoftmaxAttn``) and :mod:`diffusion.model.nets.sana_gdn_camctrl_blocks` (``CachedChunkCausalGDNUCPESinglePathLiteLA`` / ``CachedSoftmaxUCPESinglePathLiteLA``) and call into here. Cache slot layout (10 slots per attention block, shared with the scheduler). Slot 6 distinguishes GDN (state-based) from softmax (concat-based) blocks. .. list-table:: :header-rows: 1 * - Slot - GDN blocks - Softmax blocks * - 0 - S_kv state (B, H, D, D) - k post-RoPE (B, H, N, D) * - 1 - S_z state (B, H, D, 1) - v (B, H, N, D) * - 2 - cam_S_kv state (B, H_c, D_c, D_c) - cam_k post-UCPE (B, H_c, N, D_c) * - 3 - camera K ShortConv state (B*S, K-1, C) - cam_v post-UCPE (B, H_c, N, D_c) * - 4 - ShortConv K state (B*S, K-1, C) - None (no conv for softmax) * - 5 - reserved - reserved * - 6 - type flag: 1.0 - type flag: 0.0 * - 7-8 - reserved - reserved * - 9 - tconv state (handled by CachedGLUMBConvTemp) - tconv state """ from __future__ import annotations import os import torch import torch.nn as nn import torch.nn.functional as F import triton from fla.modules import ShortConvolution from diffusion.model.nets.sana_camctrl_blocks import _prepare_ray_apply_fns from diffusion.model.ops.fused_cam_gdn import ( _invert_SE3, _prepare_ucpe_rope_tables, _process_camera_conditions_raymats_only, _torch_cam_scan_reference, _torch_cam_scan_single_chunk, cam_prep_func, cam_prep_func_with_grad, ) from diffusion.model.ops.fused_gdn import fused_qk_inv_rms, prepare_rope_tables from diffusion.model.ops.fused_gdn_chunkwise import ( _default_dot_prec, cam_scan_chunkwise, phase_a, phase_b_triton, phase_c, ) from diffusion.model.ops.fused_gdn_chunkwise_cuda import cam_scan_chunkwise_cuda from diffusion.model.ops.fused_gdn_chunkwise_stateful_raw import fused_gdn_chunkwise_stateful_raw_autograd # --------------------------------------------------------------------------- # Cache slot indices (must match scheduler constants) # --------------------------------------------------------------------------- _SLOT_FWD_KV = 0 _SLOT_FWD_Z = 1 _SLOT_CAM = 2 _SLOT_CAM_AUX = 3 _SLOT_SHORTCONV = 4 _SLOT_TCONV = 5 # NOTE: CachedGLUMBConvTemp actually writes to kv_cache[-1] (slot 9), not slot 5! _SLOT_TYPE_FLAG = 6 _TYPE_STATE = 1.0 # GDN: state-based cache _TYPE_CONCAT = 0.0 # Softmax: concat-based cache def _slice_rope_to_current_chunk(rotary_emb: torch.Tensor, current_n: int) -> torch.Tensor: """Slice rotary embedding freqs to the trailing ``current_n`` token positions. When ``sink_token=true``, upstream rope is built for sink + current chunk positions (covers ``frame_index.numel()`` frames). But q/k inside the cached chunk-causal attention only cover the current chunk — sink K is either pre-rotated in S_kv (linear attn) or pre-rotated in kv_cache K (softmax attn). Slicing the trailing portion of ``rotary_emb`` aligns it with current-chunk q/k. If sizes already match (e.g. rolling_rope path that generates rope only for the current chunk's frame range), this is a no-op. """ rope_n = rotary_emb.shape[-2] if rope_n == current_n: return rotary_emb if rope_n < current_n: raise RuntimeError( f"rotary_emb has {rope_n} positions, smaller than current chunk's " f"{current_n}; cannot slice." ) return rotary_emb[..., -current_n:, :] # --------------------------------------------------------------------------- # Cached temporal short convolution # --------------------------------------------------------------------------- def _cached_temporal_short_conv( x: torch.Tensor, conv: ShortConvolution, HW: tuple[int, int, int], conv_cache: torch.Tensor | None, save_cache: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Short conv with cached left context: forward-cached + backward-isolated. Mirrors ``ChunkCausalGDN._apply_temporal_short_conv`` but replaces the global forward causal conv with a cache-aware version. Uses the same ``ShortConvolution.forward()`` (Triton/CUDA backend) as the training path for bit-exact numerical parity. The only difference is that the forward causal pass prepends cached left context instead of starting from zeros. Args: x: ``(B, N, C)`` where ``N = T * S``. conv: FLA ``ShortConvolution`` (depthwise causal Conv1d). HW: ``(T, H, W)``. conv_cache: ``(B*S, K-1, C)`` from previous chunk, or None. save_cache: Whether to return a new cache for the next chunk. Returns: (output, new_cache): output ``(B, N, C)``, new_cache ``(B*S, K-1, C)`` or None. """ T, H, W = HW S = H * W B_orig, N, C = x.shape dtype_in = x.dtype K = conv.weight.shape[-1] # Reshape to temporal: (B*S, T, C). x_t = x.reshape(B_orig, T, S, C).permute(0, 2, 1, 3).contiguous().reshape(B_orig * S, T, C) # --- Forward causal conv with cache --- # Use ShortConvolution.forward() (Triton/CUDA kernel) for exact numerical # parity with ChunkCausalGDN._apply_temporal_short_conv. if conv_cache is not None: # Prepend cached left context and run full causal conv, then slice. x_fwd_in = torch.cat([conv_cache.to(x_t.dtype), x_t], dim=1) y_fwd_full, _ = conv(x_fwd_in) y_fwd = y_fwd_full[:, K - 1 :, :] # drop positions from cached prefix else: y_fwd, _ = conv(x_t) # --- Backward conv (isolated within current chunk) --- # Same as ChunkCausalGDN._backward_causal_conv_per_chunk for a single chunk: # flip → causal conv → flip back. y_bwd_flipped, _ = conv(x_t.flip(1)) y_bwd = y_bwd_flipped.flip(1) # (B*S, T, C) # --- Center tap --- w_center = conv.weight[:, 0, -1] # (C,) center_term = x_t * w_center.unsqueeze(0).unsqueeze(0) y = y_fwd + y_bwd - center_term # Save cache: last K-1 timesteps of the conv INPUT (for next chunk's left context). new_cache: torch.Tensor | None = None if save_cache and K > 1: new_cache = x_t[:, -(K - 1) :, :].detach().clone() # Reshape back to (B, N, C). y = y.reshape(B_orig, S, T, C).permute(0, 2, 1, 3).reshape(B_orig, N, C) if y.dtype != dtype_in: y = y.to(dtype_in) return y, new_cache # --------------------------------------------------------------------------- # Fused Triton scan (chunk-causal main GDN branch) # --------------------------------------------------------------------------- def _gdn_main_triton( layer, qkv: torch.Tensor, beta: torch.Tensor, decay: torch.Tensor, rotary_emb: torch.Tensor | None, HW: tuple[int, int, int], S_kv_prev: torch.Tensor | None, S_z_prev: torch.Tensor | None, save_kv_cache: bool, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: """Run the chunk-causal main GDN scan through the fused Triton chunkwise kernels. Uses two fused Triton calls (forward-with-state + per-chunk reverse) on a shared ``qkv`` prep. The chunk-causal layout (forward seeded from cached state, backward isolated per chunk) is NOT what the bidir convenience wrapper does (``fused_bigdn_bidi_chunkwise`` sums both directions in-kernel and assumes neither has state). We call ``phase_a`` once, ``phase_b_triton`` twice (direction=1 stateful, direction=2 stateless), accumulate in ``phase_c`` via ``accumulate=True``, then divide. Args: layer: A :class:`CachedChunkCausalGDN` instance — for q_norm / k_norm weights, eps, kernel_func, etc. qkv: ``(B, N, 3, H, D)`` raw QKV (post short-conv K). Triton applies RMSNorm + ReLU + K-scale + RoPE inside the kernel. beta: ``(B, H, F)`` or ``(B, H, F, S)`` per-frame gates (input dtype). decay: ``(B, H, F)`` per-frame gates. rotary_emb: complex rotary frequencies for the current chunk. HW: ``(T=F, H, W)``; with ``S = H * W``. S_kv_prev: cached forward-scan ``(B, H, D, D)`` state from the prior chunk, or ``None`` for the first chunk. S_z_prev: cached forward-scan ``(B, H, D, 1)`` state, or ``None``. save_kv_cache: when ``True``, return ``(out, S_kv_new, S_z_new)``; otherwise return ``(out, None, None)``. Returns: ``(out, S_kv_new, S_z_new)`` where ``out`` is ``(B, N, H, D)`` post-divide in the kernel's dot_precision dtype (fp32 or bf16). """ B, N, three, H, D = qkv.shape assert three == 3, f"qkv last-3 dim must be 3 (q,k,v); got shape {qkv.shape}" T, H_sp, W_sp = HW S = H_sp * W_sp assert N == T * S, f"N={N} != T*S={T * S} for HW={HW}" C = H * D # ---- RMS norm parameters ---- if isinstance(layer.q_norm, nn.Identity): q_nw = torch.ones(C, device=qkv.device, dtype=torch.float32) k_nw = torch.ones(C, device=qkv.device, dtype=torch.float32) norm_eps = 1e-5 else: q_nw = layer.q_norm.weight.float().contiguous() k_nw = layer.k_norm.weight.float().contiguous() norm_eps = float(getattr(layer.q_norm, "eps", 1e-5)) # ---- RoPE tables for the current chunk ---- if rotary_emb is None: rope_cos = torch.ones(N, D, device=qkv.device, dtype=torch.float32) rope_sin = torch.zeros(N, D, device=qkv.device, dtype=torch.float32) else: rope_cur = _slice_rope_to_current_chunk(rotary_emb, N) rope_cos, rope_sin = prepare_rope_tables(rope_cur, N, D, qkv.device) # ---- K scale (same convention as torch path: D^-1/2 * S^-1/2) ---- k_scale = (D**-0.5) * (S**-0.5) # ---- Beta broadcast convention: kernels accept (B,H,F) or (B,H,F,S) ---- beta_c = beta.contiguous() decay_c = decay.contiguous() dot_prec = _default_dot_prec() if torch.is_grad_enabled() and qkv.requires_grad: forward = fused_gdn_chunkwise_stateful_raw_autograd( qkv, beta_c, decay_c, q_nw, k_nw, rope_cos, rope_sin, T, S, init_state_kv=S_kv_prev, init_state_z=S_z_prev, k_scale=k_scale, norm_eps=norm_eps, dot_precision=dot_prec, save_final_state=save_kv_cache, direction=1, ) reverse = fused_gdn_chunkwise_stateful_raw_autograd( qkv, beta_c, decay_c, q_nw, k_nw, rope_cos, rope_sin, T, S, k_scale=k_scale, norm_eps=norm_eps, dot_precision=dot_prec, direction=2, ) num_fwd, den_fwd = forward[:2] num_rev, den_rev = reverse denominator = (den_fwd.float() + den_rev.float()).permute(0, 2, 1).unsqueeze(-1) out = ((num_fwd.float() + num_rev.float()) / (denominator + float(layer.eps))).to(qkv.dtype) if save_kv_cache: return out, forward[2], forward[3] return out, None, None # ---- inv RMS (single fused launch over Q and K halves of qkv) ---- q_inv_rms, k_inv_rms = fused_qk_inv_rms(qkv, eps=norm_eps) # ---- Phase A: shared prep for both directions ---- I_P_kv, A_buf, I_P_z, B_z = phase_a( qkv, beta_c, q_inv_rms, k_inv_rms, q_nw, k_nw, rope_cos, rope_sin, F=T, S=S, k_scale=k_scale, norm_eps=norm_eps, dot_precision=dot_prec, ) # ---- Pad caller-supplied (B,H,D,D)/(B,H,D,1) state to padded layout ---- BLOCK_D = I_P_kv.shape[-1] init_kv_padded = None init_z_padded = None if S_kv_prev is not None: sk = S_kv_prev sk = sk.to(torch.float32) if sk.dtype != torch.float32 else sk B_s, H_s, D_in, D_out = sk.shape if D_in != BLOCK_D or D_out != BLOCK_D: init_kv_padded = F.pad( sk.transpose(-1, -2).reshape(B_s * H_s, D_out, D_in), (0, BLOCK_D - D_in, 0, BLOCK_D - D_out), ).contiguous() else: init_kv_padded = sk.transpose(-1, -2).reshape(B_s * H_s, BLOCK_D, BLOCK_D).contiguous() sz = S_z_prev.squeeze(-1) if S_z_prev.dim() == 4 else S_z_prev sz = sz.to(torch.float32) if sz.dtype != torch.float32 else sz Bz, Hz, Dz = sz.shape if Dz != BLOCK_D: init_z_padded = F.pad(sz.reshape(Bz * Hz, Dz), (0, BLOCK_D - Dz)).contiguous() else: init_z_padded = sz.reshape(Bz * Hz, BLOCK_D).contiguous() # ---- Phase B forward (direction=1) with state ---- if save_kv_cache: M_fwd, z_fwd, _, _, final_kv, final_z = phase_b_triton( I_P_kv, A_buf, I_P_z, B_z, decay_c, F=T, dot_precision=dot_prec, direction=1, init_state_kv=init_kv_padded, init_state_z=init_z_padded, return_final_state=True, ) else: M_fwd, z_fwd, _, _ = phase_b_triton( I_P_kv, A_buf, I_P_z, B_z, decay_c, F=T, dot_precision=dot_prec, direction=1, init_state_kv=init_kv_padded, init_state_z=init_z_padded, ) # ---- Phase B reverse (direction=2) — per-chunk, no state ---- _, _, M_rev, z_rev = phase_b_triton( I_P_kv, A_buf, I_P_z, B_z, decay_c, F=T, dot_precision=dot_prec, direction=2, ) del I_P_kv, A_buf, I_P_z, B_z # ---- Phase C: fwd output, then accumulate rev output into same buffers ---- num_out, den_out = phase_c( qkv, q_inv_rms, q_nw, rope_cos, rope_sin, M_fwd, z_fwd, F=T, S=S, dot_precision=dot_prec, ) phase_c( qkv, q_inv_rms, q_nw, rope_cos, rope_sin, M_rev, z_rev, F=T, S=S, dot_precision=dot_prec, num_out=num_out, den_out=den_out, accumulate=True, ) del M_fwd, z_fwd, M_rev, z_rev # ---- Divide: (B, H, N) -> (B, N, H, 1) for broadcast over D ---- eps = float(layer.eps) total_den = den_out.float().permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) out = (num_out.float() / (total_den + eps)).to(qkv.dtype) # (B, N, H, D) del num_out, den_out, total_den # ---- Unpad final state back to (B, H, D, D) / (B, H, D, 1) ---- if save_kv_cache: S_kv_new = final_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].transpose(-1, -2).contiguous() S_z_new = final_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() return out, S_kv_new, S_z_new return out, None, None def _cam_main_triton( q_cam_trans: torch.Tensor, k_cam_trans: torch.Tensor, v_cam_trans: torch.Tensor, beta: torch.Tensor, decay: torch.Tensor, cam_S_kv_prev: torch.Tensor | None, save_kv_cache: bool, T: int, S: int, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Run the camera-branch single-path delta-rule chunk-causal scan with two ``cam_scan_chunkwise`` calls (forward seeded from cached state + per-chunk reverse). The kernel expects fp32 q/k/v in ``(B, H, D, N)`` layout (already cam-prep'd: RMSNorm + ReLU + UCPE + RoPE). State is stored as ``(B*H, BLOCK_D, BLOCK_D)`` fp32 internally; we accept/return the ``(B, H, D, D)`` torch-format used by the kv_cache slot. Args: q_cam_trans, k_cam_trans, v_cam_trans: ``(B, H, D, N)`` — cam-prep'd inputs. Cast to fp32 if not already. beta: ``(B, H, F)`` or ``(B, H, F, S)``. decay: ``(B, H, F)``. cam_S_kv_prev: ``(B, H, D, D)`` fp32 cached state, or ``None``. save_kv_cache: when ``True``, return ``(out, cam_S_kv_new)``. T, S: frames and spatial-tokens-per-frame (``N = T * S``). Returns: ``(out, cam_S_kv_new)`` with ``out`` shaped ``(B, H, D, N)`` fp32 (fwd + per-chunk bwd combined) and ``cam_S_kv_new`` shaped ``(B, H, D, D)`` fp32 (or ``None`` when not saving). """ if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in (q_cam_trans, k_cam_trans, v_cam_trans)): beta_f = beta.float().contiguous() if beta_f.ndim == 3: beta_f = beta_f.unsqueeze(-1).expand(-1, -1, -1, S).contiguous() decay_f = decay.float().contiguous() # The fused cache stores M as (K feature, V feature), while the # torch recurrence uses state @ k and therefore keeps (V, K). init_state = cam_S_kv_prev.transpose(-1, -2) if cam_S_kv_prev is not None else None forward, final_state = _torch_cam_scan_single_chunk( q_cam_trans.float().contiguous(), k_cam_trans.float().contiguous(), v_cam_trans.float().contiguous(), beta_f, decay_f, init_state=init_state, return_final_state=True, ) reverse = _torch_cam_scan_reference( q_cam_trans.float().contiguous(), k_cam_trans.float().contiguous(), v_cam_trans.float().contiguous(), beta_f, decay_f, reverse=True, ) final_state = final_state.transpose(-1, -2).contiguous() if save_kv_cache else None return forward + reverse, final_state # CUDA cam scan on by default; set SANA_GDN_CUDA=0 to force Triton. scan = cam_scan_chunkwise_cuda if os.environ.get("SANA_GDN_CUDA", "1") != "0" else None scan = scan or cam_scan_chunkwise B, H, D, N = q_cam_trans.shape assert N == T * S, f"N={N} != T*S={T * S}" BLOCK_D = triton.next_power_of_2(D) # ---- Inputs: fp32 contiguous (kernel hard-requires this). ---- q32 = q_cam_trans.float().contiguous() if q_cam_trans.dtype != torch.float32 else q_cam_trans.contiguous() k32 = k_cam_trans.float().contiguous() if k_cam_trans.dtype != torch.float32 else k_cam_trans.contiguous() v32 = v_cam_trans.float().contiguous() if v_cam_trans.dtype != torch.float32 else v_cam_trans.contiguous() # ---- Beta: kernel expects (B, H, F, S) per docstring. ---- if beta.ndim == 3: beta_c = beta.unsqueeze(-1).expand(B, H, T, S).contiguous().float() else: beta_c = beta.contiguous().float() decay_c = decay.contiguous().float() # ---- Pad caller-supplied (B, H, D, D) to (B*H, BLOCK_D, BLOCK_D) fp32. ---- init_state = None if cam_S_kv_prev is not None: sk = cam_S_kv_prev.to(torch.float32) if cam_S_kv_prev.dtype != torch.float32 else cam_S_kv_prev if D != BLOCK_D: init_state = F.pad( sk.reshape(B * H, D, D), (0, BLOCK_D - D, 0, BLOCK_D - D), ).contiguous() else: init_state = sk.reshape(B * H, BLOCK_D, BLOCK_D).contiguous() # ---- Forward scan with state ---- if save_kv_cache: out_fwd, final_state = scan( q32, k32, v32, beta_c, decay_c, reverse=False, init_state=init_state, save_final_state=True, ) else: out_fwd = scan( q32, k32, v32, beta_c, decay_c, reverse=False, init_state=init_state, save_final_state=False, ) final_state = None # ---- Backward scan (per-chunk isolated; no state) ---- out_bwd = scan( q32, k32, v32, beta_c, decay_c, reverse=True, init_state=None, save_final_state=False, ) out = out_fwd + out_bwd # (B, H, D, N) fp32 if final_state is None: return out, None # Cam state is stored as M[K_feat, V_feat] (row-major D_K, D_V) — NO # transpose unlike main GDN (which transposes on save/load). Unpad to # the (B, H, D, D) shape callers expect for the kv_cache slot. cam_S_kv_new = final_state.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D].contiguous() return out, cam_S_kv_new # --------------------------------------------------------------------------- # Fused Triton cam-prep (RMSNorm + ReLU + K-scale + UCPE 4x4 + RoPE) # --------------------------------------------------------------------------- def _cam_prep_triton( layer, x: torch.Tensor, HW: tuple[int, int, int], camera_conditions: torch.Tensor, rotary_emb: torch.Tensor | None, conv_cache: torch.Tensor | None, save_cache: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, callable, torch.Tensor | None]: """Streaming cam-branch QKV prep through the bidir's fused Triton kernel. Mirrors the prep section of :meth:`BidirectionalGDNUCPESinglePathLiteLABothTriton._forward_cam_branch` (QKV linear + cam-K short conv + ``cam_prep_func`` Triton kernel + ``inflation_sq`` reshape) but applies the K conv with a per-chunk cached left context on the forward half and chunk-local context on the backward half. Args: layer: A :class:`CachedChunkCausalGDNUCPESinglePathLiteLA` instance. x: ``(B, N, C)`` input activations for the current chunk. HW: ``(T, H, W)`` token layout. camera_conditions: ``(B, T, ...)`` camera-pose tensor. rotary_emb: complex RoPE frequencies for the current chunk. conv_cache: Previous camera-K short-convolution context. save_cache: Whether to return context for the next chunk. Returns: ``(q_cam_trans, k_cam_trans, v_cam_trans, inflation_sq, apply_fn_o, new_conv_cache)`` with ``q/k/v_cam_trans`` shaped ``(B, H_cam, D_cam, N)`` in the input dtype, ``inflation_sq`` shaped ``(B, H_cam, 1, N)`` fp32, and ``apply_fn_o`` a torch closure that applies the inverse UCPE+RoPE to the scan output, and the new short-convolution cache (or ``None``). """ if layer.conv_q_cam is not None or layer.conv_v_cam is not None: raise NotImplementedError( "Triton cam-prep requires k_conv_only=True on the camera branch " "(conv_q_cam / conv_v_cam must be None)." ) B, N, _ = x.shape T, H_sp, W_sp = HW S = H_sp * W_sp H_heads = layer.cam_heads D_head = layer.cam_head_dim # ---- 1. QKV linear (fused via cat) + cam-K short conv ---- qkv_w = torch.cat([layer.q_proj_cam.weight, layer.k_proj_cam.weight, layer.v_proj_cam.weight]) qkv_b = torch.cat([layer.q_proj_cam.bias, layer.k_proj_cam.bias, layer.v_proj_cam.bias]) qkv_cam = F.linear(x, qkv_w, qkv_b) q_raw, k_raw, v_raw = qkv_cam.chunk(3, dim=-1) new_conv_cache = None if layer.conv_k_cam is not None: k_raw, new_conv_cache = _cached_temporal_short_conv( k_raw, layer.conv_k_cam, HW, conv_cache, save_cache, ) q_raw = q_raw.contiguous().view(B, N, H_heads, D_head).contiguous() k_raw = k_raw.contiguous().view(B, N, H_heads, D_head).contiguous() v_raw = v_raw.contiguous().view(B, N, H_heads, D_head).contiguous() # ---- 2. UCPE projection matrices (P / P_T / P_inv) ---- raymats = _process_camera_conditions_raymats_only(camera_conditions, B, HW, layer.patch_size) raymats = raymats.reshape(B, -1, 4, 4) P = raymats P_T = P.transpose(-1, -2).contiguous() P_inv = _invert_SE3(P).contiguous() # ---- 3. Sliced cam-branch RoPE (D/2 dims; T/H/W split halved) ---- if rotary_emb is not None: # Mirror the WAN-RoPE slicing used by the bidir kernel call site. head_dim = D_head orig_t_size = head_dim // 2 - 2 * (head_dim // 6) orig_h_size = head_dim // 6 new_head_dim = head_dim // 2 new_t_size = new_head_dim // 2 - 2 * (new_head_dim // 6) new_h_size = new_head_dim // 6 new_w_size = new_head_dim // 6 t_part = rotary_emb[..., :new_t_size] h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] w_part = rotary_emb[..., orig_t_size + orig_h_size : orig_t_size + orig_h_size + new_w_size] rotary_emb_cam = torch.cat([t_part, h_part, w_part], dim=-1) # Slice trailing N positions when upstream RoPE covers sink+current. rotary_emb_cam = _slice_rope_to_current_chunk(rotary_emb_cam, N) rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, N, D_head // 2, x.device) else: rotary_emb_cam = None rope_cos = torch.ones(N, D_head // 2, device=x.device, dtype=torch.float32) rope_sin = torch.zeros(N, D_head // 2, device=x.device, dtype=torch.float32) # ---- 4. Fused Triton prep kernel ---- q_norm_w = layer.q_norm_cam.weight.float().contiguous() k_norm_w = layer.k_norm_cam.weight.float().contiguous() k_scale = (D_head**-0.5) * (S**-0.5) norm_eps_val = float(getattr(layer.q_norm_cam, "eps", getattr(layer.q_norm_cam, "variance_epsilon", 1e-6))) prep = cam_prep_func_with_grad if torch.is_grad_enabled() and q_raw.requires_grad else cam_prep_func q_cam_trans, k_cam_trans, v_cam_trans, inflation_sq = prep( q_raw, k_raw, v_raw, q_norm_weight=q_norm_w, k_norm_weight=k_norm_w, proj_q=P_T, proj_kv=P_inv, rope_cos=rope_cos, rope_sin=rope_sin, k_scale=k_scale, norm_eps=norm_eps_val, ) inflation_sq = inflation_sq.view(B, H_heads, 1, N) # ---- 5. Inverse-UCPE closure for the scan output ---- _, _, apply_fn_o = _prepare_ray_apply_fns(head_dim=D_head, P=P, P_T=P_T, P_inv=P_inv, rotary_emb=rotary_emb_cam) return q_cam_trans, k_cam_trans, v_cam_trans, inflation_sq, apply_fn_o, new_conv_cache def _cached_gdn_forward_triton( layer, x: torch.Tensor, HW: tuple[int, int, int] | None, rotary_emb: torch.Tensor | None, apply_output_gate: bool, **kwargs: object, ) -> tuple[torch.Tensor, list]: """Cached chunk-causal forward through the fused Triton scan. Same return contract as the torch ``CachedChunkCausalGDN.forward`` cached path (``(out, kv_cache)``). Recurrent state on slots ``[_SLOT_FWD_KV, _SLOT_FWD_Z]`` and the shortconv slot ``[_SLOT_SHORTCONV]`` are updated in place exactly like the torch path so the scheduler can swap between implementations chunk-by-chunk without seeing a difference. Takes ``layer`` as an explicit argument (not ``self``) so it works whether the dispatch comes from ``CachedChunkCausalGDN.forward`` called directly or from the camctrl wrapper's ``CachedChunkCausalGDNUCPESinglePathLiteLA.forward`` which invokes ``CachedChunkCausalGDN.forward(self, ...)`` against the wrapper instance (wrapper instances don't have this helper on themselves). Guards: ``conv_q``/``conv_v`` are unsupported by the fused kernel — the streaming production checkpoint uses ``k_conv_only=True`` so this is fine in practice, but raise here if anyone tries to load a non-k_conv_only configuration through this path. """ if HW is None: raise ValueError("HW (T, H, W) must be provided.") if layer.conv_q is not None or layer.conv_v is not None: raise NotImplementedError( "Triton chunk-causal scan requires k_conv_only=True; " "got conv_q / conv_v not None." ) kv_cache = kwargs["kv_cache"] save_kv_cache = kwargs.get("save_kv_cache", False) B, N, C = x.shape T, H_sp, W_sp = HW S = H_sp * W_sp if N != T * S: raise ValueError(f"N={N} != T*S={T * S} for HW={HW}") H, D = layer.heads, layer.dim # 1. QKV projection -> (B, N, 3, H, D), made contiguous so the fused # kernel can stride-iterate over it. qkv = layer.qkv(x).reshape(B, N, 3, H, D) # 2. Short conv on K (with cache). Write the post-conv K back into # qkv[:, :, 1] so the kernel sees it as the K stream. if layer.conv_k is not None: k_flat = qkv[:, :, 1].reshape(B, N, C) k_flat, new_conv_cache = _cached_temporal_short_conv( k_flat, layer.conv_k, HW, kv_cache[_SLOT_SHORTCONV], save_kv_cache ) qkv = qkv.clone() if torch.is_grad_enabled() and qkv.requires_grad else qkv.contiguous() qkv[:, :, 1].copy_(k_flat.reshape(B, N, H, D)) if save_kv_cache: kv_cache[_SLOT_SHORTCONV] = new_conv_cache else: qkv = qkv.contiguous() # 3. Frame gates — Triton accepts (B,H,F) or (B,H,F,S); same as torch. precomputed_gates = kwargs.get("precomputed_gates", None) if precomputed_gates is not None: beta, decay = precomputed_gates else: beta, decay = layer._compute_frame_gates(x, HW) # 4. Fused Triton fwd-with-state + per-chunk rev scan. S_kv_prev = kv_cache[_SLOT_FWD_KV] S_z_prev = kv_cache[_SLOT_FWD_Z] out_4d, S_kv_new, S_z_new = _gdn_main_triton( layer, qkv, beta, decay, rotary_emb, HW, S_kv_prev, S_z_prev, save_kv_cache, ) if save_kv_cache: kv_cache[_SLOT_FWD_KV] = S_kv_new.detach().clone() kv_cache[_SLOT_FWD_Z] = S_z_new.detach().clone() kv_cache[_SLOT_TYPE_FLAG] = _TYPE_STATE # 5. Output gate + projection, matching the torch path's tail. out = out_4d.reshape(B, N, C) if apply_output_gate: out = layer._apply_output_gate(out, x) out = layer.proj(out.to(layer.proj.weight.dtype)) return out, kv_cache return out, kv_cache