| 591 | |
| 592 | |
| 593 | class SelectiveScanMatrix(torch.autograd.Function): |
| 594 | @staticmethod |
| 595 | @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) |
| 596 | def forward(ctx, us, dts, As, Bs, Cs, Ds, delta_bias=None, delta_softplus=False, return_last_state=False, chunksize=64): |
| 597 | save_for_backward = [] |
| 598 | if chunksize < 1: |
| 599 | chunksize = Bs.shape[-1] |
| 600 | mask = torch.tril(torch.ones((chunksize, chunksize), device=us.device), diagonal=0) |
| 601 | |
| 602 | def ss_chunk(us, dts, As, Bs, Cs, h0, mask): |
| 603 | # BHLD, BHLN, HND, BHDN |
| 604 | cL = us.shape[2] |
| 605 | _mask = (mask[:cL,:cL].contiguous() if cL < mask.shape[0] else mask).view(1, 1, cL, cL, 1) |
| 606 | |
| 607 | w_log = As[None, :, None, :, :] * (torch.cumsum(dts, dim=2)[..., None, :]) # (B, H, L, Dk, Dv) |
| 608 | v = us * dts # (B,H,L,Dv) |
| 609 | k = Bs # (B,H,L,Dk) |
| 610 | q = Cs # (B,H,L,Dk) |
| 611 | w = w_log.exp() |
| 612 | |
| 613 | k_div_w = k[..., None] / w |
| 614 | q_mul_w = q[..., None] * w |
| 615 | |
| 616 | # h0 independent ==================== |
| 617 | next_h_1 = w[:,:,-1] * torch.einsum("bhlkv,bhlv->bhkv", k_div_w, v) |
| 618 | y_1 = torch.einsum("bhlrv,bhrv->bhlv", torch.einsum("bhlkv,bhrkv->bhlrv", q_mul_w, k_div_w) * _mask, v) |
| 619 | |
| 620 | # h0 dependent ====================== |
| 621 | y_0 = torch.einsum("bhlkv,bhkv->bhlv", q_mul_w, h0) |
| 622 | next_h_0 = w[:,:, -1] * h0 |
| 623 | |
| 624 | next_h = next_h_0 + next_h_1 |
| 625 | y = y_0 + y_1 |
| 626 | |
| 627 | return y, next_h |
| 628 | |
| 629 | dtype = torch.float32 |
| 630 | # dtype = torch.float16 |
| 631 | inp_dtype = us.dtype |
| 632 | has_D = Ds is not None |
| 633 | dts = dts.to(dtype) |
| 634 | |
| 635 | if delta_bias is not None: |
| 636 | dts = dts + delta_bias.view(1, -1, 1).to(dtype) |
| 637 | if delta_softplus: |
| 638 | dts = torch.nn.functional.softplus(dts) |
| 639 | |
| 640 | if len(Bs.shape) == 3: |
| 641 | Bs = Bs.unsqueeze(1) |
| 642 | if len(Cs.shape) == 3: |
| 643 | Cs = Cs.unsqueeze(1) |
| 644 | |
| 645 | B, GD, L = us.shape |
| 646 | B, G, N, L = Bs.shape |
| 647 | D = GD // G |
| 648 | us = us.view(B, G, -1, L).permute(0, 1, 3, 2).to(dtype) |
| 649 | dts = dts.view(B, G, -1, L).permute(0, 1, 3, 2).to(dtype) |
| 650 | As = As.view(G, D, N).permute(0, 2, 1).to(dtype) |
nothing calls this directly
no outgoing calls
no test coverage detected