AutoAWQ_kernels/awq_ext/vllm/topk_softmax_kernels.h

#pragma once

#include <torch/extension.h>

void topk_softmax(
  torch::Tensor& topk_weights,
  torch::Tensor& topk_indices,
  torch::Tensor& token_expert_indices,
  torch::Tensor& gating_output);