(ctx, expt_scal, expt_indx, bitmatrix, n_expts_tot, simulated_ep)
| 174 | |
| 175 | @staticmethod |
| 176 | def forward(ctx, expt_scal, expt_indx, bitmatrix, n_expts_tot, simulated_ep): |
| 177 | from .compaction import compaction |
| 178 | n_tokens_pad = expt_scal.shape[0] |
| 179 | assert n_expts_tot % simulated_ep == 0 |
| 180 | _routing_clear_bitmatrix[(n_tokens_pad, )]( |
| 181 | bitmatrix.storage.data, |
| 182 | bitmatrix.storage.data.stride(0), |
| 183 | bitmatrix.storage.data.stride(1), |
| 184 | bitmatrix.storage.data.shape[1], |
| 185 | n_expts_tot // simulated_ep, |
| 186 | BLOCK_N=512, |
| 187 | ) |
| 188 | # perform compaction to update expt_scal / expt_indx |
| 189 | expt_scal, expt_indx = compaction(expt_scal, expt_indx, bitmatrix) |
| 190 | n_expts_tot = n_expts_tot // simulated_ep |
| 191 | bitmatrix.shape[-1] = n_expts_tot |
| 192 | return expt_scal, expt_indx, bitmatrix |
| 193 | |
| 194 | |
| 195 | def prune_routing(expt_scal, expt_indx, bitmatrix, n_expts_tot, simulated_ep): |
nothing calls this directly
no test coverage detected