permutation.h

Functions

void nvte_permute(const NVTETensor input, NVTETensor output, const NVTETensor sorted_row_id, NVTETensor row_id_map, const NVTETensor prob, NVTETensor prob_grad, const NVTETensor input_fwd, const int num_rows, const int topK, const int num_cols, const int num_out_tokens, cudaStream_t stream = nullptr)

void nvte_unpermute(const NVTETensor input, NVTETensor output, NVTETensor row_id_map, const NVTETensor prob, const int num_rows, const int topK, const int num_cols, cudaStream_t stream = nullptr)

void nvte_device_radix_sort_pairs(void *temp_storage, size_t *temp_storage_bytes, int *keys_in, int *keys_out, int *values_in, int *values_out, size_t num_items)