[Misc] Fused MoE Marlin support for GPTQ (#8217)
This commit is contained in:
@@ -1737,4 +1737,4 @@ torch::Tensor marlin_gemm_moe(
|
||||
moe_block_size, dev, at::cuda::getCurrentCUDAStream(dev), thread_k,
|
||||
thread_n, sms, max_par, replicate_input, apply_weights);
|
||||
return c;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,4 +9,4 @@ torch::Tensor marlin_gemm_moe(
|
||||
const torch::Tensor& g_idx, const torch::Tensor& perm,
|
||||
torch::Tensor& workspace, int64_t size_m, int64_t size_n, int64_t size_k,
|
||||
bool is_k_full, int64_t num_experts, int64_t topk, int64_t moe_block_size,
|
||||
bool replicate_input, bool apply_weights);
|
||||
bool replicate_input, bool apply_weights);
|
||||
|
||||
@@ -16,7 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"g_idx, Tensor! perm, Tensor! workspace, int size_m, int size_n, int "
|
||||
"size_k, bool is_k_full, int num_experts, int topk, int moe_block_size, "
|
||||
"bool replicate_input, bool apply_weights) -> Tensor");
|
||||
|
||||
m.impl("marlin_gemm_moe", torch::kCUDA, &marlin_gemm_moe);
|
||||
#endif
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user