A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[Quantization][MoE] remove unused ep logic from moe marlin (#31571)
Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Jinzhen Lin committed
Jan 7, 2026 at 01:07 UTC
2f4bdee61ee0dd9358efaba720b7acc53b2ece00
6 files changed
+31
-60
csrc/moe/marlin_moe_wna16/kernel.h
+14
-14
@@ -7,20 +7,20 @@
7
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
-#define MARLIN_KERNEL_PARAMS \
11
- const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
- int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
- const int4 *__restrict__ b_bias_ptr, \
14
- const float *__restrict__ a_scales_ptr, \
15
- const int4 *__restrict__ scales_ptr, \
16
- const uint16_t *__restrict__ global_scale_ptr, \
17
- const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
- const int32_t *__restrict__ sorted_token_ids_ptr, \
19
- const int32_t *__restrict__ expert_ids_ptr, \
20
- const int32_t *__restrict__ num_tokens_past_padded_ptr, \
21
- const float *__restrict__ topk_weights_ptr, int top_k, \
22
- bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
23
- int prob_n, int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ b_bias_ptr, \
14
+ const float *__restrict__ a_scales_ptr, \
15
+ const int4 *__restrict__ scales_ptr, \
16
+ const uint16_t *__restrict__ global_scale_ptr, \
17
+ const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
+ const int32_t *__restrict__ sorted_token_ids_ptr, \
19
+ const int32_t *__restrict__ expert_ids_ptr, \
20
+ const int32_t *__restrict__ num_tokens_past_padded_ptr, \
21
+ const float *__restrict__ topk_weights_ptr, int top_k, \
22
+ bool mul_topk_weights, int num_groups, int prob_m, int prob_n, \
23
+ int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
24
bool use_fp32_reduce
25
26
namespace MARLIN_NAMESPACE_NAME {
csrc/moe/marlin_moe_wna16/marlin_template.h
+2
-26
@@ -71,7 +71,6 @@ __global__ void Marlin(
71
const float* __restrict__ topk_weights_ptr, // moe top weights
72
int top_k, // num of experts per token
73
bool mul_topk_weights, // mul topk weights or not
74
- bool is_ep, // expert parallelism
74
int num_groups, // number of scale groups per output channel
75
int prob_m, // batch dimension m
76
int prob_n, // output dimension n
@@ -273,7 +272,6 @@ __global__ void Marlin(
272
const float* __restrict__ topk_weights_ptr, // moe top weights
273
int top_k, // num of experts per token
274
bool mul_topk_weights, // mul topk weights or not
276
- bool is_ep, // expert parallelism
275
int num_groups, // number of scale groups per output channel
276
int prob_m, // batch dimension m
277
int prob_n, // output dimension n
@@ -376,14 +374,6 @@ __global__ void Marlin(
374
375
// parallel: num valid moe blocks
376
int parallel = num_tokens_past_padded / moe_block_size;
379
- int num_valid_blocks = parallel;
380
- if (is_ep) {
381
- for (int i = 0; i < parallel; i++) {
382
- if (expert_ids_ptr[i] == -1) num_valid_blocks--;
383
- }
384
- }
385
- int num_invalid_blocks = parallel - num_valid_blocks;
386
- parallel = num_valid_blocks;
377
378
int k_tiles = prob_k / 16 / thread_k_blocks;
379
int n_tiles = prob_n / 16 / thread_n_blocks;
@@ -538,22 +528,8 @@ __global__ void Marlin(
528
if (par_id >= parallel) return;
529
530
old_expert_id = expert_id;
541
- if (num_invalid_blocks > 0) {
542
- int skip_count = par_id;
543
- for (int i = 0; i < num_tokens_past_padded / moe_block_size; i++) {
544
- expert_id = expert_ids_ptr[i];
545
- if (expert_id != -1) {
546
- if (skip_count == 0) {
547
- block_id = i;
548
- break;
549
- };
550
- skip_count--;
551
- };
552
- }
553
- } else {
554
- block_id = par_id;
555
- expert_id = expert_ids_ptr[block_id];
556
- }
531
+ block_id = par_id;
532
+ expert_id = expert_ids_ptr[block_id];
533
534
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
535
uint16_t val = global_scale_ptr[expert_id];
csrc/moe/marlin_moe_wna16/ops.cu
+14
-14
@@ -336,14 +336,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
336
void* perm, void* a_tmp, void* sorted_token_ids,
337
void* expert_ids, void* num_tokens_past_padded,
338
void* topk_weights, int moe_block_size, int num_experts,
339
- int top_k, bool mul_topk_weights, bool is_ep, int prob_m,
340
- int prob_n, int prob_k, void* workspace,
341
- vllm::ScalarType const& a_type, vllm::ScalarType const& b_type,
342
- vllm::ScalarType const& c_type, vllm::ScalarType const& s_type,
343
- bool has_bias, bool has_act_order, bool is_k_full, bool has_zp,
344
- int num_groups, int group_size, int dev, cudaStream_t stream,
345
- int thread_k, int thread_n, int sms, int blocks_per_sm,
346
- bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
339
+ int top_k, bool mul_topk_weights, int prob_m, int prob_n,
340
+ int prob_k, void* workspace, vllm::ScalarType const& a_type,
341
+ vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
342
+ vllm::ScalarType const& s_type, bool has_bias,
343
+ bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
344
+ int group_size, int dev, cudaStream_t stream, int thread_k,
345
+ int thread_n, int sms, int blocks_per_sm, bool use_atomic_add,
346
+ bool use_fp32_reduce, bool is_zp_float) {
347
int thread_m_blocks = div_ceil(moe_block_size, 16);
348
bool m_block_size_8 = moe_block_size == 8;
349
bool is_a_8bit = a_type.size_bits() == 8;
@@ -523,7 +523,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
523
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
524
A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
525
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
526
- topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
526
+ topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m,
527
prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
528
// clang-format on
529
}
@@ -541,7 +541,7 @@ torch::Tensor moe_wna16_marlin_gemm(
541
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
542
torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
543
torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
544
- int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
544
+ int64_t moe_block_size, int64_t top_k, bool mul_topk_weights,
545
vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
546
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
547
bool is_zp_float, int64_t thread_k, int64_t thread_n,
@@ -855,9 +855,9 @@ torch::Tensor moe_wna16_marlin_gemm(
855
perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
856
expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
857
topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
858
- mul_topk_weights, is_ep, size_m, size_n, size_k, workspace.data_ptr(),
859
- a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
860
- has_zp, num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
858
+ mul_topk_weights, size_m, size_n, size_k, workspace.data_ptr(), a_type,
859
+ b_type, c_type, s_type, has_bias, has_act_order, is_k_full, has_zp,
860
+ num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
861
thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
862
is_zp_float);
863
@@ -866,4 +866,4 @@ torch::Tensor moe_wna16_marlin_gemm(
866
867
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
868
m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
869
-}
\ No newline at end of file
869
+}
csrc/moe/torch_bindings.cpp
+1
-1
@@ -71,7 +71,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
71
"Tensor sorted_token_ids,"
72
"Tensor! expert_ids, Tensor! num_tokens_past_padded,"
73
"Tensor! topk_weights, int moe_block_size, int top_k, "
74
- "bool mul_topk_weights, bool is_ep, int b_type_id,"
74
+ "bool mul_topk_weights, int b_type_id,"
75
"int size_m, int size_n, int size_k,"
76
"bool is_full_k, bool use_atomic_add,"
77
"bool use_fp32_reduce, bool is_zp_float,"
vllm/_custom_ops.py
-3
@@ -2162,7 +2162,6 @@ def moe_wna16_marlin_gemm(
2162
moe_block_size: int,
2163
top_k: int,
2164
mul_topk_weights: bool,
2165
- is_ep: bool,
2165
b_q_type: ScalarType,
2166
size_m: int,
2167
size_n: int,
@@ -2194,7 +2193,6 @@ def moe_wna16_marlin_gemm(
2193
moe_block_size,
2194
top_k,
2195
mul_topk_weights,
2197
- is_ep,
2196
b_q_type.id,
2197
size_m,
2198
size_n,
@@ -2256,7 +2254,6 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe")
2254
moe_block_size: int,
2255
top_k: int,
2256
mul_topk_weights: bool,
2259
- is_ep: bool,
2257
b_q_type: ScalarType,
2258
size_m: int,
2259
size_n: int,
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
-2
@@ -135,7 +135,6 @@ def _fused_marlin_moe(
135
moe_block_size=block_size_m,
136
top_k=num_topk,
137
mul_topk_weights=apply_router_weight_on_input,
138
- is_ep=expert_map is not None,
138
b_q_type=quant_type,
139
size_m=M,
140
size_n=2 * N,
@@ -187,7 +186,6 @@ def _fused_marlin_moe(
186
moe_block_size=block_size_m,
187
top_k=1,
188
mul_topk_weights=not apply_router_weight_on_input,
190
- is_ep=expert_map is not None,
189
b_q_type=quant_type,
190
size_m=M * num_topk,
191
size_n=K,