532 lines
20 KiB
Python
532 lines
20 KiB
Python
"""
|
|
AggregatorStream - Streaming causal aggregator with FlashInfer KV cache.
|
|
|
|
Provides:
|
|
- Temporal causal attention
|
|
- Sliding window support
|
|
- Scale token for scale estimation frames
|
|
- Streaming inference with FlashInfer paged KV cache
|
|
"""
|
|
|
|
import logging
|
|
import torch
|
|
import torch.nn as nn
|
|
from typing import Optional, Tuple, List
|
|
|
|
from lingbot_map.layers.block import Block, FlashInferBlock, SDPABlock
|
|
from lingbot_map.layers.rope import WanRotaryPosEmbed
|
|
from lingbot_map.aggregator.base import AggregatorBase, slice_expand_and_flatten
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AggregatorStream(AggregatorBase):
|
|
"""
|
|
Streaming causal aggregator with FlashInfer paged KV cache.
|
|
|
|
Features:
|
|
- Temporal causal attention (each frame only attends to past frames)
|
|
- Sliding window support to limit attention scope
|
|
- Scale token for scale estimation frames
|
|
- Streaming inference with FlashInfer KV cache
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
# Causal-specific parameters
|
|
sliding_window_size: int = -1,
|
|
num_frame_for_scale: int = 1,
|
|
num_random_frames: int = 0,
|
|
attend_to_special_tokens: bool = False,
|
|
attend_to_scale_frames: bool = False,
|
|
enable_3d_rope: bool = False,
|
|
max_frame_num: int = 1024,
|
|
# KV cache parameters
|
|
kv_cache_sliding_window: int = 64,
|
|
kv_cache_scale_frames: int = 8,
|
|
kv_cache_cross_frame_special: bool = True,
|
|
kv_cache_include_scale_frames: bool = True,
|
|
kv_cache_camera_only: bool = False,
|
|
# Base class parameters via **kwargs
|
|
**kwargs
|
|
):
|
|
"""
|
|
Initialize AggregatorStream.
|
|
|
|
Args:
|
|
sliding_window_size: Sliding window size in blocks (-1 for full causal)
|
|
num_frame_for_scale: Number of scale estimation frames
|
|
num_random_frames: Number of random frames for long-range dependencies
|
|
attend_to_special_tokens: Enable cross-frame special token attention
|
|
attend_to_scale_frames: Include scale frames in attention
|
|
enable_3d_rope: Enable 3D RoPE for temporal dimension in KV cache
|
|
max_frame_num: Maximum number of frames for 3D RoPE
|
|
kv_cache_sliding_window: Sliding window size for KV cache eviction
|
|
kv_cache_scale_frames: Number of scale frames to keep in KV cache
|
|
kv_cache_cross_frame_special: Keep special tokens from evicted frames
|
|
kv_cache_include_scale_frames: Include scale frames in KV cache
|
|
kv_cache_camera_only: Only keep camera tokens from evicted frames
|
|
**kwargs: Base class parameters
|
|
"""
|
|
self.sliding_window_size = sliding_window_size
|
|
self.num_frame_for_scale = num_frame_for_scale
|
|
self.num_random_frames = num_random_frames
|
|
self.attend_to_special_tokens = attend_to_special_tokens
|
|
self.attend_to_scale_frames = attend_to_scale_frames
|
|
self.enable_3d_rope = enable_3d_rope
|
|
self.max_frame_num = max_frame_num
|
|
# KV cache parameters
|
|
self.kv_cache_sliding_window = kv_cache_sliding_window
|
|
self.kv_cache_scale_frames = kv_cache_scale_frames
|
|
self.kv_cache_cross_frame_special = kv_cache_cross_frame_special
|
|
self.kv_cache_include_scale_frames = kv_cache_include_scale_frames
|
|
self.kv_cache_camera_only = kv_cache_camera_only
|
|
|
|
# Pop kwargs that are passed but not needed by base class
|
|
kwargs.pop('enable_stream_inference', None)
|
|
use_flashinfer = kwargs.pop('use_flashinfer', True)
|
|
kwargs.pop('use_flexflash', None)
|
|
use_sdpa = kwargs.pop('use_sdpa', False)
|
|
|
|
# Backend selection: SDPA (no extra deps) or FlashInfer (paged KV cache)
|
|
self.use_sdpa = use_sdpa
|
|
self.use_flashinfer = not use_sdpa # FlashInfer is default unless SDPA requested
|
|
|
|
# Call parent __init__
|
|
super().__init__(**kwargs)
|
|
|
|
# Initialize KV cache
|
|
self._init_kv_cache()
|
|
|
|
# Initialize 3D RoPE if enabled
|
|
if self.enable_3d_rope:
|
|
self._init_3d_rope()
|
|
|
|
def _build_blocks(
|
|
self,
|
|
block_fn,
|
|
depth: int,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
mlp_ratio: float,
|
|
qkv_bias: bool,
|
|
proj_bias: bool,
|
|
ffn_bias: bool,
|
|
init_values: float,
|
|
qk_norm: bool,
|
|
):
|
|
"""Build frame and global blocks for streaming causal mode."""
|
|
block_params = dict(
|
|
dim=embed_dim,
|
|
num_heads=num_heads,
|
|
mlp_ratio=mlp_ratio,
|
|
qkv_bias=qkv_bias,
|
|
proj_bias=proj_bias,
|
|
ffn_bias=ffn_bias,
|
|
init_values=init_values,
|
|
qk_norm=qk_norm,
|
|
)
|
|
|
|
# Frame blocks: Standard Block + RoPE
|
|
self.frame_blocks = nn.ModuleList([
|
|
block_fn(**block_params, rope=self.rope)
|
|
for _ in range(depth)
|
|
])
|
|
|
|
# Global blocks: FlashInferBlock (default) or SDPABlock (fallback)
|
|
GlobalBlockCls = SDPABlock if self.use_sdpa else FlashInferBlock
|
|
self.global_blocks = nn.ModuleList([
|
|
GlobalBlockCls(
|
|
**block_params,
|
|
rope=self.rope if not self.disable_global_rope else None,
|
|
kv_cache_sliding_window=self.kv_cache_sliding_window,
|
|
kv_cache_scale_frames=self.kv_cache_scale_frames,
|
|
kv_cache_cross_frame_special=self.kv_cache_cross_frame_special,
|
|
kv_cache_include_scale_frames=self.kv_cache_include_scale_frames,
|
|
kv_cache_camera_only=self.kv_cache_camera_only,
|
|
)
|
|
for _ in range(depth)
|
|
])
|
|
|
|
def _setup_special_tokens(self):
|
|
"""Setup camera, register, and scale tokens for causal mode."""
|
|
# Camera token
|
|
self.camera_token = nn.Parameter(
|
|
torch.randn(1, 2, 1, self.embed_dim)
|
|
)
|
|
|
|
# Register tokens
|
|
if self.num_register_tokens > 0:
|
|
self.register_token = nn.Parameter(
|
|
torch.randn(1, 2, self.num_register_tokens, self.embed_dim)
|
|
)
|
|
|
|
# Scale token (causal mode specific)
|
|
self.scale_token = nn.Parameter(
|
|
torch.ones(1, 2, 1, self.embed_dim)
|
|
)
|
|
|
|
# Initialize
|
|
nn.init.normal_(self.camera_token, std=1e-6)
|
|
if self.num_register_tokens > 0:
|
|
nn.init.normal_(self.register_token, std=1e-6)
|
|
nn.init.normal_(self.scale_token, std=1e-6)
|
|
|
|
# Token indexing (includes scale token)
|
|
self.patch_start_idx = 1 + self.num_register_tokens + 1 # camera + register + scale
|
|
self.num_special_tokens = 1 + self.num_register_tokens + 1
|
|
|
|
def _init_kv_cache(self):
|
|
"""Initialize KV cache for streaming inference."""
|
|
self.kv_cache_manager = None # FlashInfer (lazy-initialized)
|
|
self.kv_cache = {} # Dict-based cache for SDPA
|
|
self.total_frames_processed = 0
|
|
self._cached_pos3d = None
|
|
|
|
if self.use_sdpa:
|
|
# Dict-based KV cache for SDPA
|
|
if hasattr(self, 'depth'):
|
|
for i in range(self.depth):
|
|
self.kv_cache[f"k_{i}"] = None
|
|
self.kv_cache[f"v_{i}"] = None
|
|
self.kv_cache[f"k_{i}_special"] = None
|
|
self.kv_cache[f"v_{i}_special"] = None
|
|
logger.info(f"SDPA KV cache initialized with {self.depth} blocks")
|
|
else:
|
|
logger.info("FlashInfer KV cache will be lazily initialized on first forward")
|
|
|
|
def _get_flashinfer_manager(self, device, dtype, tokens_per_frame=None):
|
|
"""Lazily initialize FlashInferKVCacheManager on first use.
|
|
|
|
Args:
|
|
device: Device for cache tensors.
|
|
dtype: Data type for cache tensors.
|
|
tokens_per_frame: Actual number of tokens per frame (patches + specials).
|
|
If None, falls back to assuming square images of self.img_size.
|
|
"""
|
|
if self.kv_cache_manager is None:
|
|
from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager
|
|
num_heads = self.embed_dim // 64 # head_dim = 64 for ViT-L
|
|
head_dim = 64
|
|
if tokens_per_frame is None:
|
|
tokens_per_frame = (self.img_size // self.patch_size) ** 2 + self.num_special_tokens
|
|
# max_num_frames: scale + window + headroom
|
|
max_num_frames = self.kv_cache_scale_frames + self.kv_cache_sliding_window + 16
|
|
self.kv_cache_manager = FlashInferKVCacheManager(
|
|
num_blocks=self.depth,
|
|
max_num_frames=max_num_frames,
|
|
tokens_per_frame=tokens_per_frame,
|
|
num_heads=num_heads,
|
|
head_dim=head_dim,
|
|
dtype=dtype,
|
|
device=device,
|
|
num_special_tokens=self.num_special_tokens,
|
|
scale_frames=self.kv_cache_scale_frames,
|
|
sliding_window=self.kv_cache_sliding_window,
|
|
max_total_frames=self.max_frame_num + 100,
|
|
force_fp32=getattr(self, 'kv_cache_force_fp32', False),
|
|
fa3=getattr(self, 'kv_cache_fa3', False),
|
|
)
|
|
logger.info(
|
|
f"FlashInfer KV cache manager initialized: {self.depth} blocks, "
|
|
f"max_frames={max_num_frames}, tokens_per_frame={tokens_per_frame}"
|
|
)
|
|
return self.kv_cache_manager
|
|
|
|
def clean_kv_cache(self):
|
|
"""Clean KV cache (call this when starting a new sequence)."""
|
|
if self.kv_cache_manager is not None:
|
|
self.kv_cache_manager.reset()
|
|
if self.kv_cache:
|
|
for key in list(self.kv_cache.keys()):
|
|
if key == "_skip_append":
|
|
self.kv_cache[key] = False
|
|
else:
|
|
self.kv_cache[key] = None
|
|
self.total_frames_processed = 0
|
|
self._cached_pos3d = None
|
|
logger.info("KV cache cleaned")
|
|
|
|
def _init_3d_rope(self):
|
|
"""Initialize 3D RoPE for streaming inference."""
|
|
if not self.enable_3d_rope:
|
|
self.rope3d = None
|
|
return
|
|
|
|
num_heads = 16
|
|
head_dim = self.embed_dim // num_heads
|
|
|
|
self.rope3d = WanRotaryPosEmbed(
|
|
attention_head_dim=head_dim,
|
|
patch_size=(1, self.patch_size, self.patch_size),
|
|
max_seq_len=self.max_frame_num,
|
|
)
|
|
logger.info(f"3D RoPE initialized for max {self.max_frame_num} frames, head_dim={head_dim}")
|
|
|
|
def _get_3d_positions_streaming(self, num_frames, H, W, device, f_start, f_end):
|
|
"""
|
|
Generate 3D RoPE positions for streaming mode with correct global frame indices.
|
|
|
|
Args:
|
|
num_frames: Number of frames in current batch
|
|
H, W: Image height and width
|
|
device: Device to create positions on
|
|
f_start: Global start frame index
|
|
f_end: Global end frame index
|
|
|
|
Returns:
|
|
pos3d: [1, 1, num_frames * P, head_dim//2] complex tensor
|
|
"""
|
|
if self.rope3d is None:
|
|
return None
|
|
|
|
pph = H // self.patch_size
|
|
ppw = W // self.patch_size
|
|
|
|
pos3d = self.rope3d(
|
|
ppf=num_frames,
|
|
pph=pph,
|
|
ppw=ppw,
|
|
patch_start_idx=self.num_special_tokens,
|
|
device=device,
|
|
f_start=f_start,
|
|
f_end=f_end
|
|
)
|
|
return pos3d
|
|
|
|
def _prepare_special_tokens(
|
|
self,
|
|
B: int,
|
|
S_local: int,
|
|
S_global: int,
|
|
C: int,
|
|
num_frame_for_scale: Optional[int] = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Prepare camera, register, and scale tokens.
|
|
|
|
Args:
|
|
B: Batch size
|
|
S_local: Local sequence length
|
|
S_global: Global sequence length
|
|
C: Embedding dimension
|
|
num_frame_for_scale: Number of frames for scale estimation
|
|
|
|
Returns:
|
|
Special tokens [B*S_global, N_special, C]
|
|
"""
|
|
# Get effective num_frame_for_scale
|
|
scale_frames = self.num_frame_for_scale if num_frame_for_scale is None else num_frame_for_scale
|
|
|
|
# Check cache state for both backends
|
|
has_flashinfer_cache = self.kv_cache_manager is not None and self.kv_cache_manager.num_frames > 0
|
|
has_sdpa_cache = self.kv_cache is not None and self.kv_cache.get("k_0") is not None
|
|
|
|
# Determine if we're in causal inference mode based on KV cache state
|
|
causal_inference = True
|
|
|
|
if causal_inference and has_flashinfer_cache:
|
|
S_cached = self.kv_cache_manager.num_frames
|
|
S_true = S_cached + S_global
|
|
elif causal_inference and has_sdpa_cache:
|
|
_, _, S_cached, _, _ = self.kv_cache["k_0"].shape
|
|
S_true = S_cached + S_global
|
|
else:
|
|
S_true = S_global
|
|
|
|
# Expand tokens based on mode
|
|
if causal_inference and S_true > S_global:
|
|
# Streaming mode: expand with S_true, then slice to get current frames
|
|
effective_scale_frames = min(scale_frames, S_true)
|
|
|
|
camera_token_full = slice_expand_and_flatten(self.camera_token, B, S_true)
|
|
camera_token = camera_token_full[-S_global:, :, :]
|
|
|
|
register_token_full = slice_expand_and_flatten(self.register_token, B, S_true)
|
|
register_token = register_token_full[-S_global:, :, :]
|
|
scale_token_full = slice_expand_and_flatten(
|
|
self.scale_token, B, S_true, first_num_frame=effective_scale_frames
|
|
)
|
|
scale_token = scale_token_full[-S_global:, :, :]
|
|
else:
|
|
# Batch mode or first inference: expand directly
|
|
effective_scale_frames = min(scale_frames, S_global)
|
|
|
|
camera_token = slice_expand_and_flatten(self.camera_token, B, S_global)
|
|
register_token = slice_expand_and_flatten(self.register_token, B, S_global)
|
|
scale_token = slice_expand_and_flatten(
|
|
self.scale_token, B, S_global, first_num_frame=effective_scale_frames
|
|
)
|
|
|
|
special_tokens = torch.cat([camera_token, register_token, scale_token], dim=1)
|
|
|
|
# Verify shape
|
|
expected_shape = (B * S_global, self.num_special_tokens, C)
|
|
assert special_tokens.shape == expected_shape, \
|
|
f"Expected {expected_shape}, got {special_tokens.shape}"
|
|
|
|
return special_tokens
|
|
|
|
def _process_global_attention(
|
|
self,
|
|
tokens: torch.Tensor,
|
|
B: int,
|
|
S_local: int,
|
|
S_global: int,
|
|
P: int,
|
|
C: int,
|
|
global_idx: int,
|
|
pos: Optional[torch.Tensor] = None,
|
|
# Mode-specific parameters
|
|
num_frame_for_scale: Optional[int] = None,
|
|
sliding_window_size: Optional[int] = None,
|
|
num_frame_per_block: int = 1,
|
|
**kwargs,
|
|
) -> Tuple[torch.Tensor, int, List[torch.Tensor]]:
|
|
"""
|
|
Process causal global attention via FlashInfer streaming path.
|
|
|
|
Args:
|
|
tokens: Input tokens
|
|
B: Batch size
|
|
S_local: Local sequence length
|
|
S_global: Global sequence length
|
|
P: Tokens per frame
|
|
C: Embedding dimension
|
|
global_idx: Current global block index
|
|
pos: Position embeddings
|
|
num_frame_for_scale: Number of frames for scale estimation
|
|
sliding_window_size: Sliding window size in blocks
|
|
num_frame_per_block: Number of frames per processing block
|
|
|
|
Returns:
|
|
(tokens, global_idx, intermediates)
|
|
"""
|
|
# Extract image dimensions from kwargs for 3D RoPE
|
|
image_height = kwargs.get('image_height', self.img_size)
|
|
image_width = kwargs.get('image_width', self.img_size)
|
|
|
|
return self._process_causal_stream(
|
|
tokens, B, S_local, S_global, P, C, global_idx, pos,
|
|
num_frame_per_block, sliding_window_size, num_frame_for_scale,
|
|
image_height=image_height, image_width=image_width
|
|
)
|
|
|
|
def _process_causal_stream(
|
|
self,
|
|
tokens: torch.Tensor,
|
|
B: int,
|
|
S_local: int,
|
|
S_global: int,
|
|
P: int,
|
|
C: int,
|
|
global_idx: int,
|
|
pos: Optional[torch.Tensor] = None,
|
|
num_frame_per_block: int = 1,
|
|
sliding_window_size: Optional[int] = None,
|
|
num_frame_for_scale: Optional[int] = None,
|
|
image_height: Optional[int] = None,
|
|
image_width: Optional[int] = None,
|
|
):
|
|
"""
|
|
Causal attention for streaming inference using FlashInfer KV cache.
|
|
|
|
Args:
|
|
tokens: Input tokens [B*S_local, P, C]
|
|
B: Batch size
|
|
S_local: Local sequence length
|
|
S_global: Global sequence length
|
|
P: Number of patches per frame (includes special tokens)
|
|
C: Channel dimension
|
|
global_idx: Starting block index
|
|
pos: Position embeddings [B*S_global, P, 2]
|
|
num_frame_per_block: Number of frames per block
|
|
sliding_window_size: Sliding window size in blocks
|
|
num_frame_for_scale: Number of scale frames
|
|
image_height: Image height for 3D RoPE calculation
|
|
image_width: Image width for 3D RoPE calculation
|
|
|
|
Returns:
|
|
(tokens, global_idx, intermediates): Updated tokens, next block index, intermediate outputs
|
|
"""
|
|
# Get effective parameters
|
|
scale_frames = num_frame_for_scale if num_frame_for_scale is not None else self.num_frame_for_scale
|
|
|
|
# Reshape tokens: [B*S_local, P, C] -> [B, S_local*P, C]
|
|
if tokens.shape != (B, S_local * P, C):
|
|
tokens = tokens.view(B, S_local, P, C).view(B, S_local * P, C)
|
|
|
|
# Calculate number of frames for block mask
|
|
num_frames = S_global
|
|
num_patches = P - self.num_special_tokens
|
|
|
|
# Check if this is the first block group
|
|
is_first_block_group = (global_idx < self.aa_block_size)
|
|
|
|
if self.enable_3d_rope and self.rope3d is not None:
|
|
if is_first_block_group:
|
|
f_start = self.total_frames_processed
|
|
f_end = self.total_frames_processed + S_global
|
|
|
|
H = image_height if image_height is not None else self.img_size
|
|
W = image_width if image_width is not None else self.img_size
|
|
pos3d = self._get_3d_positions_streaming(
|
|
S_global, H, W, tokens.device, f_start, f_end
|
|
)
|
|
self._cached_pos3d = pos3d
|
|
else:
|
|
pos3d = self._cached_pos3d
|
|
pos = pos3d
|
|
else:
|
|
# Reshape pos: [B*S_global, P, 2] -> [B, S_global*P, 2]
|
|
if pos is not None and pos.shape != (B, S_global * P, 2):
|
|
pos = pos.view(B, S_global, P, 2).view(B, S_global * P, 2)
|
|
|
|
intermediates = []
|
|
|
|
# Process blocks with KV cache
|
|
for _ in range(self.aa_block_size):
|
|
num_patches = P - self.num_special_tokens
|
|
if self.use_sdpa:
|
|
# SDPA: dict-based KV cache
|
|
tokens = self.global_blocks[global_idx](
|
|
tokens,
|
|
pos=pos,
|
|
enable_ulysses_cp=False,
|
|
num_patches=num_patches,
|
|
num_special=self.num_special_tokens,
|
|
num_frames=num_frames,
|
|
enable_3d_rope=self.enable_3d_rope,
|
|
kv_cache=self.kv_cache,
|
|
global_idx=global_idx,
|
|
num_frame_per_block=num_frame_per_block,
|
|
num_frame_for_scale=scale_frames,
|
|
num_register_tokens=self.num_register_tokens,
|
|
)
|
|
else:
|
|
# FlashInfer: paged KV cache manager
|
|
manager = self._get_flashinfer_manager(tokens.device, tokens.dtype, tokens_per_frame=P)
|
|
tokens = self.global_blocks[global_idx](
|
|
tokens,
|
|
pos=pos,
|
|
enable_ulysses_cp=False,
|
|
num_patches=num_patches,
|
|
num_special=self.num_special_tokens,
|
|
num_frames=num_frames,
|
|
enable_3d_rope=self.enable_3d_rope,
|
|
kv_cache=manager,
|
|
global_idx=global_idx,
|
|
num_frame_per_block=num_frame_per_block,
|
|
num_frame_for_scale=scale_frames,
|
|
num_register_tokens=self.num_register_tokens,
|
|
)
|
|
|
|
global_idx += 1
|
|
intermediates.append(tokens.view(B, S_local, P, C))
|
|
|
|
# Update total frames processed counter only on the first block group
|
|
if is_first_block_group and not (isinstance(self.kv_cache, dict) and self.kv_cache.get("_skip_append", False)):
|
|
self.total_frames_processed += S_global
|
|
|
|
return tokens, global_idx, intermediates
|