"""Transformer Decoder implementation (Pre-LN). This module implements the decoder component of the Transformer architecture: - create_causal_mask: Generate causal attention masks - TransformerDecoderLayer: Single decoder block with self-attn + cross-attn + FFN - TransformerDecoder: Full stack with embeddings, positional encoding, and generation Design notes: - Pre-LN with RMSNorm for training stability - Masks are boolean: True = attend, False = mask - Supports T5-style relative position bias Author: Oliver Perrin Date: 2025-10-23 """ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint from .attention import MultiHeadAttention, T5RelativePositionBias from .feedforward import FeedForward from .positional_encoding import LearnedPositionalEncoding, PositionalEncoding from .t5_layer_norm import T5LayerNorm def create_causal_mask(seq_len: int, device: Optional[torch.device] = None) -> torch.Tensor: """ Create a (seq_len, seq_len) causal mask where entry (i, j) is True iff j <= i (query at i may attend to keys up to i). """ # torch.triu(..., diagonal=1) is True above the diagonal. Invert to get allowed positions. mask = ~torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1) return mask # shape: (T, T) class TransformerDecoderLayer(nn.Module): """ Single decoder layer (Pre-LN): 1) Masked self-attention 2) Cross-attention (encoder -> decoder) 3) Feed-forward Returns the updated tgt and a dict of attention maps. """ def __init__( self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1, quantization: Optional[str] = None, activation: Literal["gelu", "relu", "swiglu", "gated-gelu"] = "gated-gelu", scale_attn_scores: bool = True, # T5 uses False ): super().__init__() # use internal MHA dropout = 0.0; the layer handles dropout after sublayers self.self_attn = MultiHeadAttention( d_model=d_model, num_heads=num_heads, dropout=0.0, quantization=quantization, scale_scores=scale_attn_scores, ) self.cross_attn = MultiHeadAttention( d_model=d_model, num_heads=num_heads, dropout=0.0, quantization=quantization, scale_scores=scale_attn_scores, ) self.ffn = FeedForward( d_model=d_model, d_ff=d_ff, dropout=dropout, activation=activation, quantization=quantization, ) self.norm1 = T5LayerNorm(d_model) self.norm2 = T5LayerNorm(d_model) self.norm3 = T5LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) def forward( self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, memory_mask: Optional[torch.Tensor] = None, collect_attn: bool = False, self_attn_position_bias: Optional[torch.Tensor] = None, cross_attn_position_bias: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]]: """ Args: tgt: (B, T, d_model) memory: (B, S, d_model) tgt_mask: optional mask for self-attn - shape (B, T, T) or (B, 1, T, T) memory_mask: optional mask for cross-attn - shape (B, S) or (B, 1, S) or (B, 1, T, S) collect_attn: whether to return attention weights self_attn_position_bias: optional T5 relative position bias for self-attention cross_attn_position_bias: optional T5 relative position bias for cross-attention Returns: (tgt_out, {"self": self_attn_weights, "cross": cross_attn_weights}) """ # Ensure masks are on same device and boolean if tgt_mask is not None: tgt_mask = tgt_mask.to(dtype=torch.bool, device=tgt.device) if memory_mask is not None: memory_mask = memory_mask.to(dtype=torch.bool, device=tgt.device) # If memory_mask is provided as (B, S) (per-key padding), expand to (B, 1, 1, S) if memory_mask.dim() == 2: memory_mask = memory_mask.unsqueeze(1).unsqueeze(1) # (B,1,1,S) # If it's (B, S, S) or (B, 1, S, S) leave as-is; if (B, T, S) convert to (B,1,T,S) elif memory_mask.dim() == 3 and memory_mask.shape[1] != 1: # assume (B, T, S) -> make (B, 1, T, S) memory_mask = memory_mask.unsqueeze(1) # --- Masked self-attention (Pre-LN) --- x_norm = self.norm1(tgt) self_out, self_attn = self.self_attn( x_norm, x_norm, x_norm, tgt_mask, return_attn_weights=collect_attn, position_bias=self_attn_position_bias, ) tgt = tgt + self.dropout1(self_out) # Clamp inf values for fp16/bf16 training stability (like HuggingFace T5) if tgt.dtype == torch.float16 or tgt.dtype == torch.bfloat16: clamp_value = torch.finfo(tgt.dtype).max - 1000 tgt = torch.clamp(tgt, min=-clamp_value, max=clamp_value) # --- Cross-attention (Pre-LN) --- x_norm = self.norm2(tgt) cross_out, cross_attn = self.cross_attn( x_norm, memory, memory, memory_mask, return_attn_weights=collect_attn, position_bias=cross_attn_position_bias, ) tgt = tgt + self.dropout2(cross_out) # Clamp inf values for fp16/bf16 training stability if tgt.dtype == torch.float16 or tgt.dtype == torch.bfloat16: clamp_value = torch.finfo(tgt.dtype).max - 1000 tgt = torch.clamp(tgt, min=-clamp_value, max=clamp_value) # --- Feed-forward (Pre-LN) --- x_norm = self.norm3(tgt) ffn_out = self.ffn(x_norm) tgt = tgt + self.dropout3(ffn_out) # Clamp inf values for fp16/bf16 training stability if tgt.dtype == torch.float16 or tgt.dtype == torch.bfloat16: clamp_value = torch.finfo(tgt.dtype).max - 1000 tgt = torch.clamp(tgt, min=-clamp_value, max=clamp_value) return tgt, {"self": self_attn, "cross": cross_attn} class TransformerDecoder(nn.Module): """ Decoder stack with token embeddings and positional encoding. Forward returns logits (B, T, vocab_size) by default; if collect_attn=True returns (logits, attn_list). """ def __init__( self, vocab_size: int, d_model: int = 512, num_layers: int = 6, num_heads: int = 8, d_ff: int = 2048, dropout: float = 0.1, max_len: int = 512, pad_token_id: Optional[int] = None, quantization: Optional[str] = None, use_learned_pos_enc: bool = False, activation: Literal["gelu", "relu", "swiglu", "gated-gelu"] = "gated-gelu", use_relative_position_bias: bool = False, # T5-style relative position bias gradient_checkpointing: bool = False, ): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.pad_token_id = pad_token_id self.num_heads = num_heads self.use_relative_position_bias = use_relative_position_bias self.gradient_checkpointing = gradient_checkpointing self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_token_id) # Note: T5 does NOT scale logits (scaling factor removed) # Positional encoding (disabled when using relative position bias for T5) self.self_relative_position_bias: Optional[T5RelativePositionBias] = None self.cross_relative_position_bias: Optional[T5RelativePositionBias] = None if use_relative_position_bias: # T5 uses relative position bias instead of absolute positional embeddings self.pos_encoder = None # Self-attention position bias (decoder is causal, so is_decoder=True) self.self_relative_position_bias = T5RelativePositionBias( num_heads=num_heads, num_buckets=32, max_distance=128, is_decoder=True, ) # T5 cross-attention does NOT use position bias elif use_learned_pos_enc: self.pos_encoder = LearnedPositionalEncoding( d_model=d_model, max_len=max_len + 2, dropout=dropout ) else: self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=max_len, dropout=dropout) # T5 does NOT scale attention scores by sqrt(d_k), others do scale_attn_scores = not use_relative_position_bias self.layers = nn.ModuleList( [ TransformerDecoderLayer( d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=dropout, quantization=quantization, activation=activation, scale_attn_scores=scale_attn_scores, ) for _ in range(num_layers) ] ) self.final_norm = T5LayerNorm(d_model) self.output_projection = nn.Linear(d_model, vocab_size, bias=False) # T5 has no bias self.input_dropout = nn.Dropout(dropout) def _build_padding_mask_from_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """ Convert input ids to (B, T, T) boolean mask where True = allowed. Note: For T5, pad_token_id=0 is also used as decoder_start_token_id. During generation, we should NOT mask the start token. The caller should provide an explicit mask or set tgt_mask to avoid this issue. """ assert self.pad_token_id is not None, "pad_token_id must be set to build mask from ids" pad_mask = input_ids != self.pad_token_id # (B, T) # Always allow attending to the first token (BOS), even if it is pad_token_id # Avoid in-place mutation for better torch.compile compatibility if pad_mask.size(1) > 0: # Create a mask for the first column (B, 1) first_col_mask = torch.zeros_like(pad_mask[:, :1], dtype=torch.bool) first_col_mask[:] = True # Combine: pad_mask OR (column==0) # We can do this by creating a column index tensor col_indices = torch.arange(pad_mask.size(1), device=pad_mask.device).unsqueeze(0) pad_mask = pad_mask | (col_indices == 0) attn_mask = pad_mask.unsqueeze(1) & pad_mask.unsqueeze(2) # (B, T, T) return attn_mask def forward( self, inputs: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, memory_mask: Optional[torch.Tensor] = None, collect_attn: bool = False, skip_padding_mask: bool = False, # Set True during generation to avoid masking start token ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict[str, Optional[torch.Tensor]]]]]: """ Args: inputs: (B, T) token ids or (B, T, d_model) embeddings memory: (B, S, d_model) tgt_mask: optional; if None, will create (causal [+ padding if ids available]) memory_mask: optional; if provided as (B, S) will be expanded to (B, 1, 1, S) skip_padding_mask: if True, only use causal mask (for generation where start_token=pad_token) """ # Prepare embeddings if inputs.dim() == 2: # token ids # T5/FLAN-T5 does NOT scale embeddings by sqrt(d_model) x = self.embedding(inputs) elif inputs.dim() == 3: x = inputs else: raise ValueError("inputs must be (B, T) token ids or (B, T, d_model) embeddings") # Apply positional encoding if not using relative position bias # (T5 uses relative position bias in attention instead of absolute positional embeddings) if self.pos_encoder is not None: x = self.pos_encoder(x) x = self.input_dropout(x) B, T, _ = x.shape # Build target mask if not provided: combine causal + padding (if available) if tgt_mask is None: causal = create_causal_mask(T, device=x.device) # (T, T) if inputs.dim() == 2 and self.pad_token_id is not None and not skip_padding_mask: # During training: combine causal mask with padding mask pad_pairwise = self._build_padding_mask_from_ids(inputs) # (B, T, T) combined = pad_pairwise & causal.unsqueeze(0) # (B, T, T) tgt_mask = combined.unsqueeze(1) # (B, 1, T, T) -> broadcast to heads else: # During generation (skip_padding_mask=True) or no padding info: # Use only causal mask - don't mask based on token values tgt_mask = causal.unsqueeze(0).unsqueeze(1) # (1, 1, T, T) else: # Ensure boolean and device alignment; accept (B, T, T) or (B,1,T,T) or (1,1,T,T) tgt_mask = tgt_mask.to(dtype=torch.bool, device=x.device) # If tgt_mask is just causal (T, T), expand it if tgt_mask.dim() == 2: tgt_mask = tgt_mask.unsqueeze(0).unsqueeze(0) elif tgt_mask.dim() == 3: tgt_mask = tgt_mask.unsqueeze(1) # Normalize memory_mask dtype/device and expand simple shapes if memory_mask is not None: memory_mask = memory_mask.to(dtype=torch.bool, device=x.device) if memory_mask.dim() == 2: # (B, S) -> (B, 1, 1, S) memory_mask = memory_mask.unsqueeze(1).unsqueeze(1) elif memory_mask.dim() == 3: # (B, T, S) -> (B, 1, T, S) memory_mask = memory_mask.unsqueeze(1) attn_list: List[Dict[str, Optional[torch.Tensor]]] = [] # Compute relative position biases (T5-style) # Note: T5 uses relative position bias for self-attention but NOT for cross-attention if self.use_relative_position_bias and self.self_relative_position_bias is not None: self_position_bias = self.self_relative_position_bias( T, T, x.device ) # (1, num_heads, T, T) else: self_position_bias = None # Cross-attention position bias is None for T5 (see T5 paper/implementation) cross_position_bias = None # Pass through decoder layers for layer in self.layers: if self.gradient_checkpointing and self.training: # Gradient checkpointing requires the inputs to require grad def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs, tgt_mask=tgt_mask, memory_mask=memory_mask, collect_attn=collect_attn, self_attn_position_bias=self_position_bias, cross_attn_position_bias=cross_position_bias) return custom_forward x, attn = cast( Tuple[torch.Tensor, Dict[str, Optional[torch.Tensor]]], checkpoint( create_custom_forward(layer), x, memory, use_reentrant=False, ), ) else: x, attn = layer( x, memory, tgt_mask=tgt_mask, memory_mask=memory_mask, collect_attn=collect_attn, self_attn_position_bias=self_position_bias, cross_attn_position_bias=cross_position_bias, ) if collect_attn: attn_list.append(attn) x = self.final_norm(x) # T5 does NOT scale logits - direct projection to vocabulary logits = self.output_projection(x) # (B, T, vocab) if collect_attn: return logits, attn_list return logits def greedy_decode_naive( self, memory: torch.Tensor, max_len: int, start_token_id: int, end_token_id: Optional[int] = None, device: Optional[torch.device] = None, memory_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Naive greedy decoding using full forward pass (O(N^2) but simpler). Used for debugging to verify step() correctness. """ if device is None: device = memory.device B = memory.size(0) # Initialize with start token generated = torch.full((B, 1), start_token_id, dtype=torch.long, device=device) for _ in range(max_len - 1): # Full forward pass on entire generated sequence # skip_padding_mask=True because start_token=pad_token for T5 logits = self.forward( generated, memory, memory_mask=memory_mask, skip_padding_mask=True ) if isinstance(logits, tuple): logits = logits[0] # logits: (B, T, vocab) # Get logits for last position next_logits = logits[:, -1, :] # (B, vocab) # Greedy: pick highest probability token next_token = next_logits.argmax(dim=-1, keepdim=True) # (B, 1) # Append to generated generated = torch.cat([generated, next_token], dim=1) # Check for EOS if end_token_id is not None and (next_token == end_token_id).all(): break return generated def greedy_decode( self, memory: torch.Tensor, max_len: int, start_token_id: int, end_token_id: Optional[int] = None, device: Optional[torch.device] = None, *, min_len: Optional[int] = None, ban_token_ids: Optional[List[int]] = None, no_repeat_ngram_size: int = 0, repetition_penalty: float = 1.0, memory_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Greedy decoding with KV caching for O(N) complexity. """ if device is None: device = memory.device B = memory.size(0) # Initialize generated sequence with start token generated = torch.full((B, 1), start_token_id, dtype=torch.long, device=device) # Initialize cache cache: Dict[str, Any] = {"past_length": 0} if memory_mask is not None: cache["memory_mask"] = memory_mask min_len = 0 if min_len is None else max(0, min_len) # Keep track of finished sequences finished = torch.zeros(B, dtype=torch.bool, device=device) for _ in range(max_len - 1): # Use the last generated token for the next step last_token = generated[:, -1:] # (B, 1) # Run one step of the decoder logits, cache = self.step(last_token, memory, cache) # logits: (B, vocab_size) next_step_logits = logits.clone() # Apply repetition penalty if repetition_penalty != 1.0: for b in range(B): if finished[b]: continue gen_seq = generated[b] unique_tokens = torch.unique(gen_seq) current_logits = next_step_logits[b, unique_tokens] next_step_logits[b, unique_tokens] = torch.where( current_logits < 0, current_logits * repetition_penalty, current_logits / repetition_penalty, ) # Apply constraints if end_token_id is not None and generated.size(1) < max(1, min_len): next_step_logits[:, end_token_id] = float("-inf") if ban_token_ids: next_step_logits[:, ban_token_ids] = float("-inf") # N-gram repetition blocking if no_repeat_ngram_size > 0: for b in range(B): if finished[b]: continue gen_seq = generated[b].tolist() if len(gen_seq) < no_repeat_ngram_size - 1: continue prefix = tuple(gen_seq[-(no_repeat_ngram_size - 1) :]) banned_for_this_batch = set() for i in range(len(gen_seq) - no_repeat_ngram_size + 1): window = tuple(gen_seq[i : i + no_repeat_ngram_size - 1]) if window == prefix: if i + no_repeat_ngram_size - 1 < len(gen_seq): banned_for_this_batch.add(gen_seq[i + no_repeat_ngram_size - 1]) if banned_for_this_batch: next_step_logits[b, list(banned_for_this_batch)] = float("-inf") # Greedy selection next_token = next_step_logits.argmax(dim=-1, keepdim=True) # (B, 1) # Update generated sequence generated = torch.cat([generated, next_token], dim=1) # Check for completion if end_token_id is not None: is_end = next_token.squeeze(-1) == end_token_id finished = finished | is_end if finished.all() and generated.size(1) >= max(1, min_len): break return generated # ----------------------------- # Incremental single-step API # ----------------------------- def step( self, last_token_ids: torch.Tensor, memory: torch.Tensor, cache: Optional[Dict] = None, ) -> Tuple[torch.Tensor, Dict]: """ Run one autoregressive step. Args: last_token_ids: (B, 1) last generated token ids memory: encoder outputs (B, S, d_model) cache: optional dict with previous cached keys/values and 'past_length'. Returns: logits: (B, vocab_size) logits for the next-token prediction new_cache: updated cache dictionary """ device = memory.device B = last_token_ids.size(0) if cache is None: cache = {} past_len = int(cache.get("past_length", 0)) # 1) Embed last token and add positional encoding for position `past_len` # T5/FLAN-T5 does NOT scale embeddings by sqrt(d_model) x = self.embedding(last_token_ids) # (B,1,d) # Handle positional encoding for single step # Note: When using relative position bias (T5-style), pos_encoder is None if self.pos_encoder is not None: if hasattr(self.pos_encoder, "pe"): # Sinusoidal: use buffer directly pe: torch.Tensor = self.pos_encoder.pe # type: ignore[union-attr] pos_idx = past_len if pos_idx >= pe.size(1): raise RuntimeError(f"pos_idx {pos_idx} exceeds max_len {pe.size(1)}") x = x + pe[:, pos_idx : pos_idx + 1, :].to(device) elif hasattr(self.pos_encoder, "embeddings"): # Learned: lookup specific position # Create position ids: [past_len] pos_idx_t = torch.tensor([past_len], dtype=torch.long, device=device) # Lookup embedding: (1, d_model) pos_emb = self.pos_encoder.embeddings(pos_idx_t) # type: ignore[union-attr] # Add to input: (B, 1, d_model) + (1, 1, d_model) broadcast x = x + pos_emb.unsqueeze(0) x = self.pos_encoder.dropout(x) # type: ignore[union-attr] else: # fallback: call pos_encoder (likely incorrect for step-by-step if it assumes pos 0) x = self.pos_encoder(x) # When pos_encoder is None (relative position bias mode), we skip positional encoding # The position information is provided via relative_position_bias in attention # We will update new_cache incrementally new_cache = dict(cache) # shallow copy new_cache["past_length"] = past_len + 1 # Optional: memory_mask could be supplied in cache under 'memory_mask' memory_mask = new_cache.get("memory_mask", None) if memory_mask is not None: memory_mask = memory_mask.to(dtype=torch.bool, device=device) # expand (B, S) -> (B,1,1,S) if necessary if memory_mask.dim() == 2: memory_mask = memory_mask.unsqueeze(1).unsqueeze(1) elif memory_mask.dim() == 3: memory_mask = memory_mask.unsqueeze(1) # Compute position biases for incremental step (T5-style) # For step mode: query_length=1, but actual position is past_len # Self-attention: query at position past_len attends to keys at positions 0..past_len # Note: T5 uses relative position bias for self-attention but NOT for cross-attention if self.use_relative_position_bias and self.self_relative_position_bias is not None: # Self-attention bias: query_length=1, key_length=past_len+1, offset=past_len self_position_bias = self.self_relative_position_bias( query_length=1, key_length=past_len + 1, device=device, query_position_offset=past_len, ) # (1, num_heads, 1, past_len+1) else: self_position_bias = None # Cross-attention position bias is None for T5 (see T5 paper/implementation) cross_position_bias = None # Iterate layers, updating caches and computing output for current token only layer_input = x # (B,1,d_model) for i, layer_module in enumerate(self.layers): layer = cast(TransformerDecoderLayer, layer_module) # ------------------- # 1) Self-attention (incremental) # ------------------- # Normalize input for pre-LN x_norm = layer.norm1(layer_input) # (B,1,d) # Project Q,K,V for the new token Q_new = layer.self_attn.W_Q(x_norm) # (B,1,d_model) K_new = layer.self_attn.W_K(x_norm) V_new = layer.self_attn.W_V(x_norm) # Reshape into heads: (B, num_heads, 1, d_k) B_, Lq, _ = Q_new.shape num_heads = layer.self_attn.num_heads d_k = layer.self_attn.d_k Qh = Q_new.view(B_, Lq, num_heads, d_k).transpose(1, 2) # (B, num_heads, 1, d_k) Kh = K_new.view(B_, Lq, num_heads, d_k).transpose(1, 2) Vh = V_new.view(B_, Lq, num_heads, d_k).transpose(1, 2) # Retrieve cached keys/values for self-attn (if exist) cache_k = cache.get(f"self_k_{i}", None) cache_v = cache.get(f"self_v_{i}", None) if cache_k is None or cache_v is None: K_all = Kh # (B, H, 1, d_k) V_all = Vh else: # concat along sequence dim (dim=2) K_all = torch.cat([cache_k.to(device), Kh], dim=2) V_all = torch.cat([cache_v.to(device), Vh], dim=2) # Store updated caches new_cache[f"self_k_{i}"] = K_all new_cache[f"self_v_{i}"] = V_all # Compute attention for the new token: Query length = 1, Key length = K_all.size(2) # Explicitly create mask for consistency with forward pass (though None should work) # mask=True means attend. step_mask = torch.ones(B_, 1, 1, K_all.size(2), dtype=torch.bool, device=device) attn_out_heads, self_attn_w = layer.self_attn.attention( Qh, K_all, V_all, mask=step_mask, position_bias=self_position_bias ) # attn_out_heads: (B, H, 1, d_k) # concat heads, project out attn_out = attn_out_heads.transpose(1, 2).contiguous().view(B_, 1, num_heads * d_k) attn_out = layer.self_attn.W_O(attn_out) # (B,1,d_model) attn_out = layer.self_attn.dropout(attn_out) layer_output = layer_input + layer.dropout1(attn_out) # ------------------- # 2) Cross-attention (use cached memory projections if available) # ------------------- x_norm2 = layer.norm2(layer_output) # (B,1,d) # Ensure memory K/V are cached per layer mem_k = cache.get(f"mem_k_{i}", None) mem_v = cache.get(f"mem_v_{i}", None) if mem_k is None or mem_v is None: # project memory once for this layer and cache it # memory: (B, S, d_model) MK = layer.cross_attn.W_K(memory) # (B, S, d_model) MV = layer.cross_attn.W_V(memory) Bm, S, _ = MK.shape MKh = MK.view(Bm, S, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose( 1, 2 ) # (B,H,S,d_k) MVh = MV.view(Bm, S, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose( 1, 2 ) mem_k = MKh mem_v = MVh new_cache[f"mem_k_{i}"] = mem_k new_cache[f"mem_v_{i}"] = mem_v else: mem_k = mem_k.to(device) mem_v = mem_v.to(device) Qc = layer.cross_attn.W_Q(x_norm2) # (B,1,d_model) Qch = Qc.view(B, 1, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose( 1, 2 ) # (B,H,1,d_k) cross_out_heads, cross_attn_w = layer.cross_attn.attention( Qch, mem_k, mem_v, mask=memory_mask, position_bias=cross_position_bias ) cross_out = ( cross_out_heads.transpose(1, 2) .contiguous() .view(B, 1, layer.cross_attn.num_heads * layer.cross_attn.d_k) ) cross_out = layer.cross_attn.W_O(cross_out) # (B,1,d_model) cross_out = layer.cross_attn.dropout(cross_out) layer_output = layer_output + layer.dropout2(cross_out) # ------------------- # 3) Feed-forward (incremental) # ------------------- x_norm3 = layer.norm3(layer_output) ffn_out = layer.ffn(x_norm3) # (B,1,d_model) layer_output = layer_output + layer.dropout3(ffn_out) # Prepare for next layer layer_input = layer_output # Final norm + output projection (for this single time step) out_norm = self.final_norm(layer_input) # (B,1,d_model) logits = self.output_projection(out_norm) # (B,1,vocab) logits = logits.squeeze(1) # (B, vocab) return logits, new_cache