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, 02:45 UTC
Add thread_n=64 support to Marlin MoE (#32360)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 15, 2026 at 19:45 UTC
83239ff19ac6638ee85a51c3829e25680d3543c8
3 files changed
+8
-5
csrc/moe/marlin_moe_wna16/generate_kernels.py
+1
-1
@@ -58,7 +58,7 @@ TEMPLATE = (
58
"( MARLIN_KERNEL_PARAMS );"
59
)
60
61
-THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
61
+THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
62
63
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
64
csrc/moe/marlin_moe_wna16/ops.cu
+4
-2
@@ -126,14 +126,16 @@ thread_config_t small_batch_thread_configs[] = {
126
127
// thread_k, thread_n, num_threads
128
{128, 128, 256},
129
- {64, 128, 128}};
129
+ {64, 128, 128},
130
+ {128, 64, 128}};
131
132
thread_config_t large_batch_thread_configs[] = {
133
// Ordered by priority
134
135
// thread_k, thread_n, num_threads
136
{64, 256, 256},
136
- {64, 128, 128}};
137
+ {64, 128, 128},
138
+ {128, 64, 128}};
139
140
typedef struct {
141
int blocks_per_sm;
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+3
-2
@@ -226,6 +226,7 @@ def prepare_fp8_moe_layer_for_marlin(
226
e = layer.num_experts
227
k = layer.hidden_size
228
n = layer.intermediate_size_per_partition
229
+ w13_n = w13_weight.size(1)
230
weight_block_size = getattr(layer, "weight_block_size", None)
231
232
# WORKSPACE
@@ -240,7 +241,7 @@ def prepare_fp8_moe_layer_for_marlin(
241
def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
242
tensor_list = []
243
if "w13" in name:
243
- size_n, size_k = n * 2, k
244
+ size_n, size_k = w13_n, k
245
else:
246
size_n, size_k = k, n
247
@@ -268,7 +269,7 @@ def prepare_fp8_moe_layer_for_marlin(
269
scales = scales.to(layer.orig_dtype)
270
tensor_list = []
271
if "w13" in name:
271
- size_n, size_k = n * 2, k
272
+ size_n, size_k = w13_n, k
273
else:
274
size_n, size_k = k, n
275