gpt-o3 / cudaefa360

#include "kernel.h"
#include <torch/extension.h>

/*
 *  Python-visible front-end
 *
 *  python:  run(… tensors …,
 *               local_expert_offset : int,
 *               routed_scaling_factor : float) -> bf16 tensor
 */
torch::Tensor run(torch::Tensor routing_logits,
                  torch::Tensor routing_bias,
                  torch::Tensor hidden_states,
                  torch::Tensor hidden_states_scale,
                  torch::Tensor gemm1_weights,
                  torch::Tensor gemm1_weights_scale,
                  torch::Tensor gemm2_weights,
                  torch::Tensor gemm2_weights_scale,
                  int32_t       local_expert_offset,
                  double        routed_scaling_factor)
{
  TORCH_CHECK(routing_logits.is_cuda(), "all inputs must be CUDA tensors");

  const int T = hidden_states.size(0);
  auto out = torch::empty({T, HIDDEN_SIZE},
                          torch::TensorOptions()
                              .dtype(torch::kBFloat16)
                              .device(hidden_states.device()));

  moe_forward_cuda(routing_logits,
                   routing_bias,
                   hidden_states,
                   hidden_states_scale,
                   gemm1_weights,
                   gemm1_weights_scale,
                   gemm2_weights,
                   gemm2_weights_scale,
                   local_expert_offset,
                   static_cast<float>(routed_scaling_factor),
                   out);

  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
  m.def("run", &run,
        "MoE FP8 block-scale forward – B200-optimised");
}