wehub-resource-sync
78 lines
3.0 KiB
Plaintext
78 lines
3.0 KiB
Plaintext
#include "machete_mm_launcher.cuh"
|
|
#include "machete_prepack_launcher.cuh"
|
|
#include "core/scalar_type.hpp"
|
|
|
|
#include <torch/csrc/stable/library.h>
|
|
#include <torch/csrc/stable/tensor.h>
|
|
#include <torch/headeronly/core/ScalarType.h>
|
|
|
|
namespace machete {
|
|
|
|
using namespace vllm;
|
|
|
|
std::vector<std::string> supported_schedules(
|
|
torch::headeronly::ScalarType a_type, int64_t b_type_id,
|
|
std::optional<torch::headeronly::ScalarType> maybe_group_scales_type,
|
|
std::optional<torch::headeronly::ScalarType> maybe_group_zeros_type,
|
|
std::optional<torch::headeronly::ScalarType> maybe_channel_scales_type,
|
|
std::optional<torch::headeronly::ScalarType> maybe_token_scales_type,
|
|
std::optional<torch::headeronly::ScalarType> maybe_out_type) {
|
|
ScalarType const b_type = ScalarType::from_id(b_type_id);
|
|
return supported_schedules_dispatch({
|
|
.a_type = a_type,
|
|
.b_type = b_type,
|
|
.maybe_group_scales_type = maybe_group_scales_type,
|
|
.maybe_group_zeros_type = maybe_group_zeros_type,
|
|
.maybe_channel_scales_type = maybe_channel_scales_type,
|
|
.maybe_token_scales_type = maybe_token_scales_type,
|
|
.maybe_out_type = maybe_out_type,
|
|
});
|
|
}
|
|
|
|
torch::stable::Tensor mm(
|
|
torch::stable::Tensor const& A, torch::stable::Tensor const& B,
|
|
int64_t b_type_id,
|
|
std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
|
|
std::optional<torch::stable::Tensor> const& maybe_group_scales,
|
|
std::optional<torch::stable::Tensor> const& maybe_group_zeros,
|
|
std::optional<int64_t> maybe_group_size,
|
|
std::optional<torch::stable::Tensor> const& maybe_channel_scales,
|
|
std::optional<torch::stable::Tensor> const& maybe_token_scales,
|
|
std::optional<std::string> maybe_schedule) {
|
|
ScalarType const b_type = ScalarType::from_id(b_type_id);
|
|
return mm_dispatch({.A = A,
|
|
.B = B,
|
|
.b_type = b_type,
|
|
.maybe_out_type = maybe_out_type,
|
|
.maybe_group_scales = maybe_group_scales,
|
|
.maybe_group_zeros = maybe_group_zeros,
|
|
.maybe_group_size = maybe_group_size,
|
|
.maybe_channel_scales = maybe_channel_scales,
|
|
.maybe_token_scales = maybe_token_scales,
|
|
.maybe_schedule = maybe_schedule});
|
|
}
|
|
|
|
torch::stable::Tensor prepack_B(
|
|
torch::stable::Tensor const& B, torch::headeronly::ScalarType const& a_type,
|
|
int64_t b_type_id,
|
|
std::optional<torch::headeronly::ScalarType> const&
|
|
maybe_group_scales_type) {
|
|
ScalarType const b_type = ScalarType::from_id(b_type_id);
|
|
return prepack_B_dispatch(
|
|
{.B = B,
|
|
.a_type = a_type,
|
|
.b_type = b_type,
|
|
.maybe_group_scales_type = maybe_group_scales_type});
|
|
}
|
|
|
|
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
|
m.impl("machete_prepack_B", TORCH_BOX(&prepack_B));
|
|
m.impl("machete_mm", TORCH_BOX(&mm));
|
|
}
|
|
|
|
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
|
|
m.impl("machete_supported_schedules", TORCH_BOX(&supported_schedules));
|
|
}
|
|
|
|
}; // namespace machete
|