| 401 | """ |
| 402 | |
| 403 | def step(self, x, cache): |
| 404 | # x : (B, D) |
| 405 | # cache : (h, inputs) |
| 406 | # h : (B, ED, N) |
| 407 | # inputs : (B, ED, d_conv-1) |
| 408 | |
| 409 | # y : (B, D) |
| 410 | # cache : (h, inputs) |
| 411 | |
| 412 | h, inputs = cache |
| 413 | |
| 414 | xz = self.in_proj(x) # (B, 2*ED) |
| 415 | x, z = xz.chunk(2, dim=1) # (B, ED), (B, ED) |
| 416 | |
| 417 | # x branch |
| 418 | x_cache = x.unsqueeze(2) |
| 419 | x = self.conv1d(torch.cat([inputs, x_cache], dim=2))[ |
| 420 | :, :, self.config.d_conv - 1 |
| 421 | ] # (B, ED) |
| 422 | |
| 423 | x = F.silu(x) |
| 424 | y, h = self.ssm_step(x, h) |
| 425 | |
| 426 | # z branch |
| 427 | z = F.silu(z) |
| 428 | |
| 429 | output = y * z |
| 430 | output = self.out_proj(output) # (B, D) |
| 431 | |
| 432 | # prepare cache for next call |
| 433 | inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) # (B, ED, d_conv-1) |
| 434 | cache = (h, inputs) |
| 435 | |
| 436 | return output, cache |
| 437 | |
| 438 | def ssm_step(self, x, h): |
| 439 | # x : (B, ED) |