Thank you for your advice @googlebot . I was able to achieve a significant speedup after incorporating your suggestions. Here is my current code in case it is of use to anyone. Is there anything else I could do for improving performance even more? I’m now gonna look into implementing backprop for the same.
// sparse_mul.cpp
#include <torch/extension.h>
#include <vector>

#define _OPENMP
#include <ATen/ParallelOpenMP.h>

template <typename scalar_t>
void mul2_kernel(
        torch::Tensor& out_values,
        const torch::Tensor& in_values,
        const torch::Tensor& active_out_indices,
        const torch::Tensor& weights,
        const torch::Tensor& bias){

    int64_t batch_size = out_values.size(0);
    int64_t active_out_dim = out_values.size(1);
    int64_t k_size = weights.size(1);

    auto out_values_0 = out_values.accessor<scalar_t, 2>();
    auto in_values_0 = in_values.accessor<scalar_t, 2>();
    auto active_out_indices_0 = active_out_indices.accessor<int64_t,2>();
    auto weights_0 = weights.accessor<scalar_t, 2>();
    auto bias_0 = bias.accessor<scalar_t, 1>();

    at::parallel_for(0, batch_size, 0, [&](int64_t start, int64_t end){
        for (int64_t i = start; i < end; i++){
            torch::set_num_threads(1); //is this required
            auto out_values_1 = out_values_0[i];
            auto in_values_1 = in_values_0[i];
            auto active_out_indices_1 = active_out_indices_0[i];
            
            for(int64_t j=0; j<active_out_dim; j++){
                scalar_t &res = out_values_1[j];
                int64_t out_index = active_out_indices_1[j];
                auto weights_1 = weights_0[out_index];

                res = bias_0[out_index];
                for(int64_t k=0; k<k_size; k++)
                    res += in_values_1[k]*weights_1[k];
            }
        }
        // std::cout << "in thread " << omp_get_thread_num() << std::endl;
    });
}

torch::Tensor mul2(
        const torch::Tensor& in_values,
        const torch::Tensor& active_out_indices,
        const torch::Tensor& weights,
        const torch::Tensor& bias){
    
    // std::cout<< "omp_get_max_threads() " << omp_get_max_threads() << std::endl;

    auto out_values = torch::empty(active_out_indices.sizes(), in_values.options());
    AT_DISPATCH_ALL_TYPES(
        in_values.scalar_type(), "mul2", [&] {
            mul2_kernel<scalar_t>(out_values, in_values, active_out_indices, weights, bias);
        });
    return out_values;
}