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 14, 05:30 UTC
[GPTOSS][DP/EP][Marlin] Enable GPTOSS Batched DP/EP using Marlin kernels (#25997)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Oct 16, 2025 at 15:53 UTC
fb0571b077e749d041c1d9802dda5ceb42c9e8b1
12 files changed
+1153
-314
csrc/moe/moe_align_sum_kernels.cu
+92
@@ -8,12 +8,77 @@
8
9
#include "../cuda_compat.h"
10
#include "../dispatch_utils.h"
11
+#include "core/math.hpp"
12
13
#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
14
15
namespace vllm {
16
namespace moe {
17
18
+namespace batched_moe_align_block_size {
19
+
20
+// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
21
+static constexpr int32_t num_threads = 1024;
22
+static constexpr int32_t num_blocks = 1;
23
+__global__ void batched_moe_align_block_size_kernel(
24
+ int32_t const num_batches, int32_t const max_tokens_per_batch,
25
+ int32_t const block_size, int32_t const* __restrict__ batch_num_tokens,
26
+ int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids,
27
+ int32_t* __restrict__ num_tokens_post_pad) {
28
+ // TODO(varun): This is a naive implementation. Could be optimized.
29
+
30
+ size_t const batch_id = threadIdx.x;
31
+ size_t const stride = blockDim.x * gridDim.x;
32
+ int32_t const num_blocks_per_batch =
33
+ CEILDIV(max_tokens_per_batch, block_size);
34
+ int32_t const sorted_ids_size =
35
+ num_blocks_per_batch * num_batches * block_size;
36
+ int32_t const block_ids_size = sorted_ids_size / block_size;
37
+ int32_t const SENTINEL =
38
+ num_batches * max_tokens_per_batch; // To denote invalid entries.
39
+ // Intialize sorted_ids
40
+ for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) {
41
+ sorted_ids[i] = SENTINEL;
42
+ }
43
+ // Intialize expert_ids with -1
44
+ for (size_t i = threadIdx.x; i < block_ids_size; i += stride) {
45
+ block_ids[i] = -1;
46
+ }
47
+
48
+ int32_t b_num_tokens = 0;
49
+ if (batch_id < num_batches) {
50
+ b_num_tokens = batch_num_tokens[batch_id];
51
+ }
52
+ int32_t const ceil_b_num_tokens =
53
+ CEILDIV(b_num_tokens, block_size) * block_size;
54
+
55
+ // Compute prefix sum over token counts per expert
56
+ using BlockScan = cub::BlockScan<int32_t, 1024>;
57
+ __shared__ typename BlockScan::TempStorage temp_storage;
58
+ int cumsum_val;
59
+ BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val);
60
+ __syncthreads();
61
+
62
+ bool const is_last_batch = batch_id == (num_batches - 1);
63
+ if (is_last_batch) {
64
+ *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens;
65
+ }
66
+
67
+ if (batch_id < num_batches) {
68
+ int32_t const batch_offset = batch_id * max_tokens_per_batch;
69
+ for (size_t i = 0; i < b_num_tokens; ++i) {
70
+ sorted_ids[cumsum_val + i] = batch_offset + i;
71
+ }
72
+
73
+ int32_t const block_start = cumsum_val / block_size;
74
+ int32_t const num_blocks = ceil_b_num_tokens / block_size;
75
+ for (size_t i = 0; i < num_blocks; ++i) {
76
+ block_ids[block_start + i] = batch_id;
77
+ }
78
+ }
79
+}
80
+} // namespace batched_moe_align_block_size
81
+
82
template <typename scalar_t>
83
__global__ void moe_align_block_size_kernel(
84
const scalar_t* __restrict__ topk_ids,
@@ -280,6 +345,33 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
345
});
346
}
347
348
+void batched_moe_align_block_size(int64_t max_tokens_per_batch,
349
+ int64_t block_size,
350
+ torch::Tensor const& batch_num_tokens,
351
+ torch::Tensor sorted_ids,
352
+ torch::Tensor batch_ids,
353
+ torch::Tensor num_tokens_post_pad) {
354
+ namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
355
+
356
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
357
+ int32_t const B = batch_num_tokens.size(0);
358
+ int32_t const num_blocks_per_batch =
359
+ round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size;
360
+ int32_t const num_blocks = num_blocks_per_batch * B;
361
+ int64_t const sorted_ids_size = num_blocks * block_size;
362
+
363
+ TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
364
+ TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
365
+ TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
366
+ TORCH_CHECK(B <= batched_kernel::num_threads);
367
+
368
+ batched_kernel::batched_moe_align_block_size_kernel<<<
369
+ batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
370
+ B, max_tokens_per_batch, block_size, batch_num_tokens.data_ptr<int32_t>(),
371
+ sorted_ids.data_ptr<int32_t>(), batch_ids.data_ptr<int32_t>(),
372
+ num_tokens_post_pad.data_ptr<int32_t>());
373
+}
374
+
375
void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size]
376
torch::Tensor& output) // [num_tokens, hidden_size]
377
{
csrc/moe/moe_ops.h
+8
@@ -12,6 +12,14 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
12
int64_t block_size, torch::Tensor sorted_token_ids,
13
torch::Tensor experts_ids,
14
torch::Tensor num_tokens_post_pad);
15
+
16
+void batched_moe_align_block_size(int64_t max_tokens_per_batch,
17
+ int64_t block_size,
18
+ torch::Tensor const& expert_num_tokens,
19
+ torch::Tensor sorted_ids,
20
+ torch::Tensor expert_ids,
21
+ torch::Tensor num_tokens_post_pad);
22
+
23
#ifndef USE_ROCM
24
torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
25
torch::Tensor b_qweight, torch::Tensor b_scales,
csrc/moe/torch_bindings.cpp
+11
@@ -22,6 +22,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
22
" Tensor! num_tokens_post_pad) -> ()");
23
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
24
25
+ // Aligning the number of tokens to be processed by each expert such
26
+ // that it is divisible by the block size, but for the batched case.
27
+ m.def(
28
+ "batched_moe_align_block_size(int max_tokens_per_batch,"
29
+ " int block_size, Tensor expert_num_tokens,"
30
+ " Tensor! sorted_token_ids,"
31
+ " Tensor! experts_ids,"
32
+ " Tensor! num_tokens_post_pad) -> ()");
33
+ m.impl("batched_moe_align_block_size", torch::kCUDA,
34
+ &batched_moe_align_block_size);
35
+
36
#ifndef USE_ROCM
37
m.def(
38
"moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
docs/design/moe_kernel_features.md
+5
-5
@@ -92,8 +92,8 @@ To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels
92
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],</br>[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
93
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
94
| deep gemm+triton<sup>2</sup> | standard,</br>batched | all<sup>1</sup> | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],</br>[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts] |
95
-| marlin | standard | <sup>3</sup> | <sup>3</sup> | silu,</br>swigluoai | Y | N | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe] |
96
-| marlin experts | standard | N/A | N/A | silu,</br>swigluoai | Y | Y | [`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts] |
95
+| marlin | standard | <sup>3</sup> | <sup>3</sup> | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
96
+| marlin experts | standard,</br>batched | N/A | N/A | silu,</br>swigluoai | Y | Y | [`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
97
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
98
| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
99
| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
@@ -115,6 +115,6 @@ The following table shows "families" of modular kernels that are intended to wor
115
116
| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
117
|----------------------------------|------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------|
118
-| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
119
-| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8`|
120
-| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
118
+| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
119
+| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts`|
120
+| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
tests/kernels/moe/test_moe.py
+360
-202
@@ -7,6 +7,8 @@ Run `pytest tests/kernels/test_moe.py`.
7
8
import functools
9
from collections.abc import Callable
10
+from dataclasses import dataclass
11
+from typing import Any
12
13
import pytest
14
import torch
@@ -26,7 +28,10 @@ from vllm.model_executor.layers.fused_moe.config import (
28
int4_w4a16_moe_quant_config,
29
int8_w8a16_moe_quant_config,
30
)
29
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
31
+from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
32
+ batched_fused_marlin_moe,
33
+ fused_marlin_moe,
34
+)
35
from vllm.model_executor.layers.fused_moe.fused_moe import (
36
fused_topk,
37
modular_triton_fused_moe,
@@ -564,6 +569,105 @@ def marlin_moe_generate_valid_test_cases():
569
return cases
570
571
572
+@dataclass
573
+class MarlinMoEWeightData:
574
+ w_ref: torch.Tensor
575
+ qweight: torch.Tensor
576
+ scales: torch.Tensor
577
+ global_scale: torch.Tensor | None
578
+ g_idx: torch.Tensor | None
579
+ zeros: torch.Tensor | None
580
+ sort_indices: torch.Tensor | None
581
+ marlin_bias: torch.Tensor | None
582
+
583
+ @staticmethod
584
+ def make(
585
+ w: torch.Tensor,
586
+ quant_type: ScalarType,
587
+ group_size: int,
588
+ act_order: bool | None = None,
589
+ bias: torch.Tensor | None = None,
590
+ ) -> "MarlinMoEWeightData":
591
+ assert w.ndim == 3
592
+ has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
593
+ k = w.shape[-1]
594
+
595
+ w_ref_l: list[torch.Tensor] = []
596
+ qweight_l: list[torch.Tensor] = []
597
+ scales_l: list[torch.Tensor] = []
598
+ global_scale_l: list[torch.Tensor] = []
599
+ zeros_l: list[torch.Tensor] = []
600
+ g_idx_l: list[torch.Tensor] = []
601
+ sort_indices_l: list[torch.Tensor] = []
602
+ bias_l: list[torch.Tensor] = []
603
+
604
+ for i in range(w.shape[0]):
605
+ if quant_type == scalar_types.float4_e2m1f:
606
+ if group_size == 16:
607
+ w_ref, qweight, scales, global_scale = (
608
+ rand_marlin_weight_nvfp4_like(w[i], group_size)
609
+ )
610
+ else:
611
+ w_ref, qweight, scales = rand_marlin_weight_mxfp4_like(
612
+ w[i], group_size
613
+ )
614
+ global_scale = None
615
+
616
+ w_ref_l.append(w_ref.T)
617
+ qweight_l.append(qweight)
618
+ scales_l.append(scales)
619
+ if global_scale is not None:
620
+ global_scale_l.append(global_scale)
621
+ elif quant_type == scalar_types.float8_e4m3fn:
622
+ w_ref, qweight, scales = marlin_quant_fp8_torch(w[i], group_size)
623
+ w_ref_l.append(w_ref.T)
624
+ qweight_l.append(qweight)
625
+ scales_l.append(scales)
626
+ elif has_zp:
627
+ w_ref, qweight, scales, zeros = awq_marlin_quantize(
628
+ w[i].transpose(1, 0), quant_type, group_size
629
+ )
630
+
631
+ w_ref_l.append(w_ref.T)
632
+ qweight_l.append(qweight)
633
+ scales_l.append(scales)
634
+ zeros_l.append(zeros)
635
+ else:
636
+ test_perm = torch.randperm(k)
637
+ w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
638
+ w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
639
+ )
640
+
641
+ w_ref_l.append(w_ref.T)
642
+ qweight_l.append(qweight)
643
+ scales_l.append(scales)
644
+ g_idx_l.append(g_idx)
645
+ sort_indices_l.append(sort_indices)
646
+
647
+ if bias is not None:
648
+ bias_l.append(marlin_permute_bias(bias[i]))
649
+
650
+ w_ref = stack_and_dev(w_ref_l)
651
+ qweight = stack_and_dev(qweight_l).contiguous()
652
+ scales = stack_and_dev(scales_l)
653
+ global_scale = stack_and_dev(global_scale_l) if global_scale_l else None
654
+ g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
655
+ zeros = stack_and_dev(zeros_l) if zeros_l else None
656
+ sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
657
+ marlin_bias = stack_and_dev(bias_l) if bias_l else None
658
+
659
+ return MarlinMoEWeightData(
660
+ w_ref=w_ref,
661
+ qweight=qweight,
662
+ scales=scales,
663
+ global_scale=global_scale,
664
+ g_idx=g_idx,
665
+ zeros=zeros,
666
+ sort_indices=sort_indices,
667
+ marlin_bias=marlin_bias,
668
+ )
669
+
670
+
671
@pytest.mark.flaky(reruns=2)
672
@pytest.mark.parametrize(
673
("m, n, k, e, topk, ep_size, dtype, group_size,act_order, quant_type, is_k_full"),
@@ -584,7 +688,6 @@ def test_fused_marlin_moe(
688
is_k_full: bool,
689
):
690
torch.cuda.manual_seed(0)
587
- has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
691
692
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
693
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
@@ -600,152 +703,44 @@ def test_fused_marlin_moe(
703
else:
704
e_map = None
705
603
- w_ref1_l = []
604
- qweight1_l = []
605
- scales1_l = []
606
- global_scale1_l = []
607
- zeros1_l = []
608
- g_idx1_l = []
609
- sort_indices1_l = []
610
-
611
- for i in range(w1.shape[0]):
612
- if quant_type == scalar_types.float4_e2m1f:
613
- if group_size == 16:
614
- w_ref1, qweight1, scales1, global_scale1 = (
615
- rand_marlin_weight_nvfp4_like(w1[i], group_size)
616
- )
617
- else:
618
- w_ref1, qweight1, scales1 = rand_marlin_weight_mxfp4_like(
619
- w1[i], group_size
620
- )
621
- global_scale1 = None
622
-
623
- w_ref1_l.append(w_ref1.T)
624
- qweight1_l.append(qweight1)
625
- scales1_l.append(scales1)
626
- if global_scale1 is not None:
627
- global_scale1_l.append(global_scale1)
628
- elif quant_type == scalar_types.float8_e4m3fn:
629
- w_ref1, qweight1, scales1 = marlin_quant_fp8_torch(w1[i], group_size)
630
- w_ref1_l.append(w_ref1.T)
631
- qweight1_l.append(qweight1)
632
- scales1_l.append(scales1)
633
- elif has_zp:
634
- w_ref1, qweight1, scales1, zeros1 = awq_marlin_quantize(
635
- w1[i].transpose(1, 0), quant_type, group_size
636
- )
637
-
638
- w_ref1_l.append(w_ref1.T)
639
- qweight1_l.append(qweight1)
640
- scales1_l.append(scales1)
641
- zeros1_l.append(zeros1)
642
- else:
643
- test_perm = torch.randperm(k)
644
- w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = marlin_quantize(
645
- w1[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
646
- )
647
-
648
- w_ref1_l.append(w_ref1.T)
649
- qweight1_l.append(qweight1)
650
- scales1_l.append(scales1)
651
- g_idx1_l.append(g_idx1)
652
- sort_indices1_l.append(sort_indices1)
653
-
654
- w_ref1 = stack_and_dev(w_ref1_l)
655
- qweight1 = stack_and_dev(qweight1_l).contiguous()
656
- scales1 = stack_and_dev(scales1_l)
657
- global_scale1 = stack_and_dev(global_scale1_l) if global_scale1_l else None
658
- g_idx1 = stack_and_dev(g_idx1_l) if g_idx1_l else None
659
- zeros1 = stack_and_dev(zeros1_l) if zeros1_l else None
660
- sort_indices1 = stack_and_dev(sort_indices1_l) if sort_indices1_l else None
661
-
662
- w_ref2_l = []
663
- qweight2_l = []
664
- scales2_l = []
665
- global_scale2_l = []
666
- zeros2_l = []
667
- g_idx2_l = []
668
- sort_indices2_l = []
669
-
670
- for i in range(w2.shape[0]):
671
- if quant_type == scalar_types.float4_e2m1f:
672
- if group_size == 16:
673
- w_ref2, qweight2, scales2, global_scale2 = (
674
- rand_marlin_weight_nvfp4_like(w2[i], group_size)
675
- )
676
- else:
677
- w_ref2, qweight2, scales2 = rand_marlin_weight_mxfp4_like(
678
- w2[i], group_size
679
- )
680
- global_scale2 = None
681
-
682
- w_ref2_l.append(w_ref2.T)
683
- qweight2_l.append(qweight2)
684
- scales2_l.append(scales2)
685
- if global_scale2 is not None:
686
- global_scale2_l.append(global_scale2)
687
- elif quant_type == scalar_types.float8_e4m3fn:
688
- w_ref2, qweight2, scales2 = marlin_quant_fp8_torch(w2[i], group_size)
689
- w_ref2_l.append(w_ref2.T)
690
- qweight2_l.append(qweight2)
691
- scales2_l.append(scales2)
692
- elif has_zp:
693
- w_ref2, qweight2, scales2, zeros2 = awq_marlin_quantize(
694
- w2[i].transpose(1, 0), quant_type, group_size
695
- )
696
-
697
- w_ref2_l.append(w_ref2.T)
698
- qweight2_l.append(qweight2)
699
- scales2_l.append(scales2)
700
- zeros2_l.append(zeros2)
701
- else:
702
- test_perm = torch.randperm(n)
703
- w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = marlin_quantize(
704
- w2[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
705
- )
706
-
707
- w_ref2_l.append(w_ref2.T)
708
- qweight2_l.append(qweight2)
709
- scales2_l.append(scales2)
710
- g_idx2_l.append(g_idx2)
711
- sort_indices2_l.append(sort_indices2)
706
+ w1_data = MarlinMoEWeightData.make(
707
+ w=w1, quant_type=quant_type, group_size=group_size, act_order=act_order
708
+ )
709
713
- w_ref2 = stack_and_dev(w_ref2_l)
714
- qweight2 = stack_and_dev(qweight2_l).contiguous()
715
- scales2 = stack_and_dev(scales2_l)
716
- global_scale2 = stack_and_dev(global_scale2_l) if global_scale2_l else None
717
- g_idx2 = stack_and_dev(g_idx2_l) if g_idx2_l else None
718
- zeros2 = stack_and_dev(zeros2_l) if zeros2_l else None
719
- sort_indices2 = stack_and_dev(sort_indices2_l) if sort_indices2_l else None
710
+ w2_data = MarlinMoEWeightData.make(
711
+ w=w2, quant_type=quant_type, group_size=group_size, act_order=act_order
712
+ )
713
714
score = torch.randn((m, e), device="cuda", dtype=dtype)
715
716
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
717
718
with set_current_vllm_config(vllm_config):
726
- torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, expert_map=e_map)
719
+ torch_output = torch_moe(
720
+ a, w1_data.w_ref, w2_data.w_ref, score, topk, expert_map=e_map
721
+ )
722
723
marlin_output = fused_marlin_moe(
724
a,
730
- qweight1,
731
- qweight2,
725
+ w1_data.qweight,
726
+ w2_data.qweight,
727
None,
728
None,
734
- scales1,
735
- scales2,
729
+ w1_data.scales,
730
+ w2_data.scales,
731
score,
732
topk_weights,
733
topk_ids,
734
global_num_experts=e,
735
expert_map=e_map,
741
- global_scale1=global_scale1,
742
- global_scale2=global_scale2,
743
- g_idx1=g_idx1,
744
- g_idx2=g_idx2,
745
- sort_indices1=sort_indices1,
746
- sort_indices2=sort_indices2,
747
- w1_zeros=zeros1,
748
- w2_zeros=zeros2,
736
+ global_scale1=w1_data.global_scale,
737
+ global_scale2=w2_data.global_scale,
738
+ g_idx1=w1_data.g_idx,
739
+ g_idx2=w2_data.g_idx,
740
+ sort_indices1=w1_data.sort_indices,
741
+ sort_indices2=w2_data.sort_indices,
742
+ w1_zeros=w1_data.zeros,
743
+ w2_zeros=w2_data.zeros,
744
quant_type_id=quant_type.id,
745
is_k_full=is_k_full,
746
)
@@ -773,92 +768,52 @@ def test_fused_marlin_moe_with_bias(m):
768
b_bias1 = torch.randn((e, 2 * n), device="cuda", dtype=dtype) / 10
769
b_bias2 = torch.randn((e, k), device="cuda", dtype=dtype) / 10
770
776
- b_bias1_l = []
777
- w_ref1_l = []
778
- qweight1_l = []
779
- scales1_l = []
780
- g_idx1_l = []
781
- sort_indices1_l = []
782
-
783
- for i in range(w1.shape[0]):
784
- test_perm = torch.randperm(k)
785
- w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = marlin_quantize(
786
- w1[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
787
- )
788
-
789
- w_ref1_l.append(w_ref1.T)
790
- qweight1_l.append(qweight1)
791
- scales1_l.append(scales1)
792
- g_idx1_l.append(g_idx1)
793
- sort_indices1_l.append(sort_indices1)
794
- b_bias1_l.append(marlin_permute_bias(b_bias1[i]))
795
-
796
- w_ref1 = stack_and_dev(w_ref1_l)
797
- qweight1 = stack_and_dev(qweight1_l).contiguous()
798
- scales1 = stack_and_dev(scales1_l)
799
- global_scale1 = None
800
- g_idx1 = stack_and_dev(g_idx1_l) if g_idx1_l else None
801
- zeros1 = None
802
- sort_indices1 = stack_and_dev(sort_indices1_l) if sort_indices1_l else None
803
- marlin_bias1 = stack_and_dev(b_bias1_l) if b_bias1_l else None
804
-
805
- b_bias2_l = []
806
- w_ref2_l = []
807
- qweight2_l = []
808
- scales2_l = []
809
- g_idx2_l = []
810
- sort_indices2_l = []
811
-
812
- for i in range(w2.shape[0]):
813
- test_perm = torch.randperm(n)
814
- w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = marlin_quantize(
815
- w2[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
816
- )
771
+ w1_data = MarlinMoEWeightData.make(
772
+ w=w1,
773
+ quant_type=quant_type,
774
+ group_size=group_size,
775
+ act_order=act_order,
776
+ bias=b_bias1,
777
+ )
778
818
- w_ref2_l.append(w_ref2.T)
819
- qweight2_l.append(qweight2)
820
- scales2_l.append(scales2)
821
- g_idx2_l.append(g_idx2)
822
- sort_indices2_l.append(sort_indices2)
823
- b_bias2_l.append(marlin_permute_bias(b_bias2[i]))
824
-
825
- w_ref2 = stack_and_dev(w_ref2_l)
826
- qweight2 = stack_and_dev(qweight2_l).contiguous()
827
- scales2 = stack_and_dev(scales2_l)
828
- global_scale2 = None
829
- g_idx2 = stack_and_dev(g_idx2_l) if g_idx2_l else None
830
- zeros2 = None
831
- sort_indices2 = stack_and_dev(sort_indices2_l) if sort_indices2_l else None
832
- marlin_bias2 = stack_and_dev(b_bias2_l) if b_bias2_l else None
779
+ w2_data = MarlinMoEWeightData.make(
780
+ w=w2,
781
+ quant_type=quant_type,
782
+ group_size=group_size,
783
+ act_order=act_order,
784
+ bias=b_bias2,
785
+ )
786
787
score = torch.randn((m, e), device="cuda", dtype=dtype)
788
789
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
790
791
with set_current_vllm_config(vllm_config):
839
- torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, b_bias1, b_bias2)
792
+ torch_output = torch_moe(
793
+ a, w1_data.w_ref, w2_data.w_ref, score, topk, b_bias1, b_bias2
794
+ )
795
796
marlin_output = fused_marlin_moe(
797
a,
843
- qweight1,
844
- qweight2,
845
- marlin_bias1,
846
- marlin_bias2,
847
- scales1,
848
- scales2,
798
+ w1_data.qweight,
799
+ w2_data.qweight,
800
+ w1_data.marlin_bias,
801
+ w2_data.marlin_bias,
802
+ w1_data.scales,
803
+ w2_data.scales,
804
score,
805
topk_weights,
806
topk_ids,
807
global_num_experts=e,
808
expert_map=None,
854
- global_scale1=global_scale1,
855
- global_scale2=global_scale2,
856
- g_idx1=g_idx1,
857
- g_idx2=g_idx2,
858
- sort_indices1=sort_indices1,
859
- sort_indices2=sort_indices2,
860
- w1_zeros=zeros1,
861
- w2_zeros=zeros2,
809
+ global_scale1=w1_data.global_scale,
810
+ global_scale2=w2_data.global_scale,
811
+ g_idx1=w1_data.g_idx,
812
+ g_idx2=w2_data.g_idx,
813
+ sort_indices1=w1_data.sort_indices,
814
+ sort_indices2=w2_data.sort_indices,
815
+ w1_zeros=w1_data.zeros,
816
+ w2_zeros=w2_data.zeros,
817
quant_type_id=quant_type.id,
818
is_k_full=is_k_full,
819
)
@@ -895,6 +850,41 @@ def test_moe_align_block_size_opcheck():
850
)
851
852
853
+def test_batched_moe_align_block_size_opcheck():
854
+ max_tokens_per_batch = 512
855
+ num_experts = 4
856
+ block_size = 16
857
+
858
+ expert_num_tokens = torch.randint(
859
+ low=0,
860
+ high=max_tokens_per_batch,
861
+ size=(num_experts,),
862
+ dtype=torch.int32,
863
+ device="cuda",
864
+ )
865
+
866
+ max_num_tokens_padded = num_experts * max(max_tokens_per_batch, block_size)
867
+ sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
868
+
869
+ assert max_num_tokens_padded % block_size == 0
870
+ max_num_m_blocks = max_num_tokens_padded // block_size
871
+ expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda")
872
+
873
+ num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device="cuda")
874
+
875
+ opcheck(
876
+ torch.ops._moe_C.batched_moe_align_block_size,
877
+ (
878
+ max_tokens_per_batch,
879
+ block_size,
880
+ expert_num_tokens,
881
+ sorted_ids,
882
+ expert_ids,
883
+ num_tokens_post_pad,
884
+ ),
885
+ )
886
+
887
+
888
@pytest.mark.parametrize("m", [1, 33, 64, 222])
889
@pytest.mark.parametrize("topk", TOP_KS)
890
@pytest.mark.parametrize("k", [128, 511, 1024])
@@ -979,3 +969,171 @@ def test_cpu_fused_moe_basic(m, n, k, e, topk, dtype, with_bias, activation):
969
else:
970
atol = 5e-2
971
torch.testing.assert_close(out, ref, atol=atol, rtol=0)
972
+
973
+
974
+@pytest.mark.parametrize("m", [16, 32, 64])
975
+@pytest.mark.parametrize("n", [128])
976
+@pytest.mark.parametrize("k", [128])
977
+@pytest.mark.parametrize("e", [8, 12, 16, 32])
978
+@pytest.mark.parametrize("topk", [2, 4])
979
+@pytest.mark.parametrize("max_tokens_per_batch", [16, 32, 64])
980
+@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
981
+def test_batched_fused_marlin_moe(
982
+ m: int, n: int, k: int, e: int, topk: int, max_tokens_per_batch: int
983
+):
984
+ print(
985
+ f"testing m={m}, n={n}, k={k}, e={e}, "
986
+ f"topk={topk}, "
987
+ f"max_tokens_per_batch={max_tokens_per_batch}"
988
+ )
989
+ torch.cuda.manual_seed(0)
990
+
991
+ dtype = torch.bfloat16
992
+ quant_dtype = scalar_types.float4_e2m1f
993
+ group_size = 32
994
+
995
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
996
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
997
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
998
+
999
+ w1_data = MarlinMoEWeightData.make(
1000
+ w=w1, quant_type=quant_dtype, group_size=group_size, act_order=None
1001
+ )
1002
+ w2_data = MarlinMoEWeightData.make(
1003
+ w=w2, quant_type=quant_dtype, group_size=group_size, act_order=None
1004
+ )
1005
+
1006
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
1007
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
1008
+
1009
+ class BatchedRun:
1010
+ @staticmethod
1011
+ def _make_expert_num_tokens_cpu(
1012
+ e: int, # num_experts
1013
+ topk_ids_cpu: torch.Tensor,
1014
+ ) -> torch.Tensor:
1015
+ expert_num_tokens_cpu = torch.zeros((e,), dtype=torch.int32, device="cpu")
1016
+ for topk_id in torch.flatten(topk_ids_cpu):
1017
+ expert_num_tokens_cpu[topk_id] += 1
1018
+ return expert_num_tokens_cpu
1019
+
1020
+ def __init__(
1021
+ self,
1022
+ max_tokens_per_batch: int,
1023
+ num_experts: int,
1024
+ _topk_ids: torch.Tensor,
1025
+ _topk_weights: torch.Tensor,
1026
+ ):
1027
+ self.max_tokens_per_batch = max_tokens_per_batch
1028
+ self.e = num_experts
1029
+ self.topk_ids_cpu = _topk_ids.to("cpu")
1030
+ self.topk_weights_cpu = _topk_weights.to("cpu")
1031
+ self.expert_num_tokens_cpu = self._make_expert_num_tokens_cpu(
1032
+ self.e, self.topk_ids_cpu
1033
+ )
1034
+
1035
+ def is_valid(self):
1036
+ """
1037
+ Return True only if the input can be represented in a Batched
1038
+ format.
1039
+ """
1040
+ return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch)
1041
+
1042
+ def _scatter(self, hidden_states: torch.Tensor) -> torch.Tensor:
1043
+ hidden_states_cpu = hidden_states.to("cpu")
1044
+ K = hidden_states_cpu.size(1)
1045
+ batched_hidden_states_cpu = torch.empty(
1046
+ (e, max_tokens_per_batch, K),
1047
+ dtype=hidden_states_cpu.dtype,
1048
+ device="cpu",
1049
+ )
1050
+
1051
+ counter_cpu = torch.zeros_like(self.expert_num_tokens_cpu)
1052
+ for t_idx, token in enumerate(hidden_states_cpu):
1053
+ for topk_id in self.topk_ids_cpu[t_idx]:
1054
+ pos_in_batch = counter_cpu[topk_id]
1055
+ batched_hidden_states_cpu[topk_id, pos_in_batch] = token
1056
+ counter_cpu[topk_id] += 1
1057
+ assert torch.allclose(counter_cpu, self.expert_num_tokens_cpu)
1058
+ return batched_hidden_states_cpu.to("cuda")
1059
+
1060
+ def _gather(
1061
+ self, batched_outputs: torch.Tensor, gather_outputs: torch.Tensor
1062
+ ) -> torch.Tensor:
1063
+ batched_outputs_cpu = batched_outputs.to("cpu")
1064
+ gather_outputs_cpu = torch.zeros_like(gather_outputs)
1065
+
1066
+ counter_cpu = torch.zeros((e,), device="cpu", dtype=torch.int32)
1067
+ md = gather_outputs_cpu.size(0)
1068
+ for t_idx in range(md):
1069
+ token = None
1070
+ for topk_id, topk_weight in zip(
1071
+ self.topk_ids_cpu[t_idx], self.topk_weights_cpu[t_idx]
1072
+ ):
1073
+ pos_in_batch = counter_cpu[topk_id]
1074
+ t = batched_outputs_cpu[topk_id, pos_in_batch] * topk_weight
1075
+ if token is None:
1076
+ token = t
1077
+ else:
1078
+ token += t
1079
+ counter_cpu[topk_id] += 1
1080
+ assert token is not None
1081
+ gather_outputs_cpu[t_idx] = token
1082
+ gather_outputs.copy_(gather_outputs_cpu)
1083
+ return gather_outputs
1084
+
1085
+ def run(
1086
+ self, hidden_states: torch.Tensor, fused_marlin_moe_kwargs: dict[Any, Any]
1087
+ ) -> torch.Tensor:
1088
+ assert hidden_states.ndim == 2
1089
+ assert self.is_valid()
1090
+
1091
+ batched_hidden_states = self._scatter(hidden_states)
1092
+
1093
+ kwargs = fused_marlin_moe_kwargs | {
1094
+ "hidden_states": batched_hidden_states,
1095
+ "expert_num_tokens": self.expert_num_tokens_cpu.to("cuda"),
1096
+ }
1097
+ batched_outputs = batched_fused_marlin_moe(**kwargs)
1098
+
1099
+ output = torch.zeros_like(hidden_states)
1100
+ output = self._gather(batched_outputs, output)
1101
+ return output
1102
+
1103
+ kwargs = {
1104
+ "w1": w1_data.qweight,
1105
+ "w2": w2_data.qweight,
1106
+ "bias1": None,
1107
+ "bias2": None,
1108
+ "w1_scale": w1_data.scales,
1109
+ "w2_scale": w2_data.scales,
1110
+ "gating_output": score,
1111
+ "global_num_experts": e,
1112
+ "expert_map": None,
1113
+ "global_scale1": w1_data.global_scale,
1114
+ "global_scale2": w2_data.global_scale,
1115
+ "g_idx1": w1_data.g_idx,
1116
+ "g_idx2": w2_data.g_idx,
1117
+ "sort_indices1": w1_data.sort_indices,
1118
+ "sort_indices2": w2_data.sort_indices,
1119
+ "w1_zeros": w1_data.zeros,
1120
+ "w2_zeros": w2_data.zeros,
1121
+ "quant_type_id": quant_dtype.id,
1122
+ "is_k_full": True,
1123
+ }
1124
+
1125
+ # Reference
1126
+ fused_marlin_moe_kwargs = kwargs | {
1127
+ "hidden_states": a,
1128
+ "topk_ids": topk_ids,
1129
+ "topk_weights": topk_weights,
1130
+ }
1131
+ ref_marlin_output = fused_marlin_moe(**fused_marlin_moe_kwargs)
1132
+
1133
+ # Batched
1134
+ br = BatchedRun(max_tokens_per_batch, e, topk_ids, topk_weights)
1135
+ if not br.is_valid():
1136
+ pytest.skip("Cannot represent data in Batched Format.")
1137
+ marlin_output = br.run(a, kwargs)
1138
+
1139
+ torch.testing.assert_close(marlin_output, ref_marlin_output, atol=1e-3, rtol=0)
tests/kernels/moe/test_moe_align_block_size.py
+94
@@ -9,6 +9,7 @@ import pytest
9
import torch
10
11
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
12
+ batched_moe_align_block_size,
13
moe_align_block_size,
14
)
15
from vllm.platforms import current_platform
@@ -300,3 +301,96 @@ def test_moe_align_block_size_deterministic():
301
assert torch.equal(results[0][2], results[i][2]), (
302
"num_tokens should be deterministic"
303
)
304
+
305
+
306
+@pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512])
307
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
308
+@pytest.mark.parametrize("block_size", [8, 16, 32, 64])
309
+@pytest.mark.parametrize("simulate_empty_batches", [False, True])
310
+def test_batched_moe_align_block_size(
311
+ max_tokens_per_batch: int,
312
+ num_experts: int,
313
+ block_size: int,
314
+ simulate_empty_batches: bool,
315
+):
316
+ def ref_outputs(
317
+ expert_num_tokens: torch.Tensor,
318
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
319
+ E = expert_num_tokens.size(0)
320
+
321
+ # Round up so each batch can be split to blocks evenly.
322
+ Msum = round_up(max_tokens_per_batch, block_size) * E
323
+ ref_sorted_ids = torch.empty((Msum,), dtype=torch.int32)
324
+ ref_expert_ids = torch.empty((Msum // block_size,), dtype=torch.int32)
325
+ ref_num_tokens_post_pad = torch.empty((1,), dtype=torch.int32)
326
+
327
+ # Intialize
328
+ sentinel = E * max_tokens_per_batch
329
+ ref_sorted_ids.fill_(sentinel)
330
+ ref_expert_ids.fill_(-1)
331
+
332
+ # Fill ref_sorted_ids
333
+ i = 0
334
+ for expert_id, expert_nt in enumerate(expert_num_tokens):
335
+ token_offset = expert_id * max_tokens_per_batch
336
+ for j in range(expert_nt):
337
+ ref_sorted_ids[i] = token_offset + j
338
+ i += 1
339
+ # round up i to the next block_size
340
+ i = round_up(i, block_size)
341
+
342
+ ref_num_tokens_post_pad[0] = i
343
+
344
+ # Fill expert_ids
345
+ nt_ceil_sum = 0
346
+ for expert_id, expert_nt in enumerate(expert_num_tokens):
347
+ expert_ids_offset = nt_ceil_sum // block_size
348
+ ceil_expert_nt = round_up(int(expert_nt.item()), block_size)
349
+ num_blocks = ceil_expert_nt // block_size
350
+ for x in range(num_blocks):
351
+ ref_expert_ids[expert_ids_offset + x] = expert_id
352
+ nt_ceil_sum += ceil_expert_nt
353
+
354
+ return (
355
+ ref_sorted_ids.to("cuda"),
356
+ ref_expert_ids.to("cuda"),
357
+ ref_num_tokens_post_pad.to("cuda"),
358
+ )
359
+
360
+ # Compute expert_num_tokens
361
+ expert_num_tokens = torch.randint(
362
+ low=0,
363
+ high=max_tokens_per_batch,
364
+ size=(num_experts,),
365
+ device="cpu",
366
+ dtype=torch.int32,
367
+ )
368
+ if simulate_empty_batches:
369
+ # mark half the batches to have 0 tokens
370
+ zero_batches = torch.randperm(num_experts)[: num_experts // 2]
371
+ expert_num_tokens[zero_batches] = 0
372
+
373
+ # ref outputs
374
+ ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs(
375
+ expert_num_tokens
376
+ )
377
+
378
+ # outputs
379
+ sorted_ids, expert_ids, num_tokens_post_pad = batched_moe_align_block_size(
380
+ max_tokens_per_batch, block_size, expert_num_tokens.to("cuda")
381
+ )
382
+
383
+ assert ref_sorted_ids.size() == sorted_ids.size(), (
384
+ f"{ref_sorted_ids.size()} vs {sorted_ids.size()}"
385
+ )
386
+ assert ref_expert_ids.size() == expert_ids.size(), (
387
+ f"{ref_expert_ids.size()} vs {expert_ids.size()}"
388
+ )
389
+ assert ref_num_tokens_post_pad.size() == num_tokens_post_pad.size(), (
390
+ f"{ref_num_tokens_post_pad.size()} vs {num_tokens_post_pad.size()}"
391
+ )
392
+ torch.testing.assert_close(ref_sorted_ids, sorted_ids, atol=0, rtol=0)
393
+ torch.testing.assert_close(ref_expert_ids, expert_ids, atol=0, rtol=0)
394
+ torch.testing.assert_close(
395
+ ref_num_tokens_post_pad, num_tokens_post_pad, atol=0, rtol=0
396
+ )
vllm/_custom_ops.py
+18
@@ -1789,6 +1789,24 @@ def moe_align_block_size(
1789
)
1790
1791
1792
+def batched_moe_align_block_size(
1793
+ max_tokens_per_batch: int,
1794
+ block_size: int,
1795
+ expert_num_tokens: torch.Tensor,
1796
+ sorted_ids: torch.Tensor,
1797
+ expert_ids: torch.Tensor,
1798
+ num_tokens_post_pad: torch.Tensor,
1799
+) -> None:
1800
+ torch.ops._moe_C.batched_moe_align_block_size(
1801
+ max_tokens_per_batch,
1802
+ block_size,
1803
+ expert_num_tokens,
1804
+ sorted_ids,
1805
+ expert_ids,
1806
+ num_tokens_post_pad,
1807
+ )
1808
+
1809
+
1810
def moe_wna16_gemm(
1811
input: torch.Tensor,
1812
output: torch.Tensor,
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+25
-1
@@ -50,7 +50,31 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
50
51
# DeepEP low-latency kernels are compiled only for certain
52
# specific hidden sizes.
53
- SUPPORTED_HIDDEN_SIZES = [2048, 2560, 4096, 5120, 6144, 7168]
53
+ # NOTE: Keep this list sorted, maybe_roundup_layer_hidden_size depends
54
+ # on it.
55
+ SUPPORTED_HIDDEN_SIZES = [2048, 2560, 4096, 5120, 6144, 7168, 8192]
56
+
57
+ @staticmethod
58
+ def maybe_roundup_layer_hidden_size(hidden_size: int) -> int:
59
+ # Round up hidden size to the closest supported hidden size.
60
+ _supported_hs = DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES
61
+ # Check sorted
62
+ num_supported_hs = len(_supported_hs)
63
+ assert all(
64
+ [
65
+ _supported_hs[i] < _supported_hs[i + 1]
66
+ for i in range(num_supported_hs - 1)
67
+ ]
68
+ )
69
+
70
+ for x in _supported_hs:
71
+ if x >= hidden_size:
72
+ return x
73
+
74
+ raise ValueError(
75
+ f"Hidden Size {hidden_size} is greater than the "
76
+ f"maximum supported hidden size {_supported_hs[-1]}"
77
+ )
78
79
def __init__(
80
self,
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+432
-103
@@ -3,13 +3,16 @@
3
"""Fused MoE utilities for GPTQ."""
4
5
import torch
6
-from typing_extensions import override
6
7
import vllm._custom_ops as ops
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
-from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
10
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
11
+ batched_moe_align_block_size,
12
+ moe_align_block_size,
13
+)
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
+ TopKWeightAndReduceDelegate,
16
TopKWeightAndReduceNoOP,
17
)
18
from vllm.model_executor.layers.fused_moe.utils import _resize_cache, disable_inplace
@@ -21,7 +24,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
24
from vllm.scalar_type import ScalarType, scalar_types
25
26
24
-def fused_marlin_moe(
27
+def _fused_marlin_moe(
28
hidden_states: torch.Tensor,
29
w1: torch.Tensor,
30
w2: torch.Tensor,
@@ -29,14 +32,16 @@ def fused_marlin_moe(
32
bias2: torch.Tensor | None,
33
w1_scale: torch.Tensor,
34
w2_scale: torch.Tensor,
32
- gating_output: torch.Tensor | None,
35
topk_weights: torch.Tensor,
34
- topk_ids: torch.Tensor,
35
- quant_type_id: int,
36
- apply_router_weight_on_input: bool = False,
37
- global_num_experts: int = -1,
38
- activation: str | None = "silu",
39
- expert_map: torch.Tensor | None = None,
36
+ num_topk: int,
37
+ quant_type: ScalarType,
38
+ apply_router_weight_on_input: bool,
39
+ activation: str,
40
+ expert_map: torch.Tensor | None,
41
+ block_size_m: int,
42
+ sorted_token_ids: torch.Tensor,
43
+ expert_ids: torch.Tensor,
44
+ num_tokens_post_padded: torch.Tensor,
45
global_scale1: torch.Tensor | None = None,
46
global_scale2: torch.Tensor | None = None,
47
g_idx1: torch.Tensor | None = None,
@@ -48,106 +53,35 @@ def fused_marlin_moe(
53
workspace: torch.Tensor | None = None,
54
intermediate_cache13: torch.Tensor | None = None,
55
intermediate_cache2: torch.Tensor | None = None,
51
- is_k_full: bool = True,
56
output: torch.Tensor | None = None,
53
- inplace: bool = False,
57
+ is_k_full: bool = True,
58
) -> torch.Tensor:
55
- """
56
- This function computes a Mixture of Experts (MoE) layer using two sets of
57
- weights, w1 and w2, and top-k gating mechanism.
58
-
59
- Parameters:
60
- - hidden_states (torch.Tensor): The input tensor to the MoE layer.
61
- - w1 (torch.Tensor): The first set of expert weights.
62
- - w2 (torch.Tensor): The second set of expert weights.
63
- - w1_scale (torch.Tensor): Scale to be used for w1.
64
- - w2_scale (torch.Tensor): Scale to be used for w2.
65
- - gating_output (Optional[torch.Tensor]): The output of the gating
66
- operation (before softmax).
67
- - g_idx1 (Optional[torch.Tensor]): The first set of act_order indices.
68
- - g_idx2 (Optional[torch.Tensor]): The second set of act_order indices.
69
- - sort_indices1 (Optional[torch.Tensor]): The first act_order input
70
- permutation.
71
- - sort_indices2 (Optional[torch.Tensor]): The second act_order input
72
- permutation.
73
- - topk_weights (torch.Tensor): Top-k weights.
74
- - topk_ids (torch.Tensor): Indices of topk-k elements.
75
- - w1_zeros (Optional[torch.Tensor]): Optional zero points to be used for w1.
76
- - w2_zeros (Optional[torch.Tensor]): Optional zero points to be used for w2.
77
- - num_bits (bool): The number of bits in expert weights quantization.
78
-
79
- Returns:
80
- - torch.Tensor: The output tensor after applying the MoE layer.
81
- """
82
- quant_type = ScalarType.from_id(quant_type_id)
83
- assert quant_type in [
84
- scalar_types.uint4,
85
- scalar_types.uint8b128,
86
- scalar_types.uint4b8,
87
- scalar_types.float8_e4m3fn,
88
- scalar_types.float4_e2m1f,
89
- ]
90
-
91
- bit4_scalar_types = [
92
- scalar_types.uint4,
93
- scalar_types.uint4b8,
94
- scalar_types.float4_e2m1f,
95
- ]
96
- num_bits = 4 if quant_type in bit4_scalar_types else 8
97
-
98
- # Check constraints.
99
- if gating_output is not None:
100
- assert hidden_states.shape[0] == gating_output.shape[0], (
101
- "Number of tokens mismatch"
102
- )
103
- assert hidden_states.shape[1] == w1.shape[1] * 16, "Hidden size mismatch w1"
104
- assert hidden_states.shape[1] == w2.shape[2] // (num_bits // 2), (
105
- "Hidden size mismatch w2"
106
- )
107
- assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
108
- assert w1.is_contiguous(), "Expert weights1 must be contiguous"
109
- assert w2.is_contiguous(), "Expert weights2 must be contiguous"
110
- assert hidden_states.dtype in [torch.float16, torch.bfloat16]
111
- assert num_bits in [4, 8]
112
- assert topk_weights.dtype == torch.float32
113
-
114
- M, K = hidden_states.shape
115
- E = w1.shape[0]
59
+ assert hidden_states.ndim == 2
60
+ M, K = hidden_states.size()
61
N = marlin_moe_intermediate_size(w1, w2)
117
- topk = topk_ids.shape[1]
118
-
119
- # M block size selection logic
120
- # TODO: tune this further for specific models
121
- for block_size_m in [8, 16, 32, 48, 64]:
122
- if M * topk / E / block_size_m < 0.9:
123
- break
124
-
125
- if global_num_experts == -1:
126
- global_num_experts = E
127
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
128
- topk_ids, block_size_m, global_num_experts, expert_map
129
- )
62
63
if workspace is None:
64
workspace = marlin_make_workspace_new(hidden_states.device, 4)
65
134
- if intermediate_cache2 is None:
135
- intermediate_cache2 = torch.empty(
136
- (M * topk, N),
66
+ if intermediate_cache13 is None:
67
+ intermediate_cache13 = torch.empty(
68
+ (M * num_topk * max(2 * N, K),),
69
device=hidden_states.device,
70
dtype=hidden_states.dtype,
71
)
72
141
- if intermediate_cache13 is None:
142
- intermediate_cache13 = torch.empty(
143
- (M * topk * max(2 * N, K),),
73
+ if intermediate_cache2 is None:
74
+ intermediate_cache2 = torch.empty(
75
+ (M * num_topk, N),
76
device=hidden_states.device,
77
dtype=hidden_states.dtype,
78
)
79
148
- intermediate_cache1 = _resize_cache(intermediate_cache13, (M * topk, 2 * N))
149
- intermediate_cache3 = _resize_cache(intermediate_cache13, (M * topk, K))
150
- intermediate_cache2 = _resize_cache(intermediate_cache2, (M * topk, N))
80
+ intermediate_cache1 = _resize_cache(intermediate_cache13, (M * num_topk, 2 * N))
81
+
82
+ intermediate_cache3 = _resize_cache(intermediate_cache13, (M * num_topk, K))
83
+
84
+ intermediate_cache2 = _resize_cache(intermediate_cache2, (M * num_topk, N))
85
86
maybe_warn_marlin_atomic_add(hidden_states.device, hidden_states.dtype)
87
use_atomic_add = (
@@ -171,7 +105,7 @@ def fused_marlin_moe(
105
num_tokens_post_padded,
106
topk_weights,
107
moe_block_size=block_size_m,
174
- top_k=topk,
108
+ top_k=num_topk,
109
mul_topk_weights=apply_router_weight_on_input,
110
is_ep=expert_map is not None,
111
b_q_type=quant_type,
@@ -199,12 +133,15 @@ def fused_marlin_moe(
133
"Only silu and swigluoai activations are supported."
134
)
135
136
+ if output is None:
137
+ output = intermediate_cache3
138
+
139
if expert_map is not None:
203
- intermediate_cache3.zero_()
140
+ output.zero_()
141
205
- intermediate_cache3 = ops.moe_wna16_marlin_gemm(
142
+ output = ops.moe_wna16_marlin_gemm(
143
intermediate_cache2,
207
- intermediate_cache3,
144
+ output,
145
w2,
146
bias2,
147
w2_scale,
@@ -222,13 +159,156 @@ def fused_marlin_moe(
159
mul_topk_weights=not apply_router_weight_on_input,
160
is_ep=expert_map is not None,
161
b_q_type=quant_type,
225
- size_m=M * topk,
162
+ size_m=M * num_topk,
163
size_n=K,
164
size_k=N,
165
is_k_full=is_k_full,
166
use_atomic_add=use_atomic_add,
167
use_fp32_reduce=True,
168
is_zp_float=False,
169
+ )
170
+
171
+ return output
172
+
173
+
174
+def fused_marlin_moe(
175
+ hidden_states: torch.Tensor,
176
+ w1: torch.Tensor,
177
+ w2: torch.Tensor,
178
+ bias1: torch.Tensor | None,
179
+ bias2: torch.Tensor | None,
180
+ w1_scale: torch.Tensor,
181
+ w2_scale: torch.Tensor,
182
+ gating_output: torch.Tensor | None,
183
+ topk_weights: torch.Tensor,
184
+ topk_ids: torch.Tensor,
185
+ quant_type_id: int,
186
+ apply_router_weight_on_input: bool = False,
187
+ global_num_experts: int = -1,
188
+ activation: str | None = "silu",
189
+ expert_map: torch.Tensor | None = None,
190
+ global_scale1: torch.Tensor | None = None,
191
+ global_scale2: torch.Tensor | None = None,
192
+ g_idx1: torch.Tensor | None = None,
193
+ g_idx2: torch.Tensor | None = None,
194
+ sort_indices1: torch.Tensor | None = None,
195
+ sort_indices2: torch.Tensor | None = None,
196
+ w1_zeros: torch.Tensor | None = None,
197
+ w2_zeros: torch.Tensor | None = None,
198
+ workspace: torch.Tensor | None = None,
199
+ intermediate_cache13: torch.Tensor | None = None,
200
+ intermediate_cache2: torch.Tensor | None = None,
201
+ is_k_full: bool = True,
202
+ output: torch.Tensor | None = None,
203
+ inplace: bool = False,
204
+) -> torch.Tensor:
205
+ """
206
+ This function computes a Mixture of Experts (MoE) layer using two sets of
207
+ weights, w1 and w2, and top-k gating mechanism.
208
+
209
+ Parameters:
210
+ - hidden_states (torch.Tensor): The input tensor to the MoE layer.
211
+ - w1 (torch.Tensor): The first set of expert weights.
212
+ - w2 (torch.Tensor): The second set of expert weights.
213
+ - w1_scale (torch.Tensor): Scale to be used for w1.
214
+ - w2_scale (torch.Tensor): Scale to be used for w2.
215
+ - gating_output (torch.Tensor|None): The output of the gating
216
+ operation (before softmax).
217
+ - g_idx1 (torch.Tensor|None): The first set of act_order indices.
218
+ - g_idx2 (torch.Tensor|None): The second set of act_order indices.
219
+ - sort_indices1 (torch.Tensor|None): The first act_order input
220
+ permutation.
221
+ - sort_indices2 (torch.Tensor|None): The second act_order input
222
+ permutation.
223
+ - topk_weights (torch.Tensor): Top-k weights.
224
+ - topk_ids (torch.Tensor): Indices of topk-k elements.
225
+ - w1_zeros (torch.Tensor|None): Optional zero points to be used for w1.
226
+ - w2_zeros (torch.Tensor|None): Optional zero points to be used for w2.
227
+ - num_bits (bool): The number of bits in expert weights quantization.
228
+
229
+ Returns:
230
+ - torch.Tensor: The output tensor after applying the MoE layer.
231
+ """
232
+
233
+ if inplace:
234
+ assert output is None, "Conflicting request"
235
+
236
+ quant_type = ScalarType.from_id(quant_type_id)
237
+ assert quant_type in [
238
+ scalar_types.uint4,
239
+ scalar_types.uint8b128,
240
+ scalar_types.uint4b8,
241
+ scalar_types.float8_e4m3fn,
242
+ scalar_types.float4_e2m1f,
243
+ ]
244
+
245
+ bit4_scalar_types = [
246
+ scalar_types.uint4,
247
+ scalar_types.uint4b8,
248
+ scalar_types.float4_e2m1f,
249
+ ]
250
+ num_bits = 4 if quant_type in bit4_scalar_types else 8
251
+
252
+ M, K = hidden_states.size()
253
+ E = w1.size(0)
254
+ topk = topk_ids.size(1)
255
+
256
+ # Check constraints.
257
+ if gating_output is not None:
258
+ assert gating_output.size(0) == M, "Number of tokens mismatch"
259
+ assert w1.size(1) * 16 == K, "Hidden size mismatch w1"
260
+ assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2"
261
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
262
+ assert w1.is_contiguous(), "Expert weights1 must be contiguous"
263
+ assert w2.is_contiguous(), "Expert weights2 must be contiguous"
264
+ assert hidden_states.dtype in [torch.float16, torch.bfloat16]
265
+ assert num_bits in [4, 8]
266
+ assert topk_weights.dtype == torch.float32
267
+
268
+ # M block size selection logic
269
+ # TODO: tune this further for specific models
270
+ for block_size_m in [8, 16, 32, 48, 64]:
271
+ if M * topk / E / block_size_m < 0.9:
272
+ break
273
+
274
+ if global_num_experts == -1:
275
+ global_num_experts = E
276
+ sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
277
+ topk_ids, block_size_m, global_num_experts, expert_map
278
+ )
279
+
280
+ assert activation is not None
281
+ moe_output = _fused_marlin_moe(
282
+ hidden_states=hidden_states,
283
+ w1=w1,
284
+ w2=w2,
285
+ bias1=bias1,
286
+ bias2=bias2,
287
+ w1_scale=w1_scale,
288
+ w2_scale=w2_scale,
289
+ topk_weights=topk_weights,
290
+ num_topk=topk,
291
+ quant_type=quant_type,
292
+ apply_router_weight_on_input=apply_router_weight_on_input,
293
+ activation=activation,
294
+ expert_map=expert_map,
295
+ block_size_m=block_size_m,
296
+ sorted_token_ids=sorted_token_ids,
297
+ expert_ids=expert_ids,
298
+ num_tokens_post_padded=num_tokens_post_padded,
299
+ global_scale1=global_scale1,
300
+ global_scale2=global_scale2,
301
+ g_idx1=g_idx1,
302
+ g_idx2=g_idx2,
303
+ sort_indices1=sort_indices1,
304
+ sort_indices2=sort_indices2,
305
+ w1_zeros=w1_zeros,
306
+ w2_zeros=w2_zeros,
307
+ workspace=workspace,
308
+ intermediate_cache13=intermediate_cache13,
309
+ intermediate_cache2=intermediate_cache2,
310
+ output=None,
311
+ is_k_full=is_k_full,
312
).view(-1, topk, K)
313
314
if output is None:
@@ -237,16 +317,173 @@ def fused_marlin_moe(
317
else:
318
output = torch.empty_like(hidden_states)
319
240
- return torch.sum(intermediate_cache3.view(-1, topk, K), dim=1, out=output)
320
+ return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output)
321
+
322
+
323
+def batched_fused_marlin_moe(
324
+ hidden_states: torch.Tensor,
325
+ expert_num_tokens: torch.Tensor,
326
+ w1: torch.Tensor,
327
+ w2: torch.Tensor,
328
+ bias1: torch.Tensor | None,
329
+ bias2: torch.Tensor | None,
330
+ w1_scale: torch.Tensor,
331
+ w2_scale: torch.Tensor,
332
+ gating_output: torch.Tensor | None,
333
+ quant_type_id: int,
334
+ apply_router_weight_on_input: bool = False,
335
+ global_num_experts: int = -1,
336
+ activation: str | None = "silu",
337
+ expert_map: torch.Tensor | None = None,
338
+ global_scale1: torch.Tensor | None = None,
339
+ global_scale2: torch.Tensor | None = None,
340
+ g_idx1: torch.Tensor | None = None,
341
+ g_idx2: torch.Tensor | None = None,
342
+ sort_indices1: torch.Tensor | None = None,
343
+ sort_indices2: torch.Tensor | None = None,
344
+ w1_zeros: torch.Tensor | None = None,
345
+ w2_zeros: torch.Tensor | None = None,
346
+ workspace: torch.Tensor | None = None,
347
+ intermediate_cache13: torch.Tensor | None = None,
348
+ intermediate_cache2: torch.Tensor | None = None,
349
+ is_k_full: bool = True,
350
+ output: torch.Tensor | None = None,
351
+ inplace: bool = False,
352
+) -> torch.Tensor:
353
+ """
354
+ This function massages the inputs so the batched hidden_states can be
355
+ presented as a 2D contiguous tensor that could be used with
356
+ _fused_marlin_moe.
357
+
358
+ Note that both batched_fused_marlin_moe and fused_marlin_moe ultimately
359
+ use `ops.moe_wna16_marlin_gemm` for the gemm operation and
360
+ `ops.moe_mna16_marlin_gemm` supports only 2D contiguous hidden_states.
361
+ Note that the moe_align_block_size function indicates,
362
+ - What rows of the A matrix (hidden_states) to access during the
363
+ matmul, via sorted_ids output.
364
+ - What expert_id to use for each block matmul, via expert_ids ouptut.
365
+
366
+ In the batched version, the tokens are already grouped/batched by experts
367
+ they subscribe to. Due to this, we can represent the batched hidden_states
368
+ tensor of shape [B, MAX_TOKENS_PER_BATCH, K] as a 2D tensor of shape,
369
+ [B * MAX_TOKENS_PER_BATCH, K]. We may treat this a 2D contiguous tensor
370
+ with topk=1 as each token (row in the tensor) subscribes to exactly one
371
+ expert_id (which is the batch_id). With the expert_num_tokens tensor, that
372
+ indicates how many tokens are actually valid in each batch, the
373
+ batched_moe_align_block_size function constructs the sorted_ids and
374
+ expert_ids tensors, so only relevant/valid rows of A (hidden_states)
375
+ are accessed and are processed with the correct expert_ids.
376
+ """
377
+
378
+ assert hidden_states.ndim == 3, (
379
+ f"hidden states must be batched. e.g. [B, MAX_TOKENS, K]."
380
+ f"But got {hidden_states.size()}"
381
+ )
382
+ if inplace:
383
+ assert output is None, "Conflicting request."
384
+
385
+ quant_type = ScalarType.from_id(quant_type_id)
386
+ assert quant_type in [
387
+ scalar_types.uint4,
388
+ scalar_types.uint8b128,
389
+ scalar_types.uint4b8,
390
+ scalar_types.float8_e4m3fn,
391
+ scalar_types.float4_e2m1f,
392
+ ]
393
+
394
+ bit4_scalar_types = [
395
+ scalar_types.uint4,
396
+ scalar_types.uint4b8,
397
+ scalar_types.float4_e2m1f,
398
+ ]
399
+ num_bits = 4 if quant_type in bit4_scalar_types else 8
400
+
401
+ B, BATCH_TOKENS_MAX, K = hidden_states.size()
402
+ M = hidden_states.view(-1, K).size(0)
403
+ E = w1.size(0)
404
+
405
+ # Check constraints.
406
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
407
+ assert hidden_states.dtype in [torch.float16, torch.bfloat16]
408
+ assert expert_num_tokens.size(0) == E
409
+ assert B == E, (
410
+ "Batch must be as big as number of experts as the tokens"
411
+ "are sorted into the batch/expert they belong to"
412
+ )
413
+ assert w1.size(1) * 16 == K, "Hidden size mismatch w1"
414
+ assert w2.size(2) // (num_bits // 2) == K, "Hidden size mismatch w2"
415
+ assert w1.is_contiguous(), "Expert weights1 must be contiguous"
416
+ assert w2.is_contiguous(), "Expert weights2 must be contiguous"
417
+ assert num_bits in [4, 8]
418
+
419
+ # Technically, the tokens are already separated by their expert ids.
420
+ # Hidden-States can just be squeezed to have just 2 dimensions,
421
+ # [B * MAX_TOKENS, K] and top_k can be interpreted as just 1.
422
+ topk = 1
423
424
+ # TODO(varun) : Choose a decent block size like in fused_marlin_moe
425
+ block_size_m = 64
426
243
-class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
427
+ sorted_token_ids, expert_ids, num_tokens_post_padded = batched_moe_align_block_size(
428
+ max_tokens_per_batch=BATCH_TOKENS_MAX,
429
+ block_size=block_size_m,
430
+ expert_num_tokens=expert_num_tokens,
431
+ )
432
+
433
+ if output is None and inplace:
434
+ output = hidden_states
435
+
436
+ # TODO (varun): This can be avoided by plumbing the marlin kernel to
437
+ # ignore topk_weights when topk_weights_ptr is a nullptr.
438
+ topk_weights = torch.ones(
439
+ (M, topk), device=hidden_states.device, dtype=torch.float32
440
+ )
441
+
442
+ assert activation is not None
443
+ output = _fused_marlin_moe(
444
+ hidden_states=hidden_states.view(-1, K),
445
+ w1=w1,
446
+ w2=w2,
447
+ bias1=bias1,
448
+ bias2=bias2,
449
+ w1_scale=w1_scale,
450
+ w2_scale=w2_scale,
451
+ topk_weights=topk_weights,
452
+ num_topk=topk,
453
+ quant_type=quant_type,
454
+ apply_router_weight_on_input=apply_router_weight_on_input,
455
+ activation=activation,
456
+ expert_map=expert_map,
457
+ block_size_m=block_size_m,
458
+ sorted_token_ids=sorted_token_ids,
459
+ expert_ids=expert_ids,
460
+ num_tokens_post_padded=num_tokens_post_padded,
461
+ global_scale1=global_scale1,
462
+ global_scale2=global_scale2,
463
+ g_idx1=g_idx1,
464
+ g_idx2=g_idx2,
465
+ sort_indices1=sort_indices1,
466
+ sort_indices2=sort_indices2,
467
+ w1_zeros=w1_zeros,
468
+ w2_zeros=w2_zeros,
469
+ workspace=workspace,
470
+ intermediate_cache13=intermediate_cache13,
471
+ intermediate_cache2=intermediate_cache2,
472
+ output=output.view(-1, K) if output is not None else output,
473
+ is_k_full=is_k_full,
474
+ )
475
+
476
+ output = output.view(B, BATCH_TOKENS_MAX, K)
477
+
478
+ return output
479
+
480
+
481
+class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
482
def __init__(self, quant_config: FusedMoEQuantConfig):
483
# TODO (varun) : Enable activation quantization
484
assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"
485
super().__init__(quant_config)
486
249
- @override
487
def moe_problem_size(
488
self,
489
a1: torch.Tensor,
@@ -274,6 +511,11 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
511
512
return E, M, N, K, topk
513
514
+
515
+class MarlinExperts(MarlinExpertsBase):
516
+ def __init__(self, quant_config: FusedMoEQuantConfig):
517
+ super().__init__(quant_config)
518
+
519
def supports_expert_map(self) -> bool:
520
return True
521
@@ -365,3 +607,90 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
607
intermediate_cache13=workspace2,
608
intermediate_cache2=workspace13,
609
)
610
+
611
+
612
+class BatchedMarlinExperts(MarlinExpertsBase):
613
+ def __init__(
614
+ self,
615
+ max_num_tokens: int,
616
+ num_dispatchers: int,
617
+ quant_config: FusedMoEQuantConfig,
618
+ ):
619
+ super().__init__(quant_config)
620
+ self.max_num_tokens = max_num_tokens
621
+ self.num_dispatchers = num_dispatchers
622
+
623
+ def supports_expert_map(self) -> bool:
624
+ return True
625
+
626
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
627
+ return TopKWeightAndReduceDelegate()
628
+
629
+ @property
630
+ def activation_formats(
631
+ self,
632
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
633
+ return (
634
+ mk.FusedMoEActivationFormat.BatchedExperts,
635
+ mk.FusedMoEActivationFormat.BatchedExperts,
636
+ )
637
+
638
+ def supports_chunking(self) -> bool:
639
+ return False
640
+
641
+ def workspace_shapes(
642
+ self,
643
+ M: int,
644
+ N: int,
645
+ K: int,
646
+ topk: int,
647
+ global_num_experts: int,
648
+ local_num_experts: int,
649
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
650
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
651
+ num_dispatchers = self.num_dispatchers
652
+ num_experts = local_num_experts
653
+ max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens
654
+ workspace13 = (num_experts * max_num_tokens * num_dispatchers, max(K, N * 2))
655
+ workspace2 = (num_experts * max_num_tokens * num_dispatchers, N)
656
+ output = (num_experts, max_num_tokens * num_dispatchers, K)
657
+ return (workspace13, workspace2, output)
658
+
659
+ def apply(
660
+ self,
661
+ output: torch.Tensor,
662
+ hidden_states: torch.Tensor,
663
+ w1: torch.Tensor,
664
+ w2: torch.Tensor,
665
+ topk_weights: torch.Tensor,
666
+ topk_ids: torch.Tensor,
667
+ activation: str,
668
+ global_num_experts: int,
669
+ expert_map: torch.Tensor | None,
670
+ a1q_scale: torch.Tensor | None,
671
+ a2_scale: torch.Tensor | None,
672
+ workspace13: torch.Tensor,
673
+ workspace2: torch.Tensor,
674
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
675
+ apply_router_weight_on_input: bool,
676
+ ):
677
+ assert expert_tokens_meta is not None, "Num valid tokens per batch is required"
678
+ return batched_fused_marlin_moe(
679
+ hidden_states=hidden_states,
680
+ expert_num_tokens=expert_tokens_meta.expert_num_tokens,
681
+ w1=w1,
682
+ w2=w2,
683
+ bias1=self.w1_bias,
684
+ bias2=self.w2_bias,
685
+ w1_scale=self.w1_scale,
686
+ w2_scale=self.w2_scale,
687
+ gating_output=None,
688
+ quant_type_id=scalar_types.float4_e2m1f.id, # works only for w4a16
689
+ apply_router_weight_on_input=apply_router_weight_on_input,
690
+ global_num_experts=global_num_experts,
691
+ activation=activation,
692
+ expert_map=expert_map,
693
+ output=output,
694
+ intermediate_cache13=workspace13,
695
+ intermediate_cache2=workspace2,
696
+ )
vllm/model_executor/layers/fused_moe/layer.py
+5
@@ -994,6 +994,11 @@ def maybe_roundup_hidden_size(
994
hidden_size, act_dtype
995
)
996
997
+ if moe_parallel_config.use_deepep_ll_kernels:
998
+ hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
999
+ hidden_size
1000
+ )
1001
+
1002
# we are padding globally so EP buffer allocation works
1003
if quant_config and quant_config.get_name() == "mxfp4":
1004
from vllm.model_executor.layers.quantization.mxfp4 import (
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+89
@@ -83,3 +83,92 @@ def moe_align_block_size(
83
expert_ids = expert_map[expert_ids]
84
85
return sorted_ids, expert_ids, num_tokens_post_pad
86
+
87
+
88
+def batched_moe_align_block_size(
89
+ max_tokens_per_batch: int, block_size: int, expert_num_tokens: torch.Tensor
90
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
91
+ """
92
+ Given num_batches, max_tokens_per_batch, block_size and the number of
93
+ valid-tokens in each batch, prepare sorted_token_ids, expert_ids and
94
+ num_tokens_post_pad. sorted_token_ids, expert_ids and num_tokens_post_pad
95
+ have the same semantics as in moe_align_block_size.
96
+
97
+ This function is intended to be a drop in replacement for
98
+ moe_align_batch_size for the batched case.
99
+
100
+ Parameters:
101
+ - max_tokens_per_batch (int): Number of tokens in each batch (both
102
+ valid and invalid).
103
+ - block_size (int): block_size to align the data to.
104
+ - expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates
105
+ the number of valid tokens in batch i.
106
+
107
+ Returns:
108
+ - sorted_token_ids (torch.Tensor): Torch tensor of size
109
+ (num_batches * max_tokens_per_batch) indicating the token indices for
110
+ that block.
111
+ - expert_ids (torch.Tensor): Torch tensor of size
112
+ ceil((num_batches * max_tokens_per_batch) / block_size) indicating
113
+ what expert to use for each block.
114
+ - num_tokens_post_pad (torch.Tensor): Torch tensor of size 1
115
+ indicating the number of valid blocks with actual data to
116
+ process. This is represented in terms of num tokens.
117
+ Example:
118
+ Let num_batches=5, max_tokens_per_batch=8, block_size=4, and
119
+ expert_num_tokens=[2, 3, 0, 6, 8]. This expert_num_tokens tensor
120
+ indicates that,
121
+ - The first 2 tokens in the 0th batch are valid and the rest 6 are
122
+ invalid (i.e. in the 2D hidden_states tensor of shape,
123
+ [num_batches * max_tokens_per_batch, K], indices 0, 1 are valid)
124
+ - The first 3 tokens in the 1st batch are valid. i.e. indices 8, 9, 10
125
+ - 0 tokens in the 2nd batch are valid
126
+ - first 6 tokens in the 3rd batch are valid. i.e. indices,
127
+ 24, 25, 26, 27, 28, 29
128
+ - so on ...
129
+
130
+ In this case,
131
+ sorted_token_ids will be [0, 1, 40, 40,
132
+ 8, 9, 10, 40,
133
+ 24, 25, 26, 27,
134
+ 28, 29, 40, 40,
135
+ 32, 33, 34, 35,
136
+ 36, 37, 38, 39,
137
+ 40, 40, 40, 40,
138
+ (rest all 40, 40, 40, 40)
139
+ ...]
140
+ Here, 40 represents an invalid index. as there is no token index 40.
141
+ The gemm kernel using this sorted_token_ids is expected to skip the
142
+ gemm computation when it encounters this invalid index.
143
+
144
+ expert_ids will be [0, 1, 3, 3, 4, 5, 5, -1, -1, (rest all -1) ...]
145
+ Here, -1 represents an invalid expert. The gemm kernel using this
146
+ expert_ids is expected to skip the gemm computation when it encounters
147
+ an expert of id -1.
148
+
149
+ num_tokens_post_pad will be 24 as sorted_token_ids has valid entries
150
+ until 24.
151
+ """
152
+
153
+ B = expert_num_tokens.size(0)
154
+ device = expert_num_tokens.device
155
+
156
+ # Round up so each batch can be split to blocks evenly.
157
+ max_num_tokens_padded = B * round_up(max_tokens_per_batch, block_size)
158
+
159
+ sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device=device)
160
+ assert max_num_tokens_padded % block_size == 0
161
+ max_num_m_blocks = max_num_tokens_padded // block_size
162
+ expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device)
163
+ num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=device)
164
+
165
+ ops.batched_moe_align_block_size(
166
+ max_tokens_per_batch,
167
+ block_size,
168
+ expert_num_tokens,
169
+ sorted_ids,
170
+ expert_ids,
171
+ num_tokens_post_pad,
172
+ )
173
+
174
+ return sorted_ids, expert_ids, num_tokens_post_pad
vllm/model_executor/layers/quantization/mxfp4.py
+14
-3
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22
ocp_mx_moe_quant_config,
23
)
24
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
25
+ BatchedMarlinExperts,
26
MarlinExperts,
27
fused_marlin_moe,
28
)
@@ -797,9 +798,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
798
prepare_finalize.activation_format
799
== mk.FusedMoEActivationFormat.BatchedExperts
800
):
800
- raise NotImplementedError(
801
- "Mxfp4 does not support batched experts format for EP"
802
- )
801
+ if self.mxfp4_backend == Mxfp4Backend.MARLIN:
802
+ max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
803
+ assert max_num_tokens_per_rank is not None
804
+ assert self.moe_quant_config is not None
805
+ return BatchedMarlinExperts(
806
+ max_num_tokens=max_num_tokens_per_rank,
807
+ num_dispatchers=prepare_finalize.num_dispatchers(),
808
+ quant_config=self.moe_quant_config,
809
+ )
810
+ else:
811
+ raise NotImplementedError(
812
+ "Incompatible Mxfp4 backend for EP batched experts format"
813
+ )
814
else:
815
assert self.moe_quant_config is not None
816
if (