biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel] Enable 8-bit weights in Fused Marlin MoE (#8032)

Browse Source 
Co-authored-by: Dipika <dipikasikka1@gmail.com>
This commit is contained in:
ElizaWszola
2024-09-16 17:47:19 +02:00
committed by GitHub
parent fc990f9795
commit a091e2da3e
12 changed files with 452 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
csrc/moe/marlin_moe_ops.h
View File

@@ -2,11 +2,14 @@
#include <torch/all.h>
#include "core/scalar_type.hpp"
torch::Tensor marlin_gemm_moe(
const torch::Tensor& a, const torch::Tensor& b_q_weights,
const torch::Tensor& sorted_ids, const torch::Tensor& topk_weights,
const torch::Tensor& topk_ids, const torch::Tensor& b_scales,
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,
torch::Tensor& workspace, vllm::ScalarTypeTorchPtr const& b_q_type,
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);