moe_sum.h

#pragma once
// SPDX-License-Identifier: MIT
 
#include <torch/extension.h>
#include "aiter_enum.h"

void asm_moe_sum(torch::Tensor& input,                   // [experts, block_size, hidden_size]
                 torch::Tensor& output,                  // [num_tokens, hidden_size]
                 torch::Tensor& sorted_ids);