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

Add thread_n=64 support to Marlin MoE (#32360)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2026-01-15 19:45:44 -05:00
committed by GitHub
parent c277fbdf31
commit 83239ff19a
3 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/moe/marlin_moe_wna16/generate_kernels.py
View File

@@ -58,7 +58,7 @@ TEMPLATE = (
"( MARLIN_KERNEL_PARAMS );"
)
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]

6
csrc/moe/marlin_moe_wna16/ops.cu
View File

@@ -126,14 +126,16 @@ thread_config_t small_batch_thread_configs[] = {
// thread_k, thread_n, num_threads
{128, 128, 256},
{64, 128, 128}};
{64, 128, 128},
{128, 64, 128}};
thread_config_t large_batch_thread_configs[] = {
// Ordered by priority
// thread_k, thread_n, num_threads
{64, 256, 256},
{64, 128, 128}};
{64, 128, 128},
{128, 64, 128}};
typedef struct {
int blocks_per_sm;

5
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
View File

@@ -226,6 +226,7 @@ def prepare_fp8_moe_layer_for_marlin(
e = layer.num_experts
k = layer.hidden_size
n = layer.intermediate_size_per_partition
w13_n = w13_weight.size(1)
weight_block_size = getattr(layer, "weight_block_size", None)
# WORKSPACE
@@ -240,7 +241,7 @@ def prepare_fp8_moe_layer_for_marlin(
def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
tensor_list = []
if "w13" in name:
size_n, size_k = n * 2, k
size_n, size_k = w13_n, k
else:
size_n, size_k = k, n
@@ -268,7 +269,7 @@ def prepare_fp8_moe_layer_for_marlin(
scales = scales.to(layer.orig_dtype)
tensor_list = []
if "w13" in name:
size_n, size_k = n * 2, k
size_n, size_k = w13_n, k
else:
size_n, size_k = k, n