u: r(B D L) delta: r(B D L) A: r(D N) B: r(B N L) C: r(B N L) D: r(D) z: r(B D L) delta_bias: r(D), fp32 ignores: [.float(), +, .softplus, .shape, new_zeros, repeat, stack, to(dtype), silu]
(
B=1,
L=256,
D=768,
N=16,
with_D=True,
with_Z=False,
with_Group=True,
with_complex=False,
)
| 1138 | |
| 1139 | |
| 1140 | def flops_selective_scan_fn( |
| 1141 | B=1, |
| 1142 | L=256, |
| 1143 | D=768, |
| 1144 | N=16, |
| 1145 | with_D=True, |
| 1146 | with_Z=False, |
| 1147 | with_Group=True, |
| 1148 | with_complex=False, |
| 1149 | ): |
| 1150 | """ |
| 1151 | u: r(B D L) |
| 1152 | delta: r(B D L) |
| 1153 | A: r(D N) |
| 1154 | B: r(B N L) |
| 1155 | C: r(B N L) |
| 1156 | D: r(D) |
| 1157 | z: r(B D L) |
| 1158 | delta_bias: r(D), fp32 |
| 1159 | |
| 1160 | ignores: |
| 1161 | [.float(), +, .softplus, .shape, new_zeros, repeat, stack, to(dtype), silu] |
| 1162 | """ |
| 1163 | assert not with_complex |
| 1164 | # https://github.com/state-spaces/mamba/issues/110 |
| 1165 | flops = 9 * B * L * D * N |
| 1166 | if with_D: |
| 1167 | flops += B * D * L |
| 1168 | if with_Z: |
| 1169 | flops += B * D * L |
| 1170 | return flops |
| 1171 | |
| 1172 | |
| 1173 | def selective_scan_flop_jit(inputs, outputs): |
no outgoing calls
no test coverage detected