Embedding Operators

CUDA Operators

void bounds_check_indices_cuda(Tensor &rows_per_table, Tensor &indices, Tensor &offsets, int64_t bounds_check_mode, Tensor &warning, const c10::optional<Tensor> &weights, const c10::optional<Tensor> &B_ofsets, const int64_t max_B)

Tensor int_nbit_split_embedding_codegen_lookup_function(Tensor dev_weights, Tensor uvm_weights, Tensor weights_placements, Tensor weights_offsets, Tensor weights_tys, Tensor D_offsets, int64_t total_D, int64_t max_int2_D, int64_t max_int4_D, int64_t max_int8_D, int64_t max_float16_D, int64_t max_float32_D, Tensor indices, Tensor offsets, int64_t pooling_mode, c10::optional<Tensor> indice_weights, int64_t output_dtype, c10::optional<Tensor> lxu_cache_weights, c10::optional<Tensor> lxu_cache_locations, c10::optional<int64_t> row_alignment, c10::optional<int64_t> max_float8_D, c10::optional<int64_t> fp8_exponent_bits, c10::optional<int64_t> fp8_exponent_bias)

Tensor int_nbit_split_embedding_uvm_caching_codegen_lookup_function(Tensor dev_weights, Tensor uvm_weights, Tensor weights_placements, Tensor weights_offsets, Tensor weights_tys, Tensor D_offsets, int64_t total_D, int64_t max_int2_D, int64_t max_int4_D, int64_t max_int8_D, int64_t max_float16_D, int64_t max_float32_D, Tensor indices, Tensor offsets, int64_t pooling_mode, c10::optional<Tensor> indice_weights, int64_t output_dtype, c10::optional<Tensor> lxu_cache_weights, c10::optional<Tensor> lxu_cache_locations, c10::optional<int64_t> row_alignment, c10::optional<int64_t> max_float8_D, c10::optional<int64_t> fp8_exponent_bits, c10::optional<int64_t> fp8_exponent_bias, c10::optional<Tensor> cache_hash_size_cumsum, c10::optional<int64_t> total_cache_hash_size, c10::optional<Tensor> cache_index_table_map, c10::optional<Tensor> lxu_cache_state, c10::optional<Tensor> lxu_state)

Simlar to int_nbit_split_embedding_codegen_lookup_function, but it does UVM_CACHING lookup.

Tensor pruned_hashmap_lookup_cuda(Tensor indices, Tensor offsets, Tensor hash_table, Tensor hash_table_offsets)

Tensor pruned_array_lookup_cuda(Tensor indices, Tensor offsets, Tensor index_remappings, Tensor index_remappings_offsets)

CPU Operators

Tensor int_nbit_split_embedding_codegen_lookup_function_cpu(Tensor dev_weights, Tensor uvm_weights, Tensor weights_placements, Tensor weights_offsets, Tensor weights_tys, Tensor D_offsets, int64_t total_D, int64_t max_int2_D, int64_t max_int4_D, int64_t max_int8_D, int64_t max_float16_D, int64_t max_float32_D, Tensor indices, Tensor offsets, int64_t pooling_mode, c10::optional<Tensor> indice_weights, int64_t output_dtype, c10::optional<Tensor> lxu_cache_weights, c10::optional<Tensor> lxu_cache_locations, c10::optional<int64_t> row_alignment, c10::optional<int64_t> max_float8_D, c10::optional<int64_t> fp8_exponent_bits, c10::optional<int64_t> fp8_exponent_bias)

Tensor int_nbit_split_embedding_uvm_caching_codegen_lookup_function_cpu(Tensor dev_weights, Tensor uvm_weights, Tensor weights_placements, Tensor weights_offsets, Tensor weights_tys, Tensor D_offsets, int64_t total_D, int64_t max_int2_D, int64_t max_int4_D, int64_t max_int8_D, int64_t max_float16_D, int64_t max_float32_D, Tensor indices, Tensor offsets, int64_t pooling_mode, c10::optional<Tensor> indice_weights, int64_t output_dtype, c10::optional<Tensor> lxu_cache_weights, c10::optional<Tensor> lxu_cache_locations, c10::optional<int64_t> row_alignment, c10::optional<int64_t> max_float8_D, c10::optional<int64_t> fp8_exponent_bits, c10::optional<int64_t> fp8_exponent_bias, c10::optional<Tensor> cache_hash_size_cumsum, c10::optional<int64_t> total_cache_hash_size, c10::optional<Tensor> cache_index_table_map, c10::optional<Tensor> lxu_cache_state, c10::optional<Tensor> lxu_state)

void pruned_hashmap_insert_unweighted_cpu(Tensor indices, Tensor dense_indices, Tensor offsets, Tensor hash_table, Tensor hash_table_offsets)

Tensor pruned_hashmap_lookup_unweighted_cpu(Tensor indices, Tensor offsets, Tensor hash_table, Tensor hash_table_offsets)

Tensor pruned_array_lookup_cpu(Tensor indices, Tensor offsets, Tensor index_remappings, Tensor index_remappings_offsets)