wehub-resource-sync
936 lines
34 KiB
Python
936 lines
34 KiB
Python
# 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
|
|
|
|
"""GDN-based UCPE camera-control attention blocks.
|
|
|
|
Each block follows a dual-branch design:
|
|
- **Main branch**: Inherited from the corresponding GDN variant
|
|
(GDN / BidirectionalGDN / ChunkCausalGDN). ``super().forward()``
|
|
is called with ``apply_output_gate=False`` to get raw attention.
|
|
- **Camera branch**: Separate QKV projections with UCPE per-ray
|
|
transforms. Camera QK normalization uses branch-specific RMSNorm
|
|
copies (initialized from main branch) to avoid cross-branch
|
|
distribution coupling.
|
|
|
|
The two raw outputs are combined, then the shared output gate and
|
|
projection are applied once. At init the camera branch contributes
|
|
zero (``out_proj_cam`` is zero-initialized), so the model starts
|
|
identical to the base GDN.
|
|
|
|
When the GDN kernels (torch / triton) are upgraded, both branches
|
|
pick up the improvement automatically.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from fla.modules import ShortConvolution
|
|
from torch.distributed.nn import functional as dist_nn
|
|
|
|
from diffusion.distributed.context_parallel.config import cp_enabled, get_cp_group
|
|
from diffusion.model.ops.fused_streaming import (
|
|
_SLOT_CAM,
|
|
_SLOT_CAM_AUX,
|
|
_cam_main_triton,
|
|
_cam_prep_triton,
|
|
)
|
|
from diffusion.model.registry import ATTENTION_BLOCKS
|
|
|
|
from .sana_camctrl_blocks import _maybe_drop_cam_branch, prepare_prope_fns
|
|
from .sana_gdn_blocks import (
|
|
GDN,
|
|
BidirectionalGDN,
|
|
CachedChunkCausalGDN,
|
|
CachedChunkCausalSoftmaxAttn,
|
|
ChunkCausalGDN,
|
|
_forward_softmax_attn,
|
|
_forward_softmax_attn_chunk_causal,
|
|
_sdpa_maybe_chunk_causal,
|
|
_sdpa_needs_head_pad,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Softmax-block KV cache helpers.
|
|
#
|
|
# Project Q/K/V for a softmax-attention block, apply RoPE (main branch) or
|
|
# UCPE per-position transforms (cam branch), and return the post-transform
|
|
# tensors without running SDPA. The AR KV-cache uses these to stash K and V
|
|
# in a per-block cache and replay them across AR sub-steps.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _prepare_softmax_main_qkv_post_rope(
|
|
block: GDN,
|
|
x: torch.Tensor,
|
|
HW: tuple[int, int, int],
|
|
rotary_emb: torch.Tensor | None,
|
|
**kwargs: object,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.dtype]:
|
|
"""Project Q/K/V for the softmax main branch, apply norm and RoPE.
|
|
|
|
Returns post-norm, post-RoPE, post-bf16 cast tensors without running
|
|
SDPA, so the caller can either run SDPA itself or stash K/V in a cache.
|
|
|
|
Args:
|
|
block: A :class:`GDN` (or subclass) that owns the softmax-attn
|
|
params (``qkv``, ``q_norm``, ``k_norm``).
|
|
x: Input tokens of shape ``(B, N, C)``.
|
|
HW: ``(T, H, W)`` token layout.
|
|
rotary_emb: Optional RoPE table; ``None`` skips RoPE.
|
|
|
|
Returns:
|
|
``(q, k, v, dtype_orig)`` where Q/K/V are shape ``(B, H, N, D)``
|
|
and ``dtype_orig`` is the original ``x.dtype``.
|
|
"""
|
|
B, N, C = x.shape
|
|
T, H_sp, W_sp = HW
|
|
S = H_sp * W_sp
|
|
|
|
frame_valid_mask = kwargs.get("frame_valid_mask", None)
|
|
token_valid_mask, _, _ = GDN._prepare_frame_valid_masks(
|
|
frame_valid_mask,
|
|
B=B,
|
|
T=T,
|
|
S=S,
|
|
device=x.device,
|
|
dtype=x.dtype,
|
|
)
|
|
if token_valid_mask is not None:
|
|
x = x * token_valid_mask.view(B, N, 1)
|
|
|
|
qkv = block.qkv(x).reshape(B, N, 3, block.heads, block.dim)
|
|
q, k, v = qkv.unbind(2)
|
|
if token_valid_mask is not None:
|
|
m = token_valid_mask.view(B, N, 1, 1)
|
|
q, k, v = q * m, k * m, v * m
|
|
|
|
q = block.q_norm(q.reshape(B, N, C)).reshape(B, N, block.heads, block.dim)
|
|
k = block.k_norm(k.reshape(B, N, C)).reshape(B, N, block.heads, block.dim)
|
|
|
|
if rotary_emb is not None:
|
|
q_perm = q.permute(0, 2, 3, 1)
|
|
k_perm = k.permute(0, 2, 3, 1)
|
|
q_perm = GDN._apply_rotary_emb(q_perm, rotary_emb)
|
|
k_perm = GDN._apply_rotary_emb(k_perm, rotary_emb)
|
|
q = q_perm.permute(0, 3, 1, 2)
|
|
k = k_perm.permute(0, 3, 1, 2)
|
|
|
|
if token_valid_mask is not None:
|
|
m = token_valid_mask.view(B, N, 1, 1)
|
|
q, k, v = q * m, k * m, v * m
|
|
|
|
q = q.transpose(1, 2) # (B, H, N, D)
|
|
k = k.transpose(1, 2)
|
|
v = v.transpose(1, 2)
|
|
|
|
dtype_orig = x.dtype
|
|
if q.dtype == torch.float32:
|
|
q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16()
|
|
|
|
return q, k, v, dtype_orig
|
|
|
|
|
|
def _sdpa_unmasked_with_pad(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Run ``F.scaled_dot_product_attention(q, k, v)`` with FA-friendly head_dim padding.
|
|
|
|
FlashAttention-2 only supports head_dim in {32, 64, 128, 256}.
|
|
Other head_dims (e.g. 112) fall back to the math backend. We pad
|
|
head_dim up to the next supported size, run SDPA, then slice back
|
|
to the original head_dim. Mirrors the no-mask path in
|
|
:func:`_forward_softmax_attn` (lines ~3034-3061).
|
|
|
|
Args:
|
|
q, k, v: ``(B, H, N_q, D)``, ``(B, H, N_kv, D)``, ``(B, H, N_kv, D)``.
|
|
|
|
Returns:
|
|
``(B, H, N_q, D)`` attention output.
|
|
"""
|
|
D = q.shape[-1]
|
|
_need_pad = _sdpa_needs_head_pad(D)
|
|
if _need_pad:
|
|
_pad_to = 128 if D <= 128 else 256
|
|
_pad_size = _pad_to - D
|
|
q = F.pad(q, (0, _pad_size))
|
|
k = F.pad(k, (0, _pad_size))
|
|
v = F.pad(v, (0, _pad_size))
|
|
out = F.scaled_dot_product_attention(q, k, v)
|
|
if _need_pad:
|
|
out = out[..., :D]
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Base class
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _GDNUCPEBase(GDN):
|
|
"""Shared camera-branch logic for all GDN + UCPE variants.
|
|
|
|
Adds a second attention branch whose positional encoding comes from
|
|
UCPE per-ray camera transforms instead of the standard RoPE used by
|
|
the main branch.
|
|
|
|
**Camera-specific parameters** (4 Linear layers per block):
|
|
``q_proj_cam``, ``k_proj_cam``, ``v_proj_cam``, ``out_proj_cam``
|
|
|
|
**Shared with main branch** (no duplication):
|
|
QK norms, GDN gates (beta/gate/dt_bias/A_log/recall_gate),
|
|
output gate, output projection.
|
|
|
|
Requires ``cam_dim == in_dim`` and ``cam_heads == heads`` so that
|
|
all shared parameters have matching dimensions.
|
|
|
|
Subclasses provide their own ``_forward_cam_branch`` (e.g. the fused
|
|
Triton pipeline in :class:`BidirectionalGDNUCPESinglePathLiteLABothTriton`
|
|
or the cached streaming pipeline in
|
|
:class:`CachedChunkCausalGDNUCPESinglePathLiteLA`).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_dim: int,
|
|
out_dim: int,
|
|
*,
|
|
cam_dim: int,
|
|
cam_heads: int,
|
|
patch_size: tuple[int, int, int] = (1, 2, 2),
|
|
**kwargs: object,
|
|
) -> None:
|
|
# Silently swallow legacy debug-stat kwargs so older configs/checkpoints
|
|
# don't blow up on construction.
|
|
kwargs.pop("cam_debug_ratios", None)
|
|
kwargs.pop("cam_debug_log_per_block", None)
|
|
super().__init__(in_dim, out_dim, **kwargs)
|
|
|
|
self.patch_size = patch_size
|
|
self.cam_dim = cam_dim
|
|
self.cam_heads = cam_heads
|
|
self.cam_head_dim = cam_dim // cam_heads
|
|
|
|
if cam_dim != in_dim:
|
|
raise ValueError(
|
|
f"Parameter sharing requires cam_dim == in_dim, " f"got cam_dim={cam_dim}, in_dim={in_dim}."
|
|
)
|
|
if cam_heads != self.heads:
|
|
raise ValueError(
|
|
f"Parameter sharing requires cam_heads == heads, " f"got cam_heads={cam_heads}, heads={self.heads}."
|
|
)
|
|
if self.cam_head_dim % 4 != 0:
|
|
raise ValueError(
|
|
"UCPE camera branch requires cam_head_dim divisible by 4, "
|
|
f"got {self.cam_head_dim} (cam_dim={cam_dim}, cam_heads={cam_heads})."
|
|
)
|
|
|
|
# ---- Camera-specific: QKV + output projections only ----
|
|
self.q_proj_cam = nn.Linear(in_dim, cam_dim, bias=True)
|
|
self.k_proj_cam = nn.Linear(in_dim, cam_dim, bias=True)
|
|
self.v_proj_cam = nn.Linear(in_dim, cam_dim, bias=True)
|
|
self.out_proj_cam = nn.Linear(cam_dim, out_dim, bias=True)
|
|
|
|
# Keep branch-specific Q/K norms so camera statistics do not disturb the
|
|
# main branch (and vice versa). Start from identical weights.
|
|
self.q_norm_cam = deepcopy(self.q_norm)
|
|
self.k_norm_cam = deepcopy(self.k_norm)
|
|
|
|
nn.init.constant_(self.out_proj_cam.weight, 0)
|
|
nn.init.constant_(self.out_proj_cam.bias, 0)
|
|
|
|
# Short convolutions for camera branch (matching base GDN variant).
|
|
if self.conv_kernel_size > 0:
|
|
self.conv_k_cam = ShortConvolution(
|
|
hidden_size=cam_dim,
|
|
kernel_size=self.conv_kernel_size,
|
|
activation=None,
|
|
)
|
|
if self.k_conv_only:
|
|
self.conv_q_cam = None
|
|
self.conv_v_cam = None
|
|
else:
|
|
self.conv_q_cam = ShortConvolution(
|
|
hidden_size=cam_dim,
|
|
kernel_size=self.conv_kernel_size,
|
|
activation=None,
|
|
)
|
|
self.conv_v_cam = ShortConvolution(
|
|
hidden_size=cam_dim,
|
|
kernel_size=self.conv_kernel_size,
|
|
activation=None,
|
|
)
|
|
self._init_cam_short_conv_for_linear_equiv()
|
|
else:
|
|
self.conv_q_cam = None
|
|
self.conv_k_cam = None
|
|
self.conv_v_cam = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Initialization helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _init_cam_short_conv_for_linear_equiv(self) -> None:
|
|
"""Initialize camera short convs as identity to match base at step 0."""
|
|
if self.conv_k_cam is None:
|
|
return
|
|
for conv in (self.conv_q_cam, self.conv_k_cam, self.conv_v_cam):
|
|
if conv is None:
|
|
continue
|
|
with torch.no_grad():
|
|
conv.weight.zero_()
|
|
conv.weight[:, 0, -1] = 1.0
|
|
if getattr(conv, "bias", None) is not None:
|
|
conv.bias.zero_()
|
|
|
|
def init_cam_branch_weights(self) -> None:
|
|
"""Copy main-branch QKV weights into the camera branch for transfer learning."""
|
|
if self.cam_dim != self.dim * self.heads:
|
|
print(
|
|
f"Warning: Skipping init_cam_branch_weights because "
|
|
f"cam_dim ({self.cam_dim}) != dim ({self.dim}) * heads ({self.heads})"
|
|
)
|
|
return
|
|
|
|
print(f"Initializing camera branch QKV from base model QKV for {self.__class__.__name__}")
|
|
w = self.qkv.weight
|
|
b = self.qkv.bias
|
|
dim = self.cam_dim
|
|
|
|
self.q_proj_cam.weight.data.copy_(w[:dim])
|
|
self.k_proj_cam.weight.data.copy_(w[dim : 2 * dim])
|
|
self.v_proj_cam.weight.data.copy_(w[2 * dim :])
|
|
if b is not None:
|
|
self.q_proj_cam.bias.data.copy_(b[:dim])
|
|
self.k_proj_cam.bias.data.copy_(b[dim : 2 * dim])
|
|
self.v_proj_cam.bias.data.copy_(b[2 * dim :])
|
|
|
|
# Mirror main-branch Q/K norm initialization into camera-specific norms.
|
|
if hasattr(self.q_norm, "state_dict") and hasattr(self.q_norm_cam, "load_state_dict"):
|
|
self.q_norm_cam.load_state_dict(self.q_norm.state_dict(), strict=False)
|
|
if hasattr(self.k_norm, "state_dict") and hasattr(self.k_norm_cam, "load_state_dict"):
|
|
self.k_norm_cam.load_state_dict(self.k_norm.state_dict(), strict=False)
|
|
|
|
# Copy short conv weights from base to camera branch.
|
|
if self.conv_k_cam is not None and self.conv_k is not None:
|
|
self.conv_k_cam.load_state_dict(self.conv_k.state_dict())
|
|
if self.conv_q_cam is not None and self.conv_q is not None:
|
|
self.conv_q_cam.load_state_dict(self.conv_q.state_dict())
|
|
if self.conv_v_cam is not None and self.conv_v is not None:
|
|
self.conv_v_cam.load_state_dict(self.conv_v.state_dict())
|
|
|
|
|
|
class BidirectionalGDNUCPESinglePathLiteLA(_GDNUCPEBase, BidirectionalGDN):
|
|
"""Bidirectional GDN with UCPE camera conditioning (single-path delta rule).
|
|
|
|
Main branch: bidirectional GDN (inherited from :class:`BidirectionalGDN`).
|
|
Camera branch: numerator-only delta-rule recurrence over UCPE-transformed
|
|
camera tensors (RMSNorm + ReLU + UCPE 4x4 + RoPE, then a single-path scan).
|
|
|
|
This is the production base for both the bidir Triton variant
|
|
(:class:`BidirectionalGDNUCPESinglePathLiteLABothTriton`) and the streaming
|
|
chunk-causal variant (:class:`ChunkCausalGDNUCPESinglePathLiteLA`).
|
|
"""
|
|
|
|
|
|
class ChunkCausalGDNUCPESinglePathLiteLA(BidirectionalGDNUCPESinglePathLiteLA, ChunkCausalGDN):
|
|
"""Chunk-causal variant of ``BidirectionalGDNUCPESinglePathLiteLA``.
|
|
|
|
Main branch: chunk-causal GDN (inherited via MRO from
|
|
:class:`ChunkCausalGDN`). Camera branch: single-path (numerator-only)
|
|
delta rule with chunk-boundary isolation in the backward pass.
|
|
|
|
All parameter names match ``BidirectionalGDNUCPESinglePathLiteLA``
|
|
exactly, so checkpoints from bidirectional training load directly.
|
|
The only behavioral difference is that the backward recurrence in the
|
|
camera branch is isolated at chunk boundaries (decay and inputs
|
|
zeroed), preventing future chunk information from leaking into past
|
|
chunks.
|
|
"""
|
|
|
|
def _apply_temporal_short_conv(
|
|
self,
|
|
x: torch.Tensor,
|
|
conv: ShortConvolution,
|
|
HW: tuple[int, int, int],
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
"""Route short conv: chunk-causal when chunk boundaries exist, else bidirectional.
|
|
|
|
For single-chunk (``chunk_size >= T``) or no chunk_size, use the
|
|
bidirectional short conv (identical forward AND backward to the
|
|
parent class). For multi-chunk, use chunk-causal short conv to
|
|
enforce causality.
|
|
"""
|
|
chunk_size = kwargs.get("chunk_size", None)
|
|
T = HW[0]
|
|
if chunk_size is not None and chunk_size < T:
|
|
return ChunkCausalGDN._apply_temporal_short_conv(self, x, conv, HW, **kwargs)
|
|
return BidirectionalGDN._apply_temporal_short_conv(self, x, conv, HW, **kwargs)
|
|
|
|
|
|
def _prepare_cam_qkv_softmax(
|
|
self,
|
|
x: torch.Tensor,
|
|
HW: tuple,
|
|
camera_conditions: torch.Tensor,
|
|
rotary_emb: torch.Tensor | None,
|
|
*,
|
|
token_valid_mask: torch.Tensor | None = None,
|
|
**kwargs,
|
|
) -> tuple:
|
|
"""Camera branch Q/K/V for softmax attention.
|
|
|
|
Mirrors ``_GDNUCPEBase._prepare_cam_qkv`` but skips the ReLU kernel and
|
|
GDN key scaling — standard softmax SDPA provides its own 1/sqrt(d_k).
|
|
Returns ``(q, k, v, apply_fn_o)`` shaped ``(B, cam_heads, cam_head_dim, N)``.
|
|
"""
|
|
B, N, C = x.shape
|
|
|
|
if token_valid_mask is not None:
|
|
x = x * token_valid_mask.view(B, N, 1)
|
|
|
|
qkv_w = torch.cat([self.q_proj_cam.weight, self.k_proj_cam.weight, self.v_proj_cam.weight])
|
|
qkv_b = torch.cat([self.q_proj_cam.bias, self.k_proj_cam.bias, self.v_proj_cam.bias])
|
|
qkv_cam = F.linear(x, qkv_w, qkv_b)
|
|
q_cam, k_cam, v_cam = qkv_cam.chunk(3, dim=-1)
|
|
|
|
if token_valid_mask is not None:
|
|
m = token_valid_mask.view(B, N, 1)
|
|
q_cam, k_cam, v_cam = q_cam * m, k_cam * m, v_cam * m
|
|
|
|
if self.conv_q_cam is not None:
|
|
q_cam = self._apply_temporal_short_conv(q_cam, self.conv_q_cam, HW, **kwargs)
|
|
if self.conv_k_cam is not None:
|
|
k_cam = self._apply_temporal_short_conv(k_cam, self.conv_k_cam, HW, **kwargs)
|
|
if self.conv_v_cam is not None:
|
|
v_cam = self._apply_temporal_short_conv(v_cam, self.conv_v_cam, HW, **kwargs)
|
|
|
|
q_cam = self.q_norm_cam(q_cam).reshape(B, N, self.cam_heads, self.cam_head_dim)
|
|
k_cam = self.k_norm_cam(k_cam).reshape(B, N, self.cam_heads, self.cam_head_dim)
|
|
v_cam = v_cam.reshape(B, N, self.cam_heads, self.cam_head_dim)
|
|
|
|
q_cam = q_cam.permute(0, 2, 3, 1).contiguous()
|
|
k_cam = k_cam.permute(0, 2, 3, 1).contiguous()
|
|
v_cam = v_cam.permute(0, 2, 3, 1).contiguous()
|
|
|
|
cached_fns = kwargs.get("prope_fns", None)
|
|
if cached_fns is not None:
|
|
apply_fn_q, apply_fn_kv, apply_fn_o = cached_fns
|
|
else:
|
|
apply_fn_q, apply_fn_kv, apply_fn_o = prepare_prope_fns(
|
|
camctrl_type="UCPE",
|
|
head_dim=self.cam_head_dim,
|
|
camera_conditions=camera_conditions,
|
|
HW=HW,
|
|
patch_size=self.patch_size,
|
|
rotary_emb=rotary_emb,
|
|
)
|
|
|
|
q_cam_trans = apply_fn_q(q_cam.transpose(-1, -2)).transpose(-1, -2).contiguous()
|
|
kv_cam = torch.cat([k_cam, v_cam], dim=1)
|
|
kv_cam_trans = apply_fn_kv(kv_cam.transpose(-1, -2)).transpose(-1, -2).contiguous()
|
|
k_cam_trans, v_cam_trans = torch.chunk(kv_cam_trans, chunks=2, dim=1)
|
|
return q_cam_trans, k_cam_trans, v_cam_trans, apply_fn_o
|
|
|
|
|
|
def _forward_cam_branch_softmax(
|
|
self,
|
|
x: torch.Tensor,
|
|
HW: tuple,
|
|
camera_conditions: torch.Tensor,
|
|
rotary_emb: torch.Tensor | None,
|
|
frame_causal: bool,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
"""Bidirectional softmax camera branch (with UCPE transforms).
|
|
|
|
Uses ``F.scaled_dot_product_attention`` with optional invalid-key masking.
|
|
"""
|
|
B, N, _ = x.shape
|
|
T, H, W = HW
|
|
S = H * W
|
|
|
|
token_valid_mask, _, _ = self._prepare_frame_valid_masks(
|
|
kwargs.get("frame_valid_mask", None),
|
|
B=B,
|
|
T=T,
|
|
S=S,
|
|
device=x.device,
|
|
dtype=x.dtype,
|
|
)
|
|
|
|
q_cam_trans, k_cam_trans, v_cam_trans, apply_fn_o = _prepare_cam_qkv_softmax(
|
|
self,
|
|
x,
|
|
HW,
|
|
camera_conditions,
|
|
rotary_emb,
|
|
token_valid_mask=token_valid_mask,
|
|
**kwargs,
|
|
)
|
|
|
|
if token_valid_mask is not None:
|
|
m = token_valid_mask.view(B, 1, 1, N)
|
|
q_cam_trans, v_cam_trans = q_cam_trans * m, v_cam_trans * m
|
|
|
|
q_sdpa = q_cam_trans.transpose(-1, -2)
|
|
k_sdpa = k_cam_trans.transpose(-1, -2)
|
|
v_sdpa = v_cam_trans.transpose(-1, -2)
|
|
|
|
dtype_orig = x.dtype
|
|
if getattr(self, "fp32_attention", True):
|
|
q_sdpa, k_sdpa, v_sdpa = q_sdpa.float(), k_sdpa.float(), v_sdpa.float()
|
|
# SDPA / FlashAttention only supports bf16/fp16; fp32 falls back to math backend.
|
|
if q_sdpa.dtype == torch.float32:
|
|
q_sdpa, k_sdpa, v_sdpa = q_sdpa.bfloat16(), k_sdpa.bfloat16(), v_sdpa.bfloat16()
|
|
|
|
key_valid_mask = token_valid_mask
|
|
attention_t = T
|
|
q_frame_offset = 0
|
|
chunk_index = kwargs.get("chunk_index")
|
|
if cp_enabled() and get_cp_group() is not None:
|
|
cp_group = get_cp_group()
|
|
cp_rank = dist.get_rank(cp_group)
|
|
cp_world = dist.get_world_size(cp_group)
|
|
k_sdpa = torch.cat(dist_nn.all_gather(k_sdpa.contiguous(), group=cp_group), dim=2)
|
|
v_sdpa = torch.cat(dist_nn.all_gather(v_sdpa.contiguous(), group=cp_group), dim=2)
|
|
if token_valid_mask is not None:
|
|
key_valid_mask = torch.cat(dist_nn.all_gather(token_valid_mask.contiguous(), group=cp_group), dim=1)
|
|
q_frame_offset = cp_rank * T
|
|
attention_t = T * cp_world
|
|
chunk_index = kwargs.get("chunk_index_global", chunk_index)
|
|
|
|
invalid_kv_logit_bias = None
|
|
if key_valid_mask is not None and not bool(key_valid_mask.all()):
|
|
invalid_kv_logit_bias = torch.where(
|
|
key_valid_mask.bool().view(B, 1, 1, -1),
|
|
torch.zeros((), dtype=q_sdpa.dtype, device=q_sdpa.device),
|
|
torch.full((), -1e9, dtype=q_sdpa.dtype, device=q_sdpa.device),
|
|
)
|
|
|
|
chunk_size = kwargs.get("chunk_size")
|
|
need_chunk_mask = chunk_size is not None and chunk_size < attention_t
|
|
if need_chunk_mask:
|
|
out = _sdpa_maybe_chunk_causal(
|
|
q_sdpa,
|
|
k_sdpa,
|
|
v_sdpa,
|
|
need_chunk_mask=True,
|
|
T=attention_t,
|
|
S=S,
|
|
chunk_size=chunk_size,
|
|
chunk_index=chunk_index,
|
|
chunk_split_strategy=str(kwargs.get("chunk_split_strategy", "uniform")),
|
|
q_frame_offset=q_frame_offset,
|
|
attn_mask=invalid_kv_logit_bias,
|
|
)
|
|
else:
|
|
# FlashAttention-2 only supports head_dim in {32, 64, 128, 256}.
|
|
D = q_sdpa.shape[-1]
|
|
_need_pad = _sdpa_needs_head_pad(D)
|
|
if _need_pad:
|
|
_pad_to = 128 if D <= 128 else 256
|
|
_pad_size = _pad_to - D
|
|
q_sdpa = F.pad(q_sdpa, (0, _pad_size))
|
|
k_sdpa = F.pad(k_sdpa, (0, _pad_size))
|
|
v_sdpa = F.pad(v_sdpa, (0, _pad_size))
|
|
out = F.scaled_dot_product_attention(
|
|
q_sdpa,
|
|
k_sdpa,
|
|
v_sdpa,
|
|
attn_mask=invalid_kv_logit_bias,
|
|
scale=D**-0.5 if _need_pad else None,
|
|
)
|
|
if _need_pad:
|
|
out = out[..., :D]
|
|
|
|
out = out.transpose(-1, -2)
|
|
if out.dtype != dtype_orig:
|
|
out = out.to(dtype_orig)
|
|
if token_valid_mask is not None:
|
|
out = out * token_valid_mask.view(B, 1, 1, N).to(out.dtype)
|
|
out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous()
|
|
out = out.reshape(B, self.cam_dim, N).permute(0, 2, 1)
|
|
if token_valid_mask is not None:
|
|
out = out * token_valid_mask.view(B, N, 1).to(out.dtype)
|
|
return out
|
|
|
|
|
|
class _SoftmaxUCPESinglePathLiteLA(
|
|
BidirectionalGDNUCPESinglePathLiteLA,
|
|
):
|
|
"""Softmax attention with UCPE camera conditioning (single-path).
|
|
|
|
Replaces GDN recurrence with ``F.scaled_dot_product_attention``.
|
|
Automatically selects the correct masking mode based on ``chunk_size``:
|
|
|
|
- ``chunk_size is None`` or ``chunk_size >= T``: full bidirectional (no mask)
|
|
- ``chunk_size < T``: chunk-causal (full within chunks, causal across)
|
|
|
|
All parameters match the GDN variants for checkpoint compatibility.
|
|
GDN-specific parameters are present but unused in forward.
|
|
"""
|
|
|
|
def __init__(self, *args, conv_kernel_size: int = 0, **kwargs):
|
|
super().__init__(*args, conv_kernel_size=0, **kwargs)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
HW: tuple[int, int, int] | None = None,
|
|
rotary_emb: torch.Tensor | None = None,
|
|
block_mask: torch.Tensor | None = None,
|
|
camera_conditions: torch.Tensor | None = None,
|
|
chunk_size: int | None = None,
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
if HW is None:
|
|
raise ValueError("HW must be provided for softmax attention.")
|
|
if chunk_size is not None:
|
|
softmax_kwargs = dict(kwargs)
|
|
chunk_split_strategy = str(softmax_kwargs.pop("chunk_split_strategy", "uniform"))
|
|
chunk_index = softmax_kwargs.pop("chunk_index", None)
|
|
main_raw = _forward_softmax_attn_chunk_causal(
|
|
self,
|
|
x,
|
|
HW,
|
|
rotary_emb,
|
|
chunk_size=chunk_size,
|
|
chunk_split_strategy=chunk_split_strategy,
|
|
chunk_index=chunk_index,
|
|
apply_output_gate=False,
|
|
**softmax_kwargs,
|
|
)
|
|
else:
|
|
main_raw = _forward_softmax_attn(
|
|
self,
|
|
x,
|
|
HW,
|
|
rotary_emb,
|
|
frame_causal=False,
|
|
apply_output_gate=False,
|
|
**kwargs,
|
|
)
|
|
|
|
cam_contrib: torch.Tensor | int = 0
|
|
camera_conditions = _maybe_drop_cam_branch(
|
|
camera_conditions,
|
|
kwargs.get("cam_branch_drop_prob", 0.0),
|
|
self.training,
|
|
x.device,
|
|
)
|
|
if camera_conditions is not None:
|
|
if HW is None:
|
|
raise ValueError("HW must be provided for UCPE camera branch.")
|
|
cam_raw = _forward_cam_branch_softmax(
|
|
self,
|
|
x,
|
|
HW,
|
|
camera_conditions,
|
|
rotary_emb,
|
|
frame_causal=False,
|
|
chunk_size=chunk_size,
|
|
**kwargs,
|
|
)
|
|
cam_contrib = self.out_proj_cam(cam_raw)
|
|
|
|
combined = main_raw + cam_contrib
|
|
combined = self._apply_output_gate(combined, x)
|
|
return self.proj(combined.to(x.dtype))
|
|
|
|
|
|
# Aliases for backward compatibility and clear intent in mappings.
|
|
BidirectionalSoftmaxUCPESinglePathLiteLA = _SoftmaxUCPESinglePathLiteLA
|
|
ChunkCausalSoftmaxUCPESinglePathLiteLA = _SoftmaxUCPESinglePathLiteLA
|
|
|
|
|
|
# ===========================================================================
|
|
# Cached streaming variants (camera-wrapped)
|
|
# ===========================================================================
|
|
#
|
|
# Streaming-inference subclasses of the chunk-causal cam classes above.
|
|
# GDN cam-wrapper dispatches main + cam branches through the fused Triton
|
|
# kernels in :mod:`diffusion.model.ops.fused_streaming`; softmax cam-wrapper
|
|
# prepends cached cam K, V to the current chunk and runs SDPA via the
|
|
# unified ``_sdpa_maybe_chunk_causal`` helper in
|
|
# :mod:`diffusion.model.nets.sana_gdn_blocks`.
|
|
|
|
|
|
@ATTENTION_BLOCKS.register_module()
|
|
class CachedChunkCausalGDNUCPESinglePathLiteLA(ChunkCausalGDNUCPESinglePathLiteLA):
|
|
"""Cached variant of :class:`ChunkCausalGDNUCPESinglePathLiteLA`.
|
|
|
|
Main branch: :class:`CachedChunkCausalGDN` (state-based cache via
|
|
``_cached_gdn_forward_triton``).
|
|
Camera branch: ``_cam_prep_triton`` + ``_cam_main_triton`` with
|
|
``cam_S_kv`` state cached in slot 2.
|
|
"""
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
HW: tuple[int, int, int] | None = None,
|
|
rotary_emb: torch.Tensor | None = None,
|
|
block_mask: torch.Tensor | None = None,
|
|
camera_conditions: torch.Tensor | None = None,
|
|
chunk_size: int | None = None,
|
|
**kwargs: object,
|
|
) -> tuple[torch.Tensor, list]:
|
|
kv_cache = kwargs.pop("kv_cache", None)
|
|
save_kv_cache = kwargs.pop("save_kv_cache", False)
|
|
if kv_cache is None:
|
|
raise RuntimeError(
|
|
"CachedChunkCausalGDNUCPESinglePathLiteLA requires kv_cache "
|
|
"to be provided (streaming inference only)."
|
|
)
|
|
if HW is None:
|
|
raise ValueError("HW (T, H, W) must be provided.")
|
|
|
|
# Pre-compute shared gates once for main + cam branches.
|
|
precomputed_gates = self._compute_frame_gates(x, HW)
|
|
|
|
main_raw, kv_cache = CachedChunkCausalGDN.forward(
|
|
self,
|
|
x,
|
|
mask=mask,
|
|
HW=HW,
|
|
rotary_emb=rotary_emb,
|
|
block_mask=block_mask,
|
|
apply_output_gate=False,
|
|
chunk_size=chunk_size,
|
|
precomputed_gates=precomputed_gates,
|
|
kv_cache=kv_cache,
|
|
save_kv_cache=save_kv_cache,
|
|
**kwargs,
|
|
)
|
|
|
|
cam_contrib: torch.Tensor | int = 0
|
|
if camera_conditions is not None:
|
|
cam_raw = self._cached_cam_branch(
|
|
x,
|
|
HW,
|
|
camera_conditions,
|
|
rotary_emb,
|
|
kv_cache,
|
|
save_kv_cache,
|
|
precomputed_gates,
|
|
)
|
|
cam_contrib = self.out_proj_cam(cam_raw)
|
|
|
|
combined = main_raw + cam_contrib
|
|
combined = self._apply_output_gate(combined, x)
|
|
output = self.proj(combined.to(self.proj.weight.dtype))
|
|
return output, kv_cache
|
|
|
|
def _cached_cam_branch(
|
|
self,
|
|
x: torch.Tensor,
|
|
HW: tuple[int, int, int],
|
|
camera_conditions: torch.Tensor,
|
|
rotary_emb: torch.Tensor | None,
|
|
kv_cache: list,
|
|
save_kv_cache: bool,
|
|
precomputed_gates: tuple | None,
|
|
) -> torch.Tensor:
|
|
"""Camera branch with cached delta-rule state."""
|
|
B, N, _ = x.shape
|
|
T, H_sp, W_sp = HW
|
|
S = H_sp * W_sp
|
|
dtype_orig = x.dtype
|
|
|
|
# Fused Triton cam prep: RMSNorm + ReLU + K-scale + UCPE 4x4 + RoPE.
|
|
q_cam_trans, k_cam_trans, v_cam_trans, inflation_sq, apply_fn_o, cam_conv_cache = _cam_prep_triton(
|
|
self,
|
|
x,
|
|
HW,
|
|
camera_conditions,
|
|
rotary_emb,
|
|
kv_cache[_SLOT_CAM_AUX],
|
|
save_kv_cache,
|
|
)
|
|
if save_kv_cache:
|
|
kv_cache[_SLOT_CAM_AUX] = cam_conv_cache
|
|
|
|
if precomputed_gates is not None:
|
|
beta, decay = precomputed_gates
|
|
else:
|
|
beta, decay = self._compute_frame_gates(x, HW)
|
|
|
|
# Dynamic beta discounting (UCPE inflation factor).
|
|
inflation_sq_spatial = inflation_sq.view(B, self.cam_heads, T, S)
|
|
frame_inflation_sq = inflation_sq_spatial.mean(dim=-1)
|
|
if beta.ndim == 3:
|
|
beta = beta / frame_inflation_sq.clamp_min(1.0)
|
|
elif beta.ndim == 4:
|
|
beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0)
|
|
|
|
out, cam_S_kv_final = _cam_main_triton(
|
|
q_cam_trans,
|
|
k_cam_trans,
|
|
v_cam_trans,
|
|
beta,
|
|
decay,
|
|
kv_cache[_SLOT_CAM],
|
|
save_kv_cache,
|
|
T,
|
|
S,
|
|
)
|
|
if save_kv_cache:
|
|
kv_cache[_SLOT_CAM] = cam_S_kv_final.detach().clone()
|
|
|
|
if getattr(self, "fp32_attention", True) and dtype_orig != torch.float32:
|
|
out = out.to(dtype_orig)
|
|
out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous()
|
|
return out.reshape(B, self.cam_dim, N).permute(0, 2, 1)
|
|
|
|
|
|
@ATTENTION_BLOCKS.register_module()
|
|
class CachedSoftmaxUCPESinglePathLiteLA(_SoftmaxUCPESinglePathLiteLA):
|
|
"""Cached softmax + UCPE camera attention for streaming inference.
|
|
|
|
Main branch: :class:`CachedChunkCausalSoftmaxAttn` (concatenate cached
|
|
post-RoPE K, V; SDPA).
|
|
Camera branch: concatenate cached post-UCPE cam K, V; SDPA on the unioned
|
|
sequence. Falls back to the parent's non-cached forward when
|
|
``kv_cache`` is absent.
|
|
"""
|
|
|
|
def __init__(self, *args, conv_kernel_size: int = 0, **kwargs):
|
|
super().__init__(*args, conv_kernel_size=0, **kwargs)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
HW: tuple[int, int, int] | None = None,
|
|
rotary_emb: torch.Tensor | None = None,
|
|
block_mask: torch.Tensor | None = None,
|
|
camera_conditions: torch.Tensor | None = None,
|
|
chunk_size: int | None = None,
|
|
**kwargs: object,
|
|
) -> torch.Tensor | tuple[torch.Tensor, list]:
|
|
kv_cache = kwargs.pop("kv_cache", None)
|
|
save_kv_cache = kwargs.pop("save_kv_cache", False)
|
|
if kv_cache is None:
|
|
return super().forward(
|
|
x,
|
|
mask=mask,
|
|
HW=HW,
|
|
rotary_emb=rotary_emb,
|
|
block_mask=block_mask,
|
|
camera_conditions=camera_conditions,
|
|
chunk_size=chunk_size,
|
|
**kwargs,
|
|
)
|
|
if HW is None:
|
|
raise ValueError("HW must be provided.")
|
|
|
|
main_raw, kv_cache = CachedChunkCausalSoftmaxAttn.forward(
|
|
self,
|
|
x,
|
|
mask=mask,
|
|
HW=HW,
|
|
rotary_emb=rotary_emb,
|
|
block_mask=block_mask,
|
|
apply_output_gate=False,
|
|
chunk_size=chunk_size,
|
|
kv_cache=kv_cache,
|
|
save_kv_cache=save_kv_cache,
|
|
**kwargs,
|
|
)
|
|
|
|
cam_contrib: torch.Tensor | int = 0
|
|
if camera_conditions is not None:
|
|
cam_raw = self._cached_cam_branch_softmax(
|
|
x,
|
|
HW,
|
|
camera_conditions,
|
|
rotary_emb,
|
|
kv_cache,
|
|
save_kv_cache,
|
|
chunk_size=chunk_size,
|
|
**kwargs,
|
|
)
|
|
cam_contrib = self.out_proj_cam(cam_raw)
|
|
|
|
combined = main_raw + cam_contrib
|
|
combined = self._apply_output_gate(combined, x)
|
|
return self.proj(combined.to(x.dtype)), kv_cache
|
|
|
|
def _cached_cam_branch_softmax(
|
|
self,
|
|
x: torch.Tensor,
|
|
HW: tuple[int, int, int],
|
|
camera_conditions: torch.Tensor,
|
|
rotary_emb: torch.Tensor | None,
|
|
kv_cache: list,
|
|
save_kv_cache: bool,
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
"""Camera branch: SDPA with cached UCPE-transformed K, V."""
|
|
B, N, _ = x.shape
|
|
T, H_sp, W_sp = HW
|
|
S = H_sp * W_sp
|
|
|
|
q_cam_trans, k_cam_trans, v_cam_trans, apply_fn_o = _prepare_cam_qkv_softmax(
|
|
self, x, HW, camera_conditions, rotary_emb, **kwargs
|
|
)
|
|
|
|
# (B, H, D, N) -> (B, H, N, D) for SDPA.
|
|
q_sdpa = q_cam_trans.transpose(-1, -2)
|
|
k_sdpa = k_cam_trans.transpose(-1, -2)
|
|
v_sdpa = v_cam_trans.transpose(-1, -2)
|
|
|
|
dtype_orig = x.dtype
|
|
if q_sdpa.dtype == torch.float32:
|
|
q_sdpa, k_sdpa, v_sdpa = q_sdpa.bfloat16(), k_sdpa.bfloat16(), v_sdpa.bfloat16()
|
|
|
|
cached_cam_k = kv_cache[_SLOT_CAM]
|
|
cached_cam_v = kv_cache[_SLOT_CAM_AUX]
|
|
if save_kv_cache:
|
|
kv_cache[_SLOT_CAM] = k_sdpa.detach().clone()
|
|
kv_cache[_SLOT_CAM_AUX] = v_sdpa.detach().clone()
|
|
if cached_cam_k is not None:
|
|
k_sdpa = torch.cat([cached_cam_k.to(k_sdpa.dtype), k_sdpa], dim=2)
|
|
v_sdpa = torch.cat([cached_cam_v.to(v_sdpa.dtype), v_sdpa], dim=2)
|
|
|
|
out = _sdpa_maybe_chunk_causal(
|
|
q_sdpa,
|
|
k_sdpa,
|
|
v_sdpa,
|
|
need_chunk_mask=False,
|
|
T=T,
|
|
S=S,
|
|
chunk_size=kwargs.get("chunk_size", None),
|
|
chunk_index=kwargs.get("chunk_index", None),
|
|
chunk_split_strategy=kwargs.get("chunk_split_strategy", "uniform"),
|
|
)
|
|
|
|
out = out.transpose(-1, -2)
|
|
if out.dtype != dtype_orig:
|
|
out = out.to(dtype_orig)
|
|
out = apply_fn_o(out.transpose(-1, -2)).transpose(-1, -2).contiguous()
|
|
return out.reshape(B, self.cam_dim, N).permute(0, 2, 1)
|