| 97 | |
| 98 | |
| 99 | class SelectiveScanEasy(torch.autograd.Function): |
| 100 | # for debug, we use it as an orinary object |
| 101 | DEBUG = (MODE == "fnDEBUG") |
| 102 | |
| 103 | if DEBUG: |
| 104 | print("DEBUG here...", flush=True) |
| 105 | saved_tensors = [] |
| 106 | |
| 107 | @classmethod |
| 108 | def save_for_backward(ctx, *args): |
| 109 | ctx.saved_tensors = args |
| 110 | |
| 111 | @staticmethod |
| 112 | @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) |
| 113 | def forward(ctx, us, dts, As, Bs, Cs, Ds, delta_bias=None, delta_softplus=False, return_last_state=False, chunksize=64): |
| 114 | has_D = Ds is not None |
| 115 | dtype = torch.float32 |
| 116 | |
| 117 | dts = dts.to(dtype) |
| 118 | if delta_bias is not None: |
| 119 | dts = dts + delta_bias.view(1, -1, 1).to(dtype) |
| 120 | if delta_softplus: |
| 121 | dts = torch.nn.functional.softplus(dts) |
| 122 | |
| 123 | B_squeeze = (len(Bs.shape) == 3) |
| 124 | C_squeeze = (len(Cs.shape) == 3) |
| 125 | if B_squeeze: |
| 126 | Bs = Bs.unsqueeze(1) |
| 127 | if C_squeeze: |
| 128 | Cs = Cs.unsqueeze(1) |
| 129 | B, G, N, L = Bs.shape |
| 130 | us = us.view(B, G, -1, L).permute(3, 0, 1, 2).to(dtype) |
| 131 | dts = dts.view(B, G, -1, L).permute(3, 0, 1, 2).to(dtype) |
| 132 | As = As.view(G, -1, N).to(dtype) |
| 133 | Bs = Bs.permute(3, 0, 1, 2).to(dtype) |
| 134 | Cs = Cs.permute(3, 0, 1, 2).to(dtype) |
| 135 | Ds = Ds.view(G, -1).to(dtype) if has_D else None |
| 136 | D = As.shape[1] |
| 137 | |
| 138 | ctx.shape = (B, G, D, N, L) |
| 139 | ctx.delta_softplus = delta_softplus |
| 140 | ctx.return_last_state = return_last_state |
| 141 | ctx.chunksize = chunksize |
| 142 | ctx.BC_squeeze = (B_squeeze, C_squeeze) |
| 143 | save_for_backward = [us, dts, As, Bs, Cs, Ds, delta_bias] |
| 144 | |
| 145 | chunks = list(range(0, L, chunksize)) |
| 146 | oys = [] |
| 147 | ohs = [] |
| 148 | hprefix = us.new_zeros((B, G, D, N), dtype=torch.float) |
| 149 | for i in chunks: |
| 150 | ts = dts[i:i+chunksize].cumsum(dim=0) |
| 151 | Ats = torch.einsum("gdn,lbgd->lbgdn", As, ts).exp() |
| 152 | # scale = Ats[-1:].detach() |
| 153 | scale = 1 |
| 154 | rAts = Ats / scale |
| 155 | duts = dts[i:i + chunksize] * us[i:i + chunksize] |
| 156 | dtBus = torch.einsum("lbgd,lbgn->lbgdn", duts, Bs[i:i + chunksize]) |
nothing calls this directly
no outgoing calls
no test coverage detected