← Walkthroughmapa/models/attention.py
Source snapshot · source SHA-256 4de5ddf3d593
1"""Pre-norm transformer block with attention isolated by array and clip.23Main path: L1Block.forward_flat() -> _attn_flat().45Input/output: (M,d); neither token count nor width changes.6Flow: normalize -> QKV -> rotate Q/K -> attend within blocks -> project -> add.7Then normalize -> token-wise MLP -> add; return to the next encoder block.89Each token represents one contact, band, and time. Within its array/clip block,10it attends to all tokens across contacts, bands, and times, including itself.11There is no causal mask and no attention across block boundaries.1213RoPE rotates Q and K using clinical contact indices and shared-clock time.14For fixed Q/K content, its positional contribution depends on coordinate15differences. Region embeddings are added once before the encoder stack.16"""1718from __future__ import annotations1920from torch import Tensor, nn2122from mapa.models.pe import L1RoPE23from mapa.models.varlen import _reference_block_diag, gpu_block_diag2425LN_EPS = 1e-6262728class _MLP(nn.Module):29 def __init__(self, d_model: int, mlp_ratio: int) -> None:30 super().__init__()31 hidden = d_model * mlp_ratio32 self.fc1 = nn.Linear(d_model, hidden) # expand features: 384 -> 153633 self.fc2 = nn.Linear(hidden, d_model) # restore width: 1536 -> 38434 self.act = nn.GELU()3536 def forward(self, x: Tensor) -> Tensor:37 # Apply the same feature transformation independently to every token.38 # Shapes below use the released encoder width d=384.39 hidden = self.fc1(x) # (M,384) -> (M,1536)40 hidden = self.act(hidden) # GELU; shape unchanged41 output = self.fc2(hidden) # (M,1536) -> (M,384)42 return output434445class L1Block(nn.Module):46 """Within-array joint spatiotemporal attention block (block-diagonal, RoPE)."""4748 def __init__(49 self, d_model: int, n_heads: int, *, mlp_ratio: int = 4, space_rope: bool = True50 ) -> None:51 super().__init__()52 if d_model % n_heads != 0:53 raise ValueError(f"d_model={d_model} not divisible by n_heads={n_heads}")54 self.n_heads = n_heads55 self.head_dim = d_model // n_heads56 # Released model: d=384, six heads of width 64; no separate Q/K norm.57 self.qkv = nn.Linear(d_model, 3 * d_model, bias=True)58 self.out = nn.Linear(d_model, d_model, bias=True)59 self.rope = L1RoPE(self.head_dim, space=space_rope)60 self.norm1 = nn.LayerNorm(d_model, eps=LN_EPS)61 self.norm2 = nn.LayerNorm(d_model, eps=LN_EPS)62 self.mlp = _MLP(d_model, mlp_ratio)636465 # qkv, then the rotary on (depth, time_pos), then block-diagonal attention over66 # cu_seqlens, then the output projection. depth is the clinical contact index and67 # time_pos the position on the shared clock; both are per token.68 def _attn_flat(69 self,70 x_flat: Tensor,71 depth: Tensor,72 time_pos: Tensor,73 cu_seqlens: Tensor,74 cu_seqlens_drop: Tensor,75 max_seqlen: int,76 *,77 rope_cs: tuple[Tensor, Tensor] | None = None,78 ) -> Tensor:79 # x_flat: (M, d). depth, time_pos: (M,) long. Token order = array-contiguous80 # (cu_seqlens delimits each array block).81 m, d = x_flat.shape82 # 1. Project Q, K, V: (M,384) -> (M,1152) -> (M,3,6,64).83 qkv = self.qkv(x_flat) # (M,3*d): three learned projections in one linear layer84 qkv = qkv.reshape(m, 3, self.n_heads, self.head_dim)85 q, k, v = qkv.unbind(dim=1) # (M, H, hd)86 # 2. Rotate Q/K. RoPE expects heads before tokens; V is unchanged.87 qh, kh = q.transpose(0, 1), k.transpose(0, 1) # (H, M, hd)88 if rope_cs is None:89 qh, kh = self.rope(qh, kh, depth, time_pos)90 else:91 qh, kh = self.rope.rotate(qh, kh, rope_cs[0], rope_cs[1])92 q, k = qh.transpose(0, 1), kh.transpose(0, 1) # (M, H, hd)93 q, k = q.to(v.dtype), k.to(v.dtype) # align to v for SDPA (no-op under fp32)94 # 3. Attend within each array/clip block; output remains (M,H,hd).95 # Per head: softmax(Q @ K.T / sqrt(hd), over keys) @ V; see varlen.py.96 if q.is_cuda:97 ctx = gpu_block_diag(q, k, v, cu_seqlens_drop, max_seqlen) # (M, H, hd)98 else:99 ctx = _reference_block_diag(q, k, v, cu_seqlens) # (M, H, hd)100 # 4. Concatenate heads, then mix them with a learned d -> d projection.101 ctx = ctx.reshape(m, d) # (M,H,hd) -> (M,d)102 return self.out(ctx)103104 def forward_flat(105 self,106 x_flat: Tensor,107 depth: Tensor,108 time_pos: Tensor,109 cu_seqlens: Tensor,110 cu_seqlens_drop: Tensor,111 max_seqlen: int,112 *,113 rope_cs: tuple[Tensor, Tensor] | None = None,114 ) -> Tensor:115 # Block overview: keep x, learn an attention update, then learn an MLP update.116 # LayerNorm acts across features per token; both additions preserve (M,d).117 normalized = self.norm1(x_flat)118 attention_update = self._attn_flat(119 normalized, depth, time_pos, cu_seqlens, cu_seqlens_drop,120 max_seqlen, rope_cs=rope_cs,121 )122 x_flat = x_flat + attention_update123 # Normalize -> Linear(d,4d) -> GELU -> Linear(4d,d) -> add to x.124 normalized = self.norm2(x_flat)125 mlp_update = self.mlp(normalized)126 x_flat = x_flat + mlp_update127 return x_flat