Spaces:
Running on Zero
Running on Zero
| """ | |
| Bidirectional Transformer decoder inspired by CUT3R. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from flash_attn import flash_attn_func | |
| from typing import Optional, Tuple, List | |
| import math | |
| import os | |
| import sys | |
| import logging | |
| from pathlib import Path | |
| logger = logging.getLogger(__name__) | |
| from src.models.physformer_rope import PhysFormerRoPE, apply_rotary_emb | |
| class RMSNorm(nn.Module): | |
| """ | |
| Root Mean Square Layer Normalization (RMSNorm). | |
| More efficient than LayerNorm, used in modern transformers like LLaMA. | |
| Reference: Zhang and Sennrich 2019 - "Root Mean Square Layer Normalization" | |
| """ | |
| def __init__(self, dim: int, eps: float = 1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| # RMS normalization: x / RMS(x) * weight | |
| # RMS(x) = sqrt(mean(x^2) + eps) | |
| rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps) | |
| x_normed = x / rms | |
| return self.weight * x_normed | |
| def get_norm_layer(norm_type: str, dim: int) -> nn.Module: | |
| """ | |
| Factory function to create normalization layer. | |
| Args: | |
| norm_type: 'layer_norm' or 'rms_norm' | |
| dim: Dimension to normalize | |
| Returns: | |
| Normalization module | |
| """ | |
| if norm_type == 'rms_norm': | |
| return RMSNorm(dim) | |
| elif norm_type == 'layer_norm': | |
| return nn.LayerNorm(dim) | |
| else: | |
| raise ValueError(f"Unknown norm_type: {norm_type}. Choose 'layer_norm' or 'rms_norm'.") | |
| class BidirectionalDecoder(nn.Module): | |
| """ | |
| Bidirectional decoder following CUT3R's architecture. | |
| - Returns all layer outputs for multi-scale fusion | |
| - Bidirectional cross-attention between state and object tokens | |
| - RMSNorm and QK-normalization for training stability (configurable) | |
| - Optional elementwise gated attention (Qwen3 style) | |
| """ | |
| def __init__( | |
| self, | |
| state_dim: int = 768, | |
| num_layers: int = 12, | |
| num_heads: int = 12, | |
| feedforward_dim: int = None, # None: 2.5x state_dim (SwiGLU) | |
| dropout: float = 0.1, | |
| activation: str = 'swiglu', | |
| return_all_layers: bool = True, # For multi-scale fusion | |
| use_learnable_pe: bool = False, | |
| use_rope: bool = True, | |
| max_state_tokens: int = 1024, | |
| max_obs_tokens: int = 256, | |
| use_self_attention: bool = False, # Optional state self-attention | |
| norm_type: str = 'rms_norm', # 'layer_norm' or 'rms_norm' | |
| qk_norm: bool = True, # QK-normalization for training stability | |
| num_register_tokens: int = 16, # Learnable register tokens (RenderFormer) | |
| center_rope_positions: bool = False, # Subtract the valid-object anchor centroid before RoPE (registers sit at the centroid) | |
| rope_set_aggregation: bool = False, # Set aggregation for multi-point RoPE | |
| rope_dim: int = None, # Custom RoPE dimension (None = auto-compute as head_dim//6*6) | |
| dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE | |
| dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE | |
| layer_indices_to_save: set = None, # Save only these layer outputs to reduce memory | |
| elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) | |
| state_update_enabled: bool = True, # False: state is fixed learnable embeddings | |
| objects_self_attention: bool = False, # Objects self-attention for object-object interaction | |
| use_se3: bool = False, # SE(3) pairwise bias for object self-attention | |
| # Block Attention Residuals (Kimi 2026, arXiv:2603.15031) | |
| block_attn_res: bool = False, # Enable Block AttnRes for inter-layer residuals | |
| block_attn_res_block_size: int = 4, # Layers per block (S in the paper) | |
| block_attn_res_final_aggregate: bool = False, # Return the final depth aggregation as the last layer's features | |
| film_conditioning: bool = False, # Per-layer FiLM conditioning on step size | |
| ): | |
| super().__init__() | |
| self.state_dim = state_dim | |
| self.num_layers = num_layers | |
| self.return_all_layers = return_all_layers | |
| # Optionally save only selected layer outputs to reduce memory | |
| if layer_indices_to_save is not None: | |
| for idx in layer_indices_to_save: | |
| assert 0 <= idx <= num_layers, \ | |
| f"Invalid layer index {idx} to save. Valid range: [0, {num_layers}]" | |
| self.layer_indices_to_save = layer_indices_to_save | |
| else: | |
| # Default: save all layers if return_all_layers is True | |
| self.layer_indices_to_save = set(range(num_layers + 1)) if return_all_layers else set() | |
| self.use_rope = bool(use_rope) | |
| self.rope_set_aggregation = rope_set_aggregation | |
| # If RoPE is enabled, disable learnable positional embeddings in decoder inputs | |
| self.use_learnable_pe = bool(use_learnable_pe and (not self.use_rope)) | |
| assert self.use_rope, "RoPE must be enabled for decoder" | |
| if feedforward_dim is None: | |
| feedforward_dim = int(2.5 * state_dim) | |
| # Register tokens (RenderFormer): learnable tokens that provide scratch | |
| # space for global information and reduce embedding noise. | |
| self.num_register_tokens = num_register_tokens | |
| self.center_rope_positions = bool(center_rope_positions) | |
| if num_register_tokens > 0: | |
| self.register_tokens = nn.Parameter(torch.randn(1, num_register_tokens, state_dim) * 0.02) | |
| # Register positions are created in forward() to match the dimensionality | |
| # of obj_anchor_pos (3D single-point or Nq*3 D multi-point RoPE). | |
| else: | |
| self.register_tokens = None | |
| # Transformer decoder layers | |
| self.state_update_enabled = state_update_enabled | |
| self.objects_self_attention = objects_self_attention | |
| if not state_update_enabled: | |
| print(f"[Decoder] State update disabled: state acts as learnable embeddings (objects attend to state only)") | |
| if objects_self_attention: | |
| print(f"[Decoder] Objects self-attention enabled: direct object-to-object interaction") | |
| self.layers = nn.ModuleList([ | |
| TransformerDecoderLayer( | |
| state_dim, | |
| num_heads, | |
| feedforward_dim, | |
| dropout, | |
| activation, | |
| use_self_attention=use_self_attention, | |
| use_rope=self.use_rope, | |
| norm_type=norm_type, | |
| qk_norm=qk_norm, | |
| rope_set_aggregation=self.rope_set_aggregation, | |
| rope_dim=rope_dim, | |
| dim_per_coord_3d=dim_per_coord_3d, | |
| dim_per_coord_multi=dim_per_coord_multi, | |
| elementwise_gate=elementwise_gate, | |
| state_update_enabled=state_update_enabled, | |
| objects_self_attention=objects_self_attention, | |
| use_se3=use_se3, | |
| film_conditioning=film_conditioning | |
| ) for _ in range(num_layers) | |
| ]) | |
| # Final layer normalization | |
| self.state_norm = get_norm_layer(norm_type, state_dim) | |
| self.obs_norm = get_norm_layer(norm_type, state_dim) | |
| # Block Attention Residuals (Kimi 2026, arXiv:2603.15031) | |
| # Replaces standard inter-layer residuals with softmax attention over block representations | |
| # h_l = sum_i alpha_{i->l} * v_i (attention over depth) instead of h_l = h_{l-1} + f_l(h_{l-1}) | |
| self.block_attn_res = block_attn_res | |
| self.block_attn_res_block_size = block_attn_res_block_size | |
| # With return_all_layers the last returned features are the last layer's own output; | |
| # block_attn_res_final_aggregate returns the final depth aggregation instead (off by default). | |
| self.block_attn_res_final_aggregate = bool(block_attn_res_final_aggregate) | |
| if self.block_attn_res_final_aggregate and not block_attn_res: | |
| raise ValueError('block_attn_res_final_aggregate requires block_attn_res') | |
| if self.block_attn_res_final_aggregate: | |
| print("[Decoder] Block AttnRes final aggregation returned as the last layer's features") | |
| if block_attn_res: | |
| N_blocks = max(1, num_layers // block_attn_res_block_size) | |
| print(f"[Decoder] Block AttnRes enabled: {num_layers} layers, block_size={block_attn_res_block_size}, N_blocks={N_blocks}") | |
| # Per-layer pseudo-query w_l (d-dimensional learnable vector) for depth attention | |
| self.attn_res_queries = nn.ParameterList([ | |
| nn.Parameter(torch.zeros(state_dim)) for _ in range(num_layers) | |
| ]) | |
| # Shared RMSNorm for keys in depth attention (prevents magnitude bias) | |
| self.attn_res_key_norm = RMSNorm(state_dim) | |
| def forward( | |
| self, | |
| state: Optional[torch.Tensor], | |
| obj_features: torch.Tensor, | |
| state_pos: Optional[torch.Tensor] = None, | |
| obj_anchor_pos: Optional[torch.Tensor] = None, | |
| mask: Optional[torch.Tensor] = None, | |
| se3_bias: Optional[torch.Tensor] = None, # (B, H, Nobj, Nobj) pairwise SE(3) bias | |
| time_features: Optional[torch.Tensor] = None # (B, 2) [s, s^2] for per-layer FiLM | |
| ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: | |
| """ | |
| Process state and per-object features through bidirectional interaction. | |
| Register tokens (RenderFormer) are prepended to object features. | |
| When state=None (no_global_state mode): | |
| - Skips all state cross-attention | |
| - Runs objects self-attention + objects FFN only | |
| - Returns dummy state entries | |
| Args: | |
| state: State tokens (B, N_state, D) or None for no-state mode | |
| obj_features: Per-object encoder features (B, N_obj, D) | |
| state_pos: State positions for RoPE (B, N_state, 3) or None | |
| obj_anchor_pos: Object anchor positions for multi-point RoPE (B, N_obj, Nq*3) | |
| mask: Optional attention mask | |
| time_features: (B, 2) [s, s^2] for per-layer FiLM conditioning | |
| Returns: | |
| If return_all_layers: (all_state_layers, all_obj_layers) | |
| Else: (final_state, final_obj_features) | |
| Register tokens are stripped from obj_features before returning. | |
| """ | |
| no_state = state is None | |
| # Store selected layer outputs for memory efficiency | |
| all_states = [] if self.return_all_layers else None | |
| all_obj_features = [] if self.return_all_layers else None | |
| # Rope-only: expect explicit obj_anchor_pos; state_pos only needed when state is used | |
| if obj_anchor_pos is None: | |
| raise ValueError("decoder.forward expects obj_anchor_pos for rope-only mode") | |
| if not no_state and state_pos is None: | |
| raise ValueError("decoder.forward expects state_pos when state is provided") | |
| if self.center_rope_positions: | |
| # Rotary angles of register tokens are zero, so register<->object attention would otherwise | |
| # depend on absolute object positions. Expressing every position relative to the centroid of | |
| # the valid objects' anchors makes the whole decoder exactly translation invariant. | |
| Bc, Nc, Pc = obj_anchor_pos.shape | |
| pos = obj_anchor_pos.view(Bc, Nc, Pc // 3, 3) | |
| if mask is not None: | |
| valid = (mask.reshape(Bc, -1)[:, -Nc:] == 0).to(pos.dtype) # (B, Nobj): additive mask 0 = valid | |
| else: | |
| valid = torch.ones(Bc, Nc, dtype=pos.dtype, device=pos.device) | |
| w = valid.view(Bc, Nc, 1, 1) | |
| centroid = (pos * w).sum(dim=(1, 2)) / (w.sum(dim=(1, 2)) * pos.shape[2]).clamp_min(1.0) # (B, 3) | |
| obj_anchor_pos = (pos - centroid.view(Bc, 1, 1, 3)).reshape(Bc, Nc, Pc) | |
| if not getattr(self, '_center_logged', False): | |
| self._center_logged = True | |
| print(f"[Decoder] RoPE positions centred at the valid-object anchor centroid ACTIVE " | |
| f"(|centroid| max={centroid.abs().max().item():.2f} m; registers at the centroid)", flush=True) | |
| # Prepend register tokens to object features (RenderFormer) | |
| if self.num_register_tokens > 0: | |
| B = obj_features.shape[0] | |
| register_tokens_expanded = self.register_tokens.expand(B, -1, -1) # (B, N_reg, D) | |
| obj_features_with_reg = torch.cat([register_tokens_expanded, obj_features], dim=1) | |
| # Register positions match the dimensionality of obj_anchor_pos (3 or Nq*3) | |
| pos_dim = obj_anchor_pos.shape[-1] # Either 3 or Nq*3 | |
| register_positions = torch.zeros(B, self.num_register_tokens, pos_dim, | |
| device=obj_anchor_pos.device, dtype=obj_anchor_pos.dtype) | |
| obj_anchor_pos_with_reg = torch.cat([register_positions, obj_anchor_pos], dim=1) | |
| # Extend mask to cover register tokens (always valid) | |
| if mask is not None: | |
| # mask: (B, 1, 1, Nobj) SDPA additive -> (B, 1, 1, Nreg+Nobj); | |
| # register entries are 0.0 (attend normally) | |
| register_mask = torch.zeros(B, 1, 1, self.num_register_tokens, | |
| device=mask.device, dtype=mask.dtype) | |
| mask_with_reg = torch.cat([register_mask, mask], dim=-1) | |
| else: | |
| mask_with_reg = None | |
| # Extend SE(3) bias to include register tokens (zero bias = attend freely) | |
| if se3_bias is not None and self.num_register_tokens > 0: | |
| N_reg = self.num_register_tokens | |
| N_obj = se3_bias.shape[-1] # original Nobj | |
| H = se3_bias.shape[1] | |
| # Pad: (B, H, Nobj, Nobj) -> (B, H, Nreg+Nobj, Nreg+Nobj) | |
| # Register tokens have zero bias (vanilla attention) | |
| se3_bias = F.pad(se3_bias, (N_reg, 0, N_reg, 0), value=0.0) | |
| else: | |
| obj_features_with_reg = obj_features | |
| obj_anchor_pos_with_reg = obj_anchor_pos | |
| mask_with_reg = mask | |
| # Process through each layer with register tokens | |
| state_current = state | |
| obj_current = obj_features_with_reg | |
| # Create a dummy state placeholder for no-state mode (used in returned lists) | |
| if no_state: | |
| B = obj_features.shape[0] | |
| _dummy_state = torch.zeros(B, 1, self.state_dim, device=obj_features.device, dtype=obj_features.dtype) | |
| # Add initial layer (layer 0) if needed | |
| if self.return_all_layers and 0 in self.layer_indices_to_save: | |
| all_states.append(_dummy_state if no_state else state) | |
| all_obj_features.append(obj_features) # Store original without register tokens | |
| # Block Attention Residuals path | |
| if self.block_attn_res: | |
| # Block AttnRes (Kimi 2026): replace inter-layer residuals with | |
| # softmax attention over completed block sums + partial block sum | |
| S = self.block_attn_res_block_size | |
| completed_blocks = [obj_current] # b_0 = initial obj features (embedding) | |
| partial_block = None # Running intra-block sum | |
| layer_in_block = 0 | |
| for layer_idx, layer in enumerate(self.layers): | |
| # Block boundary: commit partial sum | |
| if layer_in_block == S: | |
| completed_blocks.append(partial_block) | |
| partial_block = None | |
| layer_in_block = 0 | |
| # Compute input via depth attention over block representations | |
| obj_current = self._block_attn_res_aggregate( | |
| completed_blocks, partial_block, self.attn_res_queries[layer_idx] | |
| ) | |
| # Run layer (layer uses internal residuals for sublayers) | |
| state_current, obj_output = layer( | |
| state_current, obj_current, | |
| state_pos, obj_anchor_pos_with_reg, | |
| mask_with_reg, | |
| se3_bias=se3_bias, | |
| time_features=time_features | |
| ) | |
| # Compute layer delta: f_l(h_l) = output - input | |
| # The layer output includes internal residuals, so | |
| # delta = obj_output - obj_current is the "function output" for AttnRes | |
| layer_delta = obj_output - obj_current | |
| # Accumulate delta into partial block sum | |
| partial_block = layer_delta if partial_block is None else partial_block + layer_delta | |
| layer_in_block += 1 | |
| # Save layer outputs if needed | |
| if self.return_all_layers and (layer_idx + 1) in self.layer_indices_to_save: | |
| # For output, use the full obj_output (input + delta) | |
| if self.num_register_tokens > 0: | |
| obj_out_no_reg = obj_output[:, self.num_register_tokens:, :] | |
| else: | |
| obj_out_no_reg = obj_output | |
| all_states.append(_dummy_state if no_state else state_current) | |
| all_obj_features.append(obj_out_no_reg) | |
| # Final aggregation: attend over all blocks (including last partial) | |
| if partial_block is not None: | |
| completed_blocks.append(partial_block) | |
| # Reuse the last layer's query for the final aggregation | |
| obj_current = self._block_attn_res_aggregate( | |
| completed_blocks, None, self.attn_res_queries[-1] | |
| ) | |
| if self.return_all_layers and self.block_attn_res_final_aggregate: | |
| if self.num_layers not in self.layer_indices_to_save: | |
| raise ValueError('block_attn_res_final_aggregate needs the last layer among the saved layers') | |
| all_obj_features[-1] = (obj_current[:, self.num_register_tokens:, :] if self.num_register_tokens > 0 | |
| else obj_current) | |
| else: | |
| # Standard residual path | |
| for layer_idx, layer in enumerate(self.layers): | |
| state_current, obj_current = layer( | |
| state_current, obj_current, | |
| state_pos, obj_anchor_pos_with_reg, | |
| mask_with_reg, | |
| se3_bias=se3_bias, | |
| time_features=time_features | |
| ) | |
| if self.return_all_layers and (layer_idx + 1) in self.layer_indices_to_save: | |
| if self.num_register_tokens > 0: | |
| obj_current_no_reg = obj_current[:, self.num_register_tokens:, :] | |
| else: | |
| obj_current_no_reg = obj_current | |
| all_states.append(_dummy_state if no_state else state_current) | |
| all_obj_features.append(obj_current_no_reg) | |
| # Apply final normalization | |
| if self.return_all_layers: | |
| if not no_state: | |
| all_states[-1] = self.state_norm(all_states[-1]) | |
| all_obj_features[-1] = self.obs_norm(all_obj_features[-1]) | |
| return all_states, all_obj_features | |
| else: | |
| if self.num_register_tokens > 0: | |
| obj_current_no_reg = obj_current[:, self.num_register_tokens:, :] | |
| else: | |
| obj_current_no_reg = obj_current | |
| state_final = self.state_norm(state_current) if not no_state else _dummy_state | |
| obj_final = self.obs_norm(obj_current_no_reg) | |
| return state_final, obj_final | |
| def _block_attn_res_aggregate( | |
| self, | |
| completed_blocks: list, | |
| partial_block: Optional[torch.Tensor], | |
| w_l: nn.Parameter, | |
| ) -> torch.Tensor: | |
| """ | |
| Block AttnRes aggregation (Eq. 6 from Kimi 2026). | |
| Computes softmax attention over completed block representations + optional | |
| partial block sum, using per-layer pseudo-query w_l. | |
| phi(q, k) = exp(q^T RMSNorm(k)) | |
| alpha_{i->l} = phi(w_l, k_i) / sum_j phi(w_l, k_j) | |
| h_l = sum_i alpha_{i->l} * v_i | |
| Args: | |
| completed_blocks: List of block representations [b_0, b_1, ...] | |
| partial_block: Current intra-block partial sum (None if at block boundary) | |
| w_l: Learnable pseudo-query vector (d,) | |
| Returns: | |
| Aggregated input for next layer (B, N_tokens, D) | |
| """ | |
| if partial_block is not None: | |
| sources = completed_blocks + [partial_block] | |
| else: | |
| sources = completed_blocks | |
| # Stack: (num_sources, B, N_tokens, D) | |
| V = torch.stack(sources, dim=0) | |
| # RMSNorm on keys prevents magnitude bias between complete blocks and partial sums | |
| K = self.attn_res_key_norm(V) | |
| # Compute attention logits: w_l^T RMSNorm(k_i) for each source | |
| # w_l: (D,), K: (num_sources, B, N_tokens, D) -> logits: (num_sources, B, N_tokens) | |
| logits = torch.einsum('d, n b t d -> n b t', w_l, K) | |
| alpha = logits.softmax(dim=0) | |
| # Weighted sum: (num_sources, B, N_tokens) x (num_sources, B, N_tokens, D) -> (B, N_tokens, D) | |
| return torch.einsum('n b t, n b t d -> b t d', alpha, V) | |
| class TransformerDecoderLayer(nn.Module): | |
| """ | |
| Single decoder layer of bidirectional cross-attention, following CUT3R and | |
| RenderFormer (RMSNorm, QK-normalization, optional gated attention). | |
| When state_update_enabled=False: | |
| - State is treated as fixed learnable embeddings (never updated) | |
| - Only the Objects <- State cross-attention runs | |
| - State self-attention and State <- Objects are skipped, saving roughly | |
| half the computation | |
| When objects_self_attention=True: | |
| - Objects self-attention runs before Objects <- State, enabling direct | |
| object-to-object interaction | |
| """ | |
| def __init__( | |
| self, | |
| d_model: int, | |
| num_heads: int, | |
| feedforward_dim: int, | |
| dropout: float = 0.1, | |
| activation: str = 'swiglu', | |
| use_self_attention: bool = False, | |
| drop_path_prob: float = 0.1, | |
| use_rope: bool = True, | |
| norm_type: str = 'rms_norm', # 'layer_norm' or 'rms_norm' | |
| qk_norm: bool = True, | |
| rope_set_aggregation: bool = False, # Set aggregation for multi-point RoPE | |
| rope_dim: int = None, # Custom RoPE dimension | |
| dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE | |
| dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE | |
| elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) | |
| state_update_enabled: bool = True, # False: state is fixed learnable embeddings | |
| objects_self_attention: bool = False, # Objects self-attention for object-object interaction | |
| use_se3: bool = False, # SE(3) pairwise bias replaces RoPE in object self-attention | |
| film_conditioning: bool = False # Per-layer FiLM conditioning on step size | |
| ): | |
| super().__init__() | |
| self.use_self_attention = use_self_attention | |
| self.use_rope = bool(use_rope) | |
| self.state_update_enabled = state_update_enabled | |
| self.objects_self_attention = objects_self_attention | |
| self._objects_self_attention_logged = False # One-time log flag | |
| # Cross-attention with optional RoPE and QK-normalization | |
| self.state_cross_attn = CrossAttentionWithRoPE( | |
| d_model=d_model, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| use_rope=self.use_rope, | |
| qk_norm=qk_norm, | |
| rope_set_aggregation=rope_set_aggregation, | |
| rope_dim=rope_dim, | |
| dim_per_coord_3d=dim_per_coord_3d, | |
| dim_per_coord_multi=dim_per_coord_multi, | |
| elementwise_gate=elementwise_gate | |
| ) | |
| self.obs_cross_attn = CrossAttentionWithRoPE( | |
| d_model=d_model, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| use_rope=self.use_rope, | |
| qk_norm=qk_norm, | |
| rope_set_aggregation=rope_set_aggregation, | |
| rope_dim=rope_dim, | |
| dim_per_coord_3d=dim_per_coord_3d, | |
| dim_per_coord_multi=dim_per_coord_multi, | |
| elementwise_gate=elementwise_gate | |
| ) | |
| # Optional state self-attention (disabled by default) | |
| if use_self_attention: | |
| self.state_self_attention = SelfAttentionWithRoPE( | |
| d_model=d_model, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| use_rope=self.use_rope, | |
| qk_norm=qk_norm, | |
| rope_set_aggregation=rope_set_aggregation, | |
| rope_dim=rope_dim, | |
| dim_per_coord_3d=dim_per_coord_3d, | |
| dim_per_coord_multi=dim_per_coord_multi, | |
| elementwise_gate=elementwise_gate | |
| ) | |
| self.state_norm0 = get_norm_layer(norm_type, d_model) # Norm for state self-attention | |
| # Objects self-attention for object-to-object interaction. Spatial encoding | |
| # uses multi-point RoPE over anchor positions, or an SE(3) pairwise bias | |
| # when use_se3=True. | |
| if objects_self_attention: | |
| if use_se3: | |
| from src.models.se3_attention import SelfAttentionWithSE3Bias | |
| self.obj_self_attn = SelfAttentionWithSE3Bias( | |
| d_model=d_model, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| qk_norm=qk_norm, | |
| elementwise_gate=elementwise_gate | |
| ) | |
| logger.info(f"[Decoder Init] obj_self_attn = SelfAttentionWithSE3Bias (no RoPE, pairwise bias)") | |
| else: | |
| self.obj_self_attn = SelfAttentionWithRoPE( | |
| d_model=d_model, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| use_rope=self.use_rope, | |
| qk_norm=qk_norm, | |
| rope_set_aggregation=rope_set_aggregation, | |
| rope_dim=rope_dim, | |
| dim_per_coord_3d=dim_per_coord_3d, | |
| dim_per_coord_multi=dim_per_coord_multi, | |
| elementwise_gate=elementwise_gate | |
| ) | |
| self.obj_norm0 = get_norm_layer(norm_type, d_model) # Norm for objects self-attention | |
| # Feedforward networks | |
| self.state_ffn = FeedForward(d_model, feedforward_dim, dropout, activation) | |
| self.obs_ffn = FeedForward(d_model, feedforward_dim, dropout, activation) | |
| # Layer normalization (pre-norm) | |
| self.state_norm1 = get_norm_layer(norm_type, d_model) | |
| self.state_norm2 = get_norm_layer(norm_type, d_model) | |
| self.obs_norm1 = get_norm_layer(norm_type, d_model) | |
| self.obs_norm2 = get_norm_layer(norm_type, d_model) | |
| # Per-layer FiLM conditioning on step size | |
| self.layer_film = None | |
| if film_conditioning: | |
| from src.models.encoder import StepSizeFiLM | |
| self.layer_film = StepSizeFiLM(time_dim=2, feature_dim=d_model) | |
| # Dropout | |
| self.dropout = nn.Dropout(dropout) | |
| self.drop_path = DropPath(drop_path_prob) if drop_path_prob and drop_path_prob > 0.0 else nn.Identity() | |
| def forward( | |
| self, | |
| state: torch.Tensor, | |
| obj_features: torch.Tensor, | |
| state_pos: Optional[torch.Tensor] = None, | |
| obj_anchor_pos: Optional[torch.Tensor] = None, | |
| mask: Optional[torch.Tensor] = None, | |
| se3_bias: Optional[torch.Tensor] = None, # (B, H, Nobj, Nobj) pairwise SE(3) bias | |
| time_features: Optional[torch.Tensor] = None # (B, 2) [s, s^2] for per-layer FiLM | |
| ) -> tuple: | |
| """ | |
| Bidirectional or unidirectional cross-attention. | |
| When state_update_enabled=True (default): | |
| - Full bidirectional interaction: state and objects attend to each other | |
| - Optional state self-attention | |
| - Both state and objects are updated | |
| When state_update_enabled=False: | |
| - Unidirectional: Objects <- State only | |
| - State remains fixed learnable embeddings | |
| Args: | |
| state: State tokens (B, N_state, D) | |
| obj_features: Per-object encoder features (B, N_obj, D) | |
| state_pos: State 3D positions for RoPE (B, N_state, 3) | |
| obj_anchor_pos: Object anchor positions for RoPE (B, N_obj, Nq*3) | |
| mask: Optional attention mask | |
| Returns: | |
| Updated state and obj_features | |
| """ | |
| no_state = state is None | |
| if self.state_update_enabled and not no_state: | |
| # Full bidirectional mode | |
| # Optional state self-attention with RoPE. The object mask does not | |
| # apply to state tokens; pass None. | |
| if self.use_self_attention: | |
| state_sa = self.state_norm0(state) | |
| state_sa = self.state_self_attention( | |
| query=state_sa, | |
| key_value=state_sa, | |
| pos_q=state_pos, # Use 3D grid positions for RoPE | |
| pos_k=state_pos, | |
| attn_mask=None # Object mask does not apply to state tokens | |
| ) | |
| state = state + self.drop_path(self.dropout(state_sa)) | |
| # Bidirectional cross-attention | |
| # State <- Objects: state queries per-object features | |
| state_normed = self.state_norm1(state) | |
| obj_normed = self.obs_norm1(obj_features) | |
| state2 = self.state_cross_attn( | |
| query=state_normed, | |
| key_value=obj_normed, | |
| pos_q=state_pos, | |
| pos_k=obj_anchor_pos, | |
| attn_mask=mask | |
| ) | |
| state = state + self.drop_path(self.dropout(state2)) | |
| # Objects <- State: per-object features query state | |
| obj2 = self.obs_cross_attn( | |
| query=obj_normed, | |
| key_value=state_normed, | |
| pos_q=obj_anchor_pos, | |
| pos_k=state_pos, | |
| attn_mask=None # State keys are all valid, no need to mask | |
| ) | |
| obj_features = obj_features + self.drop_path(self.dropout(obj2)) | |
| # Per-layer FiLM conditioning on step size | |
| if self.layer_film is not None and time_features is not None: | |
| obj_features = self.layer_film(obj_features, time_features) | |
| # Feedforward networks with residual connections | |
| state = state + self.drop_path(self.state_ffn(self.state_norm2(state))) | |
| obj_features = obj_features + self.drop_path(self.obs_ffn(self.obs_norm2(obj_features))) | |
| else: | |
| # Unidirectional / no-state mode | |
| # no_state=True: objects self-attention + objects FFN only | |
| # no_state=False: objects self-attention + Objects <- State + objects FFN | |
| # Objects self-attention runs before Objects <- State so objects | |
| # exchange information first | |
| if self.objects_self_attention: | |
| # One-time runtime confirmation log | |
| if not self._objects_self_attention_logged: | |
| mode_str = 'NO-STATE' if no_state else 'UNIDIRECTIONAL' | |
| print(f"[Decoder Runtime] Objects self-attention active ({mode_str}): shape={obj_features.shape}, se3_bias={'yes' if se3_bias is not None else 'no'}") | |
| self._objects_self_attention_logged = True | |
| obj_sa = self.obj_norm0(obj_features) | |
| if se3_bias is not None: | |
| # SE(3) path: use pairwise bias instead of RoPE | |
| obj_sa = self.obj_self_attn( | |
| query=obj_sa, | |
| key_value=obj_sa, | |
| se3_bias=se3_bias, | |
| attn_mask=mask | |
| ) | |
| else: | |
| # ARoPE/OBB/PCA path: use RoPE with anchor positions | |
| obj_sa = self.obj_self_attn( | |
| query=obj_sa, | |
| key_value=obj_sa, | |
| pos_q=obj_anchor_pos, # Multi-point anchor positions for RoPE | |
| pos_k=obj_anchor_pos, | |
| attn_mask=mask # Mask padding objects | |
| ) | |
| obj_features = obj_features + self.drop_path(self.dropout(obj_sa)) | |
| # Objects <- State cross-attention (skipped when no_state=True) | |
| if not no_state: | |
| state_normed = self.state_norm1(state) | |
| obj_normed = self.obs_norm1(obj_features) | |
| # Objects <- State: per-object features query state (learnable context) | |
| obj2 = self.obs_cross_attn( | |
| query=obj_normed, | |
| key_value=state_normed, | |
| pos_q=obj_anchor_pos, | |
| pos_k=state_pos, | |
| attn_mask=None # State keys are all valid, no need to mask | |
| ) | |
| obj_features = obj_features + self.drop_path(self.dropout(obj2)) | |
| # Per-layer FiLM conditioning on step size | |
| # Applied after attention, before the FFN. Padding objects may receive | |
| # non-zero features from the FiLM beta bias; the SDPA -inf mask keeps | |
| # valid objects from attending to them. | |
| if self.layer_film is not None and time_features is not None: | |
| obj_features = self.layer_film(obj_features, time_features) | |
| # Objects FFN | |
| obj_features = obj_features + self.drop_path(self.obs_ffn(self.obs_norm2(obj_features))) | |
| return state, obj_features | |
| class SelfAttentionWithRoPE(nn.Module): | |
| """ | |
| Self-attention that applies RoPE using 3D positions. | |
| QK-normalization (RenderFormer) improves training stability; elementwise | |
| gated attention (Qwen3) adds non-linearity and sparsity. | |
| """ | |
| def __init__( | |
| self, | |
| d_model: int, | |
| num_heads: int, | |
| dropout: float = 0.1, | |
| use_rope: bool = True, | |
| qk_norm: bool = True, # QK-normalization (RenderFormer) | |
| rope_set_aggregation: bool = False, | |
| rope_dim: int = None, # Custom RoPE dimension (None = auto-compute) | |
| dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE | |
| dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE | |
| elementwise_gate: bool = False # Elementwise gated attention (Qwen3 style) | |
| ): | |
| super().__init__() | |
| assert d_model % num_heads == 0, "d_model must be divisible by num_heads" | |
| self.d_model = d_model | |
| self.num_heads = num_heads | |
| self.head_dim = d_model // num_heads | |
| self.scale = self.head_dim ** -0.5 | |
| self.use_rope = bool(use_rope) | |
| self.qk_norm = qk_norm | |
| self.elementwise_gate = elementwise_gate | |
| # Projections; with elementwise_gate, q_proj outputs 2x dims (Q and gate scores) | |
| if self.elementwise_gate: | |
| self.q_proj = nn.Linear(d_model, d_model * 2, bias=True) | |
| else: | |
| self.q_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.k_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.v_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.out_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.dropout = nn.Dropout(dropout) | |
| self._attn_dispatch_logged = False # One-time attention dispatch log | |
| # QK-normalization: per-head RMSNorm (RenderFormer) | |
| if self.qk_norm: | |
| self.q_norm = RMSNorm(self.head_dim) | |
| self.k_norm = RMSNorm(self.head_dim) | |
| # 3D rotary embedding; only the first rope_dim dims are rotated | |
| if rope_dim is None: | |
| # Auto-compute: multiple of 6, capped at head_dim | |
| self.rope_dim = max(6, (self.head_dim // 6) * 6) | |
| self.rope_dim = min(self.rope_dim, self.head_dim) | |
| else: | |
| # Custom rope_dim (must be a multiple of 6 for 3D RoPE) | |
| assert rope_dim % 6 == 0, f"rope_dim must be multiple of 6 for 3D RoPE, got {rope_dim}" | |
| assert rope_dim <= self.head_dim, f"rope_dim {rope_dim} exceeds head_dim {self.head_dim}" | |
| self.rope_dim = rope_dim | |
| self.rope3d = PhysFormerRoPE( | |
| dim=self.rope_dim, | |
| max_freq=10.0, | |
| use_3d=True, | |
| learnable_freqs=False, | |
| use_set_aggregation=rope_set_aggregation, | |
| dim_per_coord_3d=dim_per_coord_3d, # Pass to PhysFormerRoPE | |
| dim_per_coord_multi=dim_per_coord_multi # Pass to PhysFormerRoPE | |
| ) | |
| def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor: | |
| B, N, D = x.shape | |
| x = x.view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # (B,H,N,Hd) | |
| return x | |
| def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: | |
| B, H, N, Hd = x.shape | |
| return x.permute(0, 2, 1, 3).contiguous().view(B, N, H * Hd) | |
| def forward( | |
| self, | |
| query: torch.Tensor, | |
| key_value: torch.Tensor, | |
| pos_q: torch.Tensor, | |
| pos_k: torch.Tensor, | |
| attn_mask: Optional[torch.Tensor] = None | |
| ) -> torch.Tensor: | |
| """ | |
| Self-attention using Flash Attention, or SDPA when a mask is provided. | |
| Args: | |
| query: Input tokens (B, N, D) | |
| key_value: Same as query for self-attention (B, N, D) | |
| pos_q: 3D positions (B, N, 3) or multi-point positions (B, N, Nq*3) | |
| pos_k: Same as pos_q for self-attention | |
| attn_mask: Optional additive mask; when present, attention runs | |
| through SDPA since Flash does not support masks | |
| Returns: | |
| Output tokens (B, N, D) | |
| """ | |
| B, N, _ = query.shape | |
| # Q projection with optional elementwise gate | |
| if self.elementwise_gate: | |
| q_gate = self.q_proj(query) # (B, N, d_model * 2) | |
| q_states, gate_score = torch.split(q_gate, self.d_model, dim=-1) | |
| q = q_states.view(B, N, self.num_heads, self.head_dim) # (B, N, H, D) - Flash expects (B, S, H, D) | |
| gate_score = gate_score.view(B, N, self.num_heads, self.head_dim) # (B, N, H, D) | |
| else: | |
| q = self.q_proj(query).view(B, N, self.num_heads, self.head_dim) # (B, N, H, D) | |
| gate_score = None | |
| k = self.k_proj(key_value).view(B, N, self.num_heads, self.head_dim) # (B, N, H, D) | |
| v = self.v_proj(key_value).view(B, N, self.num_heads, self.head_dim) # (B, N, H, D) | |
| # QK-normalization (RenderFormer) | |
| if self.qk_norm: | |
| # RMSNorm normalizes the last (per-head) dimension | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| # RoPE (single-point 3D or multi-point) | |
| if pos_q is None or pos_k is None: | |
| raise ValueError("SelfAttentionWithRoPE expects pos_q and pos_k") | |
| # Single-point (3D) vs multi-point positions | |
| if pos_q.shape[-1] == 3: | |
| rope_q = self.rope3d.forward_3d(pos_q) # (B, N, rope_dim) | |
| else: | |
| rope_q = self.rope3d.forward_sa_rope(pos_q) # SA-RoPE: symmetric anchor pooling | |
| if pos_k is pos_q: | |
| rope_k = rope_q # identical positions -> identical angles (computed once) | |
| elif pos_k.shape[-1] == 3: | |
| rope_k = self.rope3d.forward_3d(pos_k) # (B, N, rope_dim) | |
| else: | |
| rope_k = self.rope3d.forward_sa_rope(pos_k) # SA-RoPE: symmetric anchor pooling | |
| # The angles stay at (B, 1, N, rope_dim) and broadcast over heads inside apply_rotary_emb. cos/sin | |
| # are elementwise, so this equals expanding to (B, H, N, rope_dim) first, at 1/num_heads of the work. | |
| rope_q_T = rope_q.unsqueeze(1) | |
| rope_k_T = rope_k.unsqueeze(1) if rope_k is not rope_q else rope_q_T | |
| # Apply rotation | |
| q_rot = q[..., :self.rope_dim] | |
| k_rot = k[..., :self.rope_dim] | |
| q_pass = q[..., self.rope_dim:] | |
| k_pass = k[..., self.rope_dim:] | |
| # apply_rotary_emb operates on (B, H, N, D): transpose for the rotation, | |
| # then transpose back to the (B, N, H, D) layout Flash expects. | |
| q_rot = q_rot.transpose(1, 2) # (B, H, N, D) | |
| k_rot = k_rot.transpose(1, 2) | |
| q_rot = apply_rotary_emb(rope_q_T, q_rot) | |
| k_rot = apply_rotary_emb(rope_k_T, k_rot) | |
| # Permute back to (B, N, H, D) | |
| q_rot = q_rot.transpose(1, 2) | |
| k_rot = k_rot.transpose(1, 2) | |
| q = torch.cat([q_rot, q_pass], dim=-1) | |
| k = torch.cat([k_rot, k_pass], dim=-1) | |
| # flash_attn_func does not support additive masks: dispatch to SDPA when a | |
| # mask is provided, Flash Attention otherwise. | |
| if attn_mask is not None: | |
| if not self._attn_dispatch_logged: | |
| print(f"[SelfAttentionWithRoPE] SDPA path (mask active): Q={q.shape}, K={k.shape}, mask={attn_mask.shape}, " | |
| f"mask_min={attn_mask.min().item():.0f}, n_masked={(attn_mask == float('-inf')).sum().item()}") | |
| self._attn_dispatch_logged = True | |
| # SDPA expects (B, H, N, D) layout | |
| q_sdpa = q.transpose(1, 2) # (B, H, N, D) | |
| k_sdpa = k.transpose(1, 2) | |
| v_sdpa = v.transpose(1, 2) | |
| out = F.scaled_dot_product_attention( | |
| q_sdpa, k_sdpa, v_sdpa, | |
| attn_mask=attn_mask, | |
| dropout_p=self.dropout.p if self.training else 0.0, | |
| scale=self.scale | |
| ) # (B, H, N, D) | |
| out = out.transpose(1, 2) # Back to (B, N, H, D) | |
| else: | |
| if not self._attn_dispatch_logged: | |
| print(f"[SelfAttentionWithRoPE] Flash path (no mask): Q={q.shape}") | |
| self._attn_dispatch_logged = True | |
| # Flash Attention runs in bf16/fp16; the output is cast back to the | |
| # original dtype immediately. | |
| original_dtype = query.dtype | |
| attn_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 | |
| q_half = q.to(attn_dtype) | |
| k_half = k.to(attn_dtype) | |
| v_half = v.to(attn_dtype) | |
| out = flash_attn_func(q_half, k_half, v_half, dropout_p=self.dropout.p if self.training else 0.0, softmax_scale=self.scale) | |
| out = out.to(original_dtype) | |
| if self.elementwise_gate: | |
| out = out * torch.sigmoid(gate_score) | |
| # Merge heads: (B, N, H, D) -> (B, N, H*D) | |
| out = out.reshape(B, N, self.num_heads * self.head_dim) | |
| out = self.out_proj(out) | |
| return out | |
| class CrossAttentionWithRoPE(nn.Module): | |
| """ | |
| Cross-attention that applies RoPE to queries and keys. | |
| QK-normalization (RenderFormer) improves training stability; elementwise | |
| gated attention (Qwen3) adds non-linearity and sparsity. | |
| """ | |
| def __init__( | |
| self, | |
| d_model: int, | |
| num_heads: int, | |
| dropout: float = 0.1, | |
| use_rope: bool = True, | |
| qk_norm: bool = True, # QK-normalization (RenderFormer) | |
| rope_set_aggregation: bool = False, | |
| rope_dim: int = None, # Custom RoPE dimension (None = auto-compute) | |
| dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE | |
| dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE | |
| elementwise_gate: bool = False # Elementwise gated attention (Qwen3 style) | |
| ): | |
| super().__init__() | |
| assert d_model % num_heads == 0, "d_model must be divisible by num_heads" | |
| self.d_model = d_model | |
| self.num_heads = num_heads | |
| self.head_dim = d_model // num_heads | |
| self.scale = self.head_dim ** -0.5 | |
| self.use_rope = bool(use_rope) | |
| self.qk_norm = qk_norm | |
| self.elementwise_gate = elementwise_gate | |
| # Projections; with elementwise_gate, q_proj outputs 2x dims (Q and gate scores) | |
| if self.elementwise_gate: | |
| self.q_proj = nn.Linear(d_model, d_model * 2, bias=True) | |
| else: | |
| self.q_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.k_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.v_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.out_proj = nn.Linear(d_model, d_model, bias=True) | |
| self.dropout = nn.Dropout(dropout) | |
| self._attn_dispatch_logged = False # One-time attention dispatch log | |
| # QK-normalization: per-head RMSNorm (RenderFormer) | |
| if self.qk_norm: | |
| self.q_norm = RMSNorm(self.head_dim) | |
| self.k_norm = RMSNorm(self.head_dim) | |
| # 3D rotary embedding; only the first rope_dim dims are rotated | |
| if rope_dim is None: | |
| # Auto-compute: multiple of 6, capped at head_dim | |
| self.rope_dim = max(6, (self.head_dim // 6) * 6) | |
| self.rope_dim = min(self.rope_dim, self.head_dim) | |
| else: | |
| # Custom rope_dim (must be a multiple of 6 for 3D RoPE) | |
| assert rope_dim % 6 == 0, f"rope_dim must be multiple of 6 for 3D RoPE, got {rope_dim}" | |
| assert rope_dim <= self.head_dim, f"rope_dim {rope_dim} exceeds head_dim {self.head_dim}" | |
| self.rope_dim = rope_dim | |
| self.rope3d = PhysFormerRoPE( | |
| dim=self.rope_dim, | |
| max_freq=10.0, | |
| use_3d=True, | |
| learnable_freqs=False, | |
| use_set_aggregation=rope_set_aggregation, | |
| dim_per_coord_3d=dim_per_coord_3d, # Pass to PhysFormerRoPE | |
| dim_per_coord_multi=dim_per_coord_multi # Pass to PhysFormerRoPE | |
| ) | |
| def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor: | |
| B, N, D = x.shape | |
| x = x.view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # (B,H,N,Hd) | |
| return x | |
| def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: | |
| B, H, N, Hd = x.shape | |
| return x.permute(0, 2, 1, 3).contiguous().view(B, N, H * Hd) | |
| def forward( | |
| self, | |
| query: torch.Tensor, | |
| key_value: torch.Tensor, | |
| pos_q: torch.Tensor, | |
| pos_k: torch.Tensor, | |
| attn_mask: Optional[torch.Tensor] = None | |
| ) -> torch.Tensor: | |
| """ | |
| Cross-attention using Flash Attention v2. | |
| """ | |
| B, Nq, _ = query.shape | |
| _, Nk, _ = key_value.shape | |
| # Projections => (B, N, H, D) for Flash | |
| if self.elementwise_gate: | |
| q_gate = self.q_proj(query) | |
| q_states, gate_score = torch.split(q_gate, self.d_model, dim=-1) | |
| q = q_states.view(B, Nq, self.num_heads, self.head_dim) | |
| gate_score = gate_score.view(B, Nq, self.num_heads, self.head_dim) | |
| else: | |
| q = self.q_proj(query).view(B, Nq, self.num_heads, self.head_dim) | |
| gate_score = None | |
| k = self.k_proj(key_value).view(B, Nk, self.num_heads, self.head_dim) | |
| v = self.v_proj(key_value).view(B, Nk, self.num_heads, self.head_dim) | |
| # QK Norm | |
| if self.qk_norm: | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| # RoPE | |
| if pos_q is None or pos_k is None: | |
| raise ValueError("CrossAttentionWithRoPE expects pos_q and pos_k") | |
| # Single-point (3D) vs multi-point positions | |
| if pos_q.shape[-1] == 3: | |
| rope_q = self.rope3d.forward_3d(pos_q) | |
| else: | |
| rope_q = self.rope3d.forward_sa_rope(pos_q) # SA-RoPE: symmetric anchor pooling | |
| if pos_k is pos_q: | |
| rope_k = rope_q # identical positions -> identical angles (computed once) | |
| elif pos_k.shape[-1] == 3: | |
| rope_k = self.rope3d.forward_3d(pos_k) | |
| else: | |
| rope_k = self.rope3d.forward_sa_rope(pos_k) # SA-RoPE: symmetric anchor pooling | |
| # (B, 1, N, rope_dim) angles broadcast over heads (see SelfAttentionWithRoPE) | |
| rope_q_T = rope_q.unsqueeze(1) | |
| rope_k_T = rope_k.unsqueeze(1) if rope_k is not rope_q else rope_q_T | |
| # apply_rotary_emb operates on (B, H, N, D): transpose for the rotation and back | |
| q_rot = q[..., :self.rope_dim].transpose(1, 2) | |
| k_rot = k[..., :self.rope_dim].transpose(1, 2) | |
| q_rot = apply_rotary_emb(rope_q_T, q_rot) | |
| k_rot = apply_rotary_emb(rope_k_T, k_rot) | |
| q_rot = q_rot.transpose(1, 2) | |
| k_rot = k_rot.transpose(1, 2) | |
| q_pass = q[..., self.rope_dim:] | |
| k_pass = k[..., self.rope_dim:] | |
| q = torch.cat([q_rot, q_pass], dim=-1) | |
| k = torch.cat([k_rot, k_pass], dim=-1) | |
| # flash_attn_func does not support additive masks: dispatch to SDPA when a | |
| # mask is provided, Flash Attention otherwise. | |
| if attn_mask is not None: | |
| if not self._attn_dispatch_logged: | |
| n_inf = (attn_mask == float('-inf')).sum().item() | |
| n_total = attn_mask.numel() | |
| print(f"[CrossAttentionWithRoPE] SDPA path (mask active): Q={q.shape}, K={k.shape}, " | |
| f"mask={attn_mask.shape}, masked={n_inf}/{n_total} ({100*n_inf/max(n_total,1):.1f}%)") | |
| self._attn_dispatch_logged = True | |
| # SDPA expects (B, H, N, D) layout | |
| q_sdpa = q.transpose(1, 2) # (B, H, Nq, D) | |
| k_sdpa = k.transpose(1, 2) # (B, H, Nk, D) | |
| v_sdpa = v.transpose(1, 2) # (B, H, Nk, D) | |
| out = F.scaled_dot_product_attention( | |
| q_sdpa, k_sdpa, v_sdpa, | |
| attn_mask=attn_mask, | |
| dropout_p=self.dropout.p if self.training else 0.0, | |
| scale=self.scale | |
| ) # (B, H, Nq, D) | |
| out = out.transpose(1, 2) # Back to (B, Nq, H, D) | |
| else: | |
| if not self._attn_dispatch_logged: | |
| print(f"[CrossAttentionWithRoPE] Flash path (no mask): Q={q.shape}, K={k.shape}") | |
| self._attn_dispatch_logged = True | |
| # Flash Attention runs in bf16/fp16; the output is cast back to the | |
| # original dtype immediately. | |
| original_dtype = query.dtype | |
| attn_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 | |
| q_half = q.to(attn_dtype) | |
| k_half = k.to(attn_dtype) | |
| v_half = v.to(attn_dtype) | |
| out = flash_attn_func(q_half, k_half, v_half, dropout_p=self.dropout.p if self.training else 0.0, softmax_scale=self.scale) | |
| out = out.to(original_dtype) | |
| if self.elementwise_gate: | |
| out = out * torch.sigmoid(gate_score) | |
| out = out.reshape(B, Nq, self.num_heads * self.head_dim) | |
| out = self.out_proj(out) | |
| return out | |
| class FeedForward(nn.Module): | |
| """Feedforward network with activation and dropout. | |
| Supports GELU/RELU and SwiGLU (default). For SwiGLU, hidden_dim is the inner dimension. | |
| """ | |
| def __init__( | |
| self, | |
| d_model: int, | |
| hidden_dim: int, | |
| dropout: float = 0.1, | |
| activation: str = 'swiglu' | |
| ): | |
| super().__init__() | |
| self.activation_name = activation | |
| self.dropout = nn.Dropout(dropout) | |
| if activation == 'swiglu': | |
| self.fc1 = nn.Linear(d_model, hidden_dim) | |
| self.fc1_gate = nn.Linear(d_model, hidden_dim) | |
| self.fc2 = nn.Linear(hidden_dim, d_model) | |
| else: | |
| self.fc1 = nn.Linear(d_model, hidden_dim) | |
| self.fc2 = nn.Linear(hidden_dim, d_model) | |
| if activation == 'gelu': | |
| self.act = nn.GELU() | |
| elif activation == 'relu': | |
| self.act = nn.ReLU() | |
| elif activation == 'silu': | |
| self.act = nn.SiLU() | |
| else: | |
| raise ValueError(f"Unknown activation: {activation}") | |
| def forward(self, x): | |
| if self.activation_name == 'swiglu': | |
| a = self.fc1(x) | |
| b = torch.nn.functional.silu(self.fc1_gate(x)) | |
| y = a * b | |
| y = self.dropout(y) | |
| y = self.fc2(y) | |
| y = self.dropout(y) | |
| return y | |
| else: | |
| y = self.fc1(x) | |
| y = self.act(y) | |
| y = self.dropout(y) | |
| y = self.fc2(y) | |
| y = self.dropout(y) | |
| return y | |
| class DropPath(nn.Module): | |
| """Stochastic Depth per sample (DropPath). | |
| """ | |
| def __init__(self, drop_prob: float = 0.0): | |
| super().__init__() | |
| self.drop_prob = drop_prob | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if self.drop_prob == 0.0 or not self.training: | |
| return x | |
| keep_prob = 1 - self.drop_prob | |
| shape = (x.shape[0],) + (1,) * (x.ndim - 1) | |
| random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) | |
| random_tensor.floor_() | |
| return x.div(keep_prob) * random_tensor | |
| if __name__ == "__main__": | |
| """Test the decoder module.""" | |
| print("Testing BidirectionalDecoder...") | |
| decoder = BidirectionalDecoder( | |
| state_dim=512, | |
| num_layers=4, | |
| num_heads=8, | |
| feedforward_dim=2048 | |
| ) | |
| # Create dummy inputs | |
| batch_size = 2 | |
| state = torch.randn(batch_size, 256, 512) # 256 state tokens | |
| observation = torch.randn(batch_size, 10, 512) # 10 observation tokens | |
| # Forward pass | |
| output = decoder(state, observation) | |
| print(f"Output shape: {output.shape}") # Should be (2, 10, 512) | |
| print("\nDecoder test successful!") | |