| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | from einops import rearrange, repeat |
| 7 | |
| 8 | |
| 9 | def selective_state_update_ref( |
| 10 | state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False |
| 11 | ): |
| 12 | """Argument: |
| 13 | state: (batch, dim, dstate) or (batch, nheads, dim, dstate) |
| 14 | x: (batch, dim) or (batch, nheads, dim) |
| 15 | dt: (batch, dim) or (batch, nheads, dim) |
| 16 | A: (dim, dstate) or (nheads, dim, dstate) |
| 17 | B: (batch, dstate) or (batch, ngroups, dstate) |
| 18 | C: (batch, dstate) or (batch, ngroups, dstate) |
| 19 | D: (dim,) or (nheads, dim) |
| 20 | z: (batch, dim) or (batch, nheads, dim) |
| 21 | dt_bias: (dim,) or (nheads, dim) |
| 22 | |
| 23 | Return: |
| 24 | out: (batch, dim) or (batch, nheads, dim) |
| 25 | |
| 26 | """ |
| 27 | has_heads = state.dim() > 3 |
| 28 | if state.dim() == 3: |
| 29 | state = state.unsqueeze(1) |
| 30 | if x.dim() == 2: |
| 31 | x = x.unsqueeze(1) |
| 32 | if dt.dim() == 2: |
| 33 | dt = dt.unsqueeze(1) |
| 34 | if A.dim() == 2: |
| 35 | A = A.unsqueeze(0) |
| 36 | if B.dim() == 2: |
| 37 | B = B.unsqueeze(1) |
| 38 | if C.dim() == 2: |
| 39 | C = C.unsqueeze(1) |
| 40 | if D is not None and D.dim() == 1: |
| 41 | D = D.unsqueeze(0) |
| 42 | if z is not None and z.dim() == 2: |
| 43 | z = z.unsqueeze(1) |
| 44 | if dt_bias is not None and dt_bias.dim() == 1: |
| 45 | dt_bias = dt_bias.unsqueeze(0) |
| 46 | batch, nheads, dim, dstate = state.shape |
| 47 | assert x.shape == (batch, nheads, dim) |
| 48 | assert dt.shape == x.shape |
| 49 | assert A.shape == (nheads, dim, dstate) |
| 50 | ngroups = B.shape[1] |
| 51 | assert nheads % ngroups == 0, "nheads must be divisible by ngroups" |
| 52 | assert B.shape == (batch, ngroups, dstate) |
| 53 | assert C.shape == B.shape |
| 54 | if D is not None: |
| 55 | assert D.shape == (nheads, dim) |
| 56 | if z is not None: |
| 57 | assert z.shape == x.shape |
| 58 | if dt_bias is not None: |
| 59 | assert dt_bias.shape == (nheads, dim) |
| 60 | dt = dt + dt_bias |
| 61 | dt = F.softplus(dt) if dt_softplus else dt |
| 62 | dA = torch.exp( |
| 63 | rearrange(dt, "b h d -> b h d 1") * A |
| 64 | ) # (batch, nheads, dim, dstate) |
| 65 | B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) |
| 66 | C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) |
| 67 | dB = rearrange(dt, "b h d -> b h d 1") * rearrange( |
| 68 | B, "b h n -> b h 1 n" |
| 69 | ) # (batch, nheads, dim, dstate) |
| 70 | state.copy_( |
| 71 | state * dA + dB * rearrange(x, "b h d -> b h d 1") |
| 72 | ) # (batch, dim, dstate |
| 73 | out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C) |
| 74 | if D is not None: |
| 75 | out += (x * D).to(out.dtype) |
| 76 | out = (out if z is None else out * F.silu(z)).to(x.dtype) |
| 77 | if not has_heads: |
| 78 | out = out.squeeze(1) |
| 79 | return out |
| 80 | |
| 81 | |
| 82 | def selective_state_update_replayssm_output_only_ref( |
| 83 | state: torch.Tensor, |
| 84 | x: torch.Tensor, |
| 85 | dt: torch.Tensor, |
| 86 | A: torch.Tensor, |
| 87 | B: torch.Tensor, |
| 88 | C: torch.Tensor, |
| 89 | D: torch.Tensor | None = None, |
| 90 | z: torch.Tensor | None = None, |
| 91 | dt_bias: torch.Tensor | None = None, |
| 92 | dt_softplus: bool = False, |
| 93 | x_cache: torch.Tensor | None = None, |
| 94 | dt_cache: torch.Tensor | None = None, |
| 95 | B_cache: torch.Tensor | None = None, |
| 96 | write_pos: torch.Tensor | None = None, |
| 97 | max_cache_len: int = 16, |
| 98 | ) -> torch.Tensor: |
| 99 | """Pure-PyTorch cached-bc reference for validation.""" |
| 100 | has_heads = state.dim() > 3 |
| 101 | if state.dim() == 3: |
| 102 | state = state.unsqueeze(1) |
| 103 | if x.dim() == 2: |
| 104 | x = x.unsqueeze(1) |
| 105 | if dt.dim() == 2: |
| 106 | dt = dt.unsqueeze(1) |
| 107 | if A.dim() == 2: |
| 108 | A = A.unsqueeze(0) |
| 109 | if B.dim() == 2: |
| 110 | B = B.unsqueeze(1) |
| 111 | if C.dim() == 2: |
| 112 | C = C.unsqueeze(1) |
| 113 | if D is not None and D.dim() == 1: |
| 114 | D = D.unsqueeze(0) |
| 115 | if z is not None and z.dim() == 2: |
| 116 | z = z.unsqueeze(1) |
| 117 | if dt_bias is not None and dt_bias.dim() == 1: |
| 118 | dt_bias = dt_bias.unsqueeze(0) |
| 119 | |
| 120 | batch, nheads, dim, dstate = state.shape |
| 121 | assert x.shape == (batch, nheads, dim) |
| 122 | assert dt.shape == x.shape |
| 123 | assert A.shape == (nheads, dim, dstate) |
| 124 | ngroups = B.shape[1] |
| 125 | assert nheads % ngroups == 0, "nheads must be divisible by ngroups" |
| 126 | assert B.shape == (batch, ngroups, dstate) |
| 127 | assert C.shape == B.shape |
| 128 | |
| 129 | ratio = nheads // ngroups |
| 130 | |
| 131 | dt_val = dt[:, :, 0].float() |
| 132 | if dt_bias is not None: |
| 133 | dt_val = dt_val + dt_bias[:, 0].float() |
| 134 | if dt_softplus: |
| 135 | dt_val = F.softplus(dt_val) |
| 136 | A_val = A[:, 0, 0].float() |
| 137 | C_heads = C.repeat_interleave(ratio, dim=1) |
| 138 | out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32) |
| 139 | |
| 140 | assert x_cache is not None |
| 141 | assert dt_cache is not None |
| 142 | assert B_cache is not None |
| 143 | assert write_pos is not None |
| 144 | |
| 145 | for b in range(batch): |
| 146 | cache_len = int(write_pos[b].item()) |
| 147 | is_flush = cache_len == max_cache_len - 1 |
| 148 | n_steps = cache_len + 1 |
| 149 | |
| 150 | dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32) |
| 151 | if cache_len > 0: |
| 152 | dt_all[:, :cache_len] = dt_cache[b, :, :cache_len] |
| 153 | dt_all[:, cache_len] = dt_val[b] |
| 154 | |
| 155 | cumsum = torch.cumsum(dt_all, dim=-1) |
| 156 | total = cumsum[:, -1] |
| 157 | dA_cumsum = A_val[:, None] * cumsum |
| 158 | dA_total = A_val * total |
| 159 | total_decay = torch.exp(dA_total) |
| 160 | scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum) |
| 161 | |
| 162 | x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype) |
| 163 | if cache_len > 0: |
| 164 | x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1) |
| 165 | x_all[..., cache_len] = x[b] |
| 166 | |
| 167 | B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype) |
| 168 | if cache_len > 0: |
| 169 | B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :] |
| 170 | B_all[:, cache_len, :] = B[b] |
| 171 | |
| 172 | B_heads = B_all.repeat_interleave(ratio, dim=0) |
| 173 | C_heads_b = C_heads[b] |
| 174 | |
| 175 | if is_flush: |
| 176 | B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype) |
| 177 | delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float()) |
| 178 | state_new = state[b].float() * total_decay[:, None, None] + delta |
| 179 | state[b].copy_(state_new.to(state.dtype)) |
| 180 | out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads_b.float()) |
| 181 | else: |
| 182 | checkpoint_out = torch.einsum( |
| 183 | "hdn,hn->hd", state[b].float(), C_heads_b.float() |
| 184 | ) |
| 185 | checkpoint_out = checkpoint_out * total_decay[:, None] |
| 186 | BC = torch.einsum("hkn,hn->hk", B_heads.float(), C_heads_b.float()) |
| 187 | cache_out = torch.einsum("hdk,hk->hd", x_all.float(), scale * BC) |
| 188 | out[b] = checkpoint_out + cache_out |
| 189 | x_cache[b, :, cache_len, :] = x[b] |
| 190 | dt_cache[b, :, cache_len] = dt_val[b] |
| 191 | B_cache[b, :, cache_len, :] = B[b] |
| 192 | |
| 193 | if D is not None: |
| 194 | out = out + (x.float() * D[None]).to(out.dtype) |
| 195 | if z is not None: |
| 196 | out = out * F.silu(z.float()) |
| 197 | out = out.to(x.dtype) |
| 198 | if not has_heads: |
| 199 | out = out.squeeze(1) |
| 200 | return out |
| 201 | |
| 202 | |
| 203 | def allocate_update_caches( |
| 204 | batch: int, |
| 205 | nheads: int, |
| 206 | ngroups: int, |
| 207 | dim: int, |
| 208 | dstate: int, |
| 209 | max_cache_len: int, |
| 210 | device: torch.device, |
| 211 | x_dtype: torch.dtype, |
| 212 | B_dtype: torch.dtype, |
| 213 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 214 | """Allocate dense reference caches for standalone validation.""" |
| 215 | x_cache = torch.zeros( |
| 216 | batch, |
| 217 | nheads, |
| 218 | max_cache_len, |
| 219 | dim, |
| 220 | device=device, |
| 221 | dtype=x_dtype, |
| 222 | ) |
| 223 | dt_cache = torch.zeros( |
| 224 | batch, |
| 225 | nheads, |
| 226 | max_cache_len, |
| 227 | device=device, |
| 228 | dtype=torch.float32, |
| 229 | ) |
| 230 | B_cache = torch.zeros( |
| 231 | batch, |
| 232 | ngroups, |
| 233 | max_cache_len, |
| 234 | dstate, |
| 235 | device=device, |
| 236 | dtype=B_dtype, |
| 237 | ) |
| 238 | write_pos = torch.zeros(batch, dtype=torch.int32, device=device) |
| 239 | return x_cache, dt_cache, B_cache, write_pos |