Spaces:
Paused
Paused
File size: 7,853 Bytes
e319a7b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """Causal streaming transformer for muscriptor inference."""
from einops import rearrange
import torch
import torch.nn as nn
from torch.nn import functional as F
from muscriptor.modules.streaming import ModelState, State, StatefulModule
def create_sin_embedding(
positions: torch.Tensor,
dim: int,
max_period: float = 10000,
dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
assert dim % 2 == 0
half_dim = dim // 2
positions = positions.to(dtype)
adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)
phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
class StreamingMultiheadAttention(StatefulModule):
"""Causal multi-head self-attention with a preallocated KV cache."""
def __init__(
self,
embed_dim: int,
num_heads: int,
device=None,
dtype=None,
):
super().__init__()
factory_kwargs = {"device": device, "dtype": dtype}
self.embed_dim = embed_dim
self.num_heads = num_heads
self.dim_per_head = embed_dim // num_heads
in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **factory_kwargs)
self.in_proj_weight = in_proj.weight
self.in_proj_bias = in_proj.bias
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs)
def init_state(self, batch_size: int, sequence_length: int) -> State:
weight = self.in_proj_weight
return {
"cache": torch.full(
(2, batch_size, sequence_length, self.num_heads, self.dim_per_head),
float("nan"),
device=weight.device,
dtype=weight.dtype,
),
# Kept as a plain host int: it is advanced deterministically by the
# host-side generate loop, and reading it from a device tensor
# (`.item()`) would force a GPU sync per layer per decode step.
"offset": 0,
}
def increment_step(self, state: State, increment: int = 1) -> None:
state["offset"] = state["offset"] + increment
def _complete_kv(self, k, v, state: State | None):
if state is None:
return k, v
cache = state["cache"]
end = state["offset"]
T = k.shape[1]
cache[0, :, end : end + T] = k
cache[1, :, end : end + T] = v
return cache[0, :, : end + T], cache[1, :, : end + T]
def forward(
self,
query: torch.Tensor,
model_state: ModelState | None = None,
):
state = self.get_state(model_state)
projected = nn.functional.linear(query, self.in_proj_weight)
packed = rearrange(projected, "b t (p h d) -> b t p h d", p=3, h=self.num_heads)
q, k, v = packed.unbind(dim=2)
k, v = self._complete_kv(k, v, state)
dtype = q.dtype
q_t = q.transpose(1, 2)
k_t = k.transpose(1, 2)
v_t = v.transpose(1, 2)
# Causality must be bottom-right aligned so streaming decode steps
# (T_q=1, T_k=cache_len) attend to all past tokens; PyTorch's
# is_causal=True is top-left aligned and would mask out all cached
# tokens except position 0 when T_q < T_k. An explicit attn_mask
# forces SDPA onto the unfused math fallback, so only build one in
# the rectangular case that actually needs it — the two shapes this
# model hits (single-token decode and square prefill) stay mask-free
# and dispatch to the fused (flash) CPU/CUDA kernels.
T_q, T_k = q_t.shape[2], k_t.shape[2]
if T_q == 1:
# One query row, bottom-right aligned: nothing is masked.
x = F.scaled_dot_product_attention(q_t, k_t, v_t, dropout_p=0.0)
elif T_q == T_k:
# Square: bottom-right and top-left alignment coincide.
x = F.scaled_dot_product_attention(
q_t, k_t, v_t, is_causal=True, dropout_p=0.0
)
else:
# Unused in practice
raise NotImplementedError(
f"Streaming attention with T_q={T_q} and T_k={T_k} is not supported; use T_q=1 or T_q=T_k."
)
x = x.transpose(1, 2).to(dtype)
x = rearrange(x, "b t h d -> b t (h d)")
x = self.out_proj(x)
return x
class StreamingTransformerLayer(nn.Module):
"""Pre-norm transformer block: self-attention + GELU FFN."""
def __init__(
self,
d_model: int,
num_heads: int,
dim_feedforward: int = 2048,
device=None,
dtype=None,
):
super().__init__()
factory_kwargs = {"device": device, "dtype": dtype}
self.self_attn = StreamingMultiheadAttention(
embed_dim=d_model, num_heads=num_heads, **factory_kwargs
)
self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)
self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)
self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False, **factory_kwargs)
self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False, **factory_kwargs)
def forward(
self,
x: torch.Tensor,
model_state: ModelState | None = None,
):
x = x + self.self_attn(self.norm1(x), model_state=model_state)
x = x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
return x
class StreamingTransformer(StatefulModule):
"""Stack of causal streaming transformer layers with sinusoidal positions."""
def __init__(
self,
d_model: int,
num_heads: int,
num_layers: int,
dim_feedforward: int = 2048,
max_period: float = 10_000,
device=None,
dtype=None,
):
super().__init__()
assert d_model % num_heads == 0
self.max_period = max_period
self.layers = nn.ModuleList(
[
StreamingTransformerLayer(
d_model=d_model,
num_heads=num_heads,
dim_feedforward=dim_feedforward,
device=device,
dtype=dtype,
)
for _ in range(num_layers)
]
)
def init_state(self, batch_size: int, sequence_length: int) -> State:
device = self.layers[0].norm2.weight.device
return {
"offsets": torch.zeros(batch_size, dtype=torch.long, device=device),
}
def increment_step(self, state: State, increment: int = 1) -> None:
state["offsets"] = state["offsets"] + increment
def forward(
self,
x: torch.Tensor,
prepend_length: int = 0,
model_state: ModelState | None = None,
):
del prepend_length # unused; positions come from state['offsets']
B, T, C = x.shape
state = self.get_state(model_state)
offsets = (
state["offsets"]
if state is not None
else torch.zeros(B, dtype=torch.long, device=x.device)
)
positions = torch.arange(T, device=x.device).view(1, -1, 1)
positions = positions + offsets.view(-1, 1, 1)
# Always compute the sinusoidal embedding in fp32: fp16 cannot even
# represent odd integers above 2048, so half-precision positions would
# collapse neighbouring timesteps to the same embedding.
pos_emb = create_sin_embedding(
positions, C, max_period=self.max_period, dtype=torch.float32
)
x = x + (pos_emb * (positions >= 0).float()).to(x.dtype)
for layer in self.layers:
x = layer(x, model_state=model_state)
return x
|