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 4, 17:30 UTC
[MoE] Move remaining PrepareAndFinalize to prepare finalize folder (#39009)
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Signed-off-by: Jackmin801 <ongjackm@gmail.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
Jackmin801 committed
Apr 23, 2026 at 17:00 UTC
447c372ac504a4696ddb54da8821a07065ef3faf
9 files changed
+184
-163
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+1
-1
@@ -223,7 +223,7 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
223
)
224
225
if has_mori():
226
- from vllm.model_executor.layers.fused_moe.mori_prepare_finalize import (
226
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
227
MoriPrepareAndFinalize,
228
)
229
tests/kernels/moe/test_batched_deepgemm.py
+3
-1
@@ -10,10 +10,12 @@ from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
10
BatchedDeepGemmExperts,
11
)
12
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
13
- BatchedPrepareAndFinalize,
13
BatchedTritonExperts,
14
)
15
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
16
+from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
17
+ BatchedPrepareAndFinalize,
18
+)
19
from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported
20
21
from .test_deepgemm import make_block_quant_fp8_weights
tests/kernels/moe/utils.py
+3
-1
@@ -18,7 +18,6 @@ from vllm.model_executor.layers.fused_moe.config import (
18
RoutingMethodType,
19
)
20
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
21
- BatchedPrepareAndFinalize,
21
BatchedTritonExperts,
22
NaiveBatchedExperts,
23
)
@@ -27,6 +26,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
26
fused_experts,
27
)
28
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
29
+from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
30
+ BatchedPrepareAndFinalize,
31
+)
32
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
33
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
34
from vllm.utils.deep_gemm import per_block_cast_to_fp8
vllm/model_executor/layers/fused_moe/all2all_utils.py
+2
-2
@@ -41,9 +41,9 @@ if current_platform.is_cuda_alike():
41
DeepEPLLPrepareAndFinalize,
42
)
43
if has_mori():
44
- from .mori_prepare_finalize import MoriPrepareAndFinalize
44
+ from .prepare_finalize.mori import MoriPrepareAndFinalize
45
if has_nixl_ep():
46
- from .nixl_ep_prepare_finalize import (
46
+ from .prepare_finalize.nixl_ep import (
47
NIXL_EP_QUANT_BLOCK_SHAPE,
48
NixlEPPrepareAndFinalize,
49
)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
-158
@@ -14,13 +14,11 @@ from vllm.model_executor.layers.fused_moe.config import (
14
from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
15
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16
TopKWeightAndReduceDelegate,
17
- TopKWeightAndReduceNaiveBatched,
17
)
18
from vllm.model_executor.layers.fused_moe.utils import (
19
_resize_cache,
20
moe_kernel_quantize_input,
21
normalize_batched_scales_shape,
23
- normalize_scales_shape,
22
)
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
QuantKey,
@@ -489,162 +487,6 @@ def invoke_moe_batched_triton_kernel(
487
)
488
489
492
-class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
493
- """
494
- A reference prepare/finalize class that reorganizes the tokens into
495
- expert batched format, i.e. E x max_num_tokens x K. This is the format
496
- that the batched dispatch/combine kernels use.
497
- """
498
-
499
- def __init__(
500
- self,
501
- max_num_tokens: int,
502
- num_local_experts: int,
503
- num_dispatchers: int,
504
- rank: int,
505
- ):
506
- super().__init__()
507
- self.max_num_tokens = max_num_tokens
508
- self.num_local_experts = num_local_experts
509
- self.rank = rank
510
- self.num_dispatchers_ = num_dispatchers
511
-
512
- @property
513
- def activation_format(self) -> mk.FusedMoEActivationFormat:
514
- return mk.FusedMoEActivationFormat.BatchedExperts
515
-
516
- def max_num_tokens_per_rank(self) -> int | None:
517
- return self.max_num_tokens
518
-
519
- def topk_indices_dtype(self) -> torch.dtype | None:
520
- return None
521
-
522
- def num_dispatchers(self) -> int:
523
- return self.num_dispatchers_
524
-
525
- def output_is_reduced(self) -> bool:
526
- return False
527
-
528
- def prepare(
529
- self,
530
- a1: torch.Tensor,
531
- topk_weights: torch.Tensor,
532
- topk_ids: torch.Tensor,
533
- num_experts: int,
534
- expert_map: torch.Tensor | None,
535
- apply_router_weight_on_input: bool,
536
- quant_config: FusedMoEQuantConfig,
537
- defer_input_quant: bool = False,
538
- ) -> mk.PrepareResultType:
539
- if defer_input_quant:
540
- raise NotImplementedError(
541
- f"{self.__class__.__name__} does not support defer_input_quant=True. "
542
- "Please select an MoE kernel that accepts quantized inputs."
543
- )
544
- assert a1.dim() == 2
545
- assert topk_ids.dim() == 2
546
- assert topk_ids.size(0) == a1.size(0)
547
-
548
- if apply_router_weight_on_input:
549
- topk = topk_ids.size(1)
550
- # TODO: this only works for topK=1, will need to update for topK>1
551
- assert topk == 1, (
552
- "apply_router_weight_on_input is only implemented for topk=1"
553
- )
554
- a1.mul_(topk_weights.to(a1.dtype))
555
-
556
- num_tokens, hidden_dim = a1.size()
557
- topk = topk_ids.size(1)
558
-
559
- tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device)
560
-
561
- num_local_experts = self.num_local_experts
562
-
563
- if quant_config.quant_dtype is None:
564
- b_type = a1.dtype
565
- else:
566
- b_type = quant_config.quant_dtype
567
-
568
- b_a1 = torch.zeros(
569
- (num_local_experts, self.max_num_tokens, hidden_dim),
570
- dtype=b_type,
571
- device=a1.device,
572
- )
573
-
574
- if quant_config.is_quantized:
575
- scale_shape = quant_config.batched_scale_shape(
576
- num_local_experts, self.max_num_tokens, hidden_dim
577
- )
578
-
579
- b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device)
580
- else:
581
- assert quant_config.a1_scale is None
582
- b_a1_scale = None
583
-
584
- first_expert = num_local_experts * self.rank
585
- last_expert = first_expert + num_local_experts
586
-
587
- a1_scale = normalize_scales_shape(quant_config.a1_scale)
588
-
589
- for expert_id in range(first_expert, last_expert):
590
- topks = torch.any(topk_ids == expert_id, dim=1).flatten()
591
- rows = torch.count_nonzero(topks.flatten())
592
- if rows == 0:
593
- continue
594
- idx = expert_id - first_expert
595
- tokens_per_expert[idx] = rows
596
- rhs = a1[: topks.numel()][topks]
597
- if quant_config.quant_dtype is not None:
598
- if a1_scale is not None:
599
- if quant_config.is_per_act_token:
600
- rhs_a1_scale = a1_scale[: topks.numel()][topks]
601
- else:
602
- rhs_a1_scale = a1_scale
603
- else:
604
- rhs_a1_scale = None
605
- b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input(
606
- rhs,
607
- rhs_a1_scale,
608
- quant_config.quant_dtype,
609
- quant_config.per_act_token_quant,
610
- quant_config.block_shape,
611
- )
612
- assert b_s is not None
613
- if quant_config.is_per_act_token:
614
- b_a1_scale[idx, :rows] = b_s[:rows]
615
- else:
616
- b_a1_scale[idx, : b_s.shape[0]] = b_s
617
- else:
618
- b_a1[idx, :rows, :] = rhs
619
-
620
- assert b_a1_scale is None or b_a1_scale.ndim == 3
621
-
622
- expert_tokens_meta = mk.ExpertTokensMetadata(
623
- expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
624
- )
625
-
626
- return b_a1, b_a1_scale, expert_tokens_meta, None, None
627
-
628
- def finalize(
629
- self,
630
- output: torch.Tensor,
631
- fused_expert_output: torch.Tensor,
632
- topk_weights: torch.Tensor,
633
- topk_ids: torch.Tensor,
634
- apply_router_weight_on_input: bool,
635
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
636
- ) -> None:
637
- if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
638
- weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
639
- weight_and_reduce_impl.apply(
640
- output=output,
641
- fused_expert_output=fused_expert_output,
642
- topk_weights=topk_weights,
643
- topk_ids=topk_ids,
644
- apply_router_weight_on_input=apply_router_weight_on_input,
645
- )
646
-
647
-
490
class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
491
"""
492
A reference MoE expert class that operates on expert batched format,
vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py
+4
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
5
+ BatchedPrepareAndFinalize,
6
+)
7
from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import (
8
MoEPrepareAndFinalizeNaiveDPEPModular,
9
MoEPrepareAndFinalizeNaiveDPEPMonolithic,
@@ -13,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import (
16
)
17
18
__all__ = [
19
+ "BatchedPrepareAndFinalize",
20
"MoEPrepareAndFinalizeNaiveDPEPMonolithic",
21
"MoEPrepareAndFinalizeNaiveDPEPModular",
22
"make_moe_prepare_and_finalize_naive_dp_ep",
vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py
new
+171
@@ -0,0 +1,171 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
8
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
9
+ TopKWeightAndReduceDelegate,
10
+ TopKWeightAndReduceNaiveBatched,
11
+)
12
+from vllm.model_executor.layers.fused_moe.utils import (
13
+ moe_kernel_quantize_input,
14
+ normalize_scales_shape,
15
+)
16
+
17
+
18
+class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
19
+ """
20
+ A reference prepare/finalize class that reorganizes the tokens into
21
+ expert batched format, i.e. E x max_num_tokens x K. This is the format
22
+ that the batched dispatch/combine kernels use.
23
+ """
24
+
25
+ def __init__(
26
+ self,
27
+ max_num_tokens: int,
28
+ num_local_experts: int,
29
+ num_dispatchers: int,
30
+ rank: int,
31
+ ):
32
+ super().__init__()
33
+ self.max_num_tokens = max_num_tokens
34
+ self.num_local_experts = num_local_experts
35
+ self.rank = rank
36
+ self.num_dispatchers_ = num_dispatchers
37
+
38
+ @property
39
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
40
+ return mk.FusedMoEActivationFormat.BatchedExperts
41
+
42
+ def max_num_tokens_per_rank(self) -> int | None:
43
+ return self.max_num_tokens
44
+
45
+ def topk_indices_dtype(self) -> torch.dtype | None:
46
+ return None
47
+
48
+ def num_dispatchers(self) -> int:
49
+ return self.num_dispatchers_
50
+
51
+ def output_is_reduced(self) -> bool:
52
+ return False
53
+
54
+ def prepare(
55
+ self,
56
+ a1: torch.Tensor,
57
+ topk_weights: torch.Tensor,
58
+ topk_ids: torch.Tensor,
59
+ num_experts: int,
60
+ expert_map: torch.Tensor | None,
61
+ apply_router_weight_on_input: bool,
62
+ quant_config: FusedMoEQuantConfig,
63
+ defer_input_quant: bool = False,
64
+ ) -> mk.PrepareResultType:
65
+ if defer_input_quant:
66
+ raise NotImplementedError(
67
+ f"{self.__class__.__name__} does not support defer_input_quant=True. "
68
+ "Please select an MoE kernel that accepts quantized inputs."
69
+ )
70
+ assert a1.dim() == 2
71
+ assert topk_ids.dim() == 2
72
+ assert topk_ids.size(0) == a1.size(0)
73
+
74
+ if apply_router_weight_on_input:
75
+ topk = topk_ids.size(1)
76
+ # TODO: this only works for topK=1, will need to update for topK>1
77
+ assert topk == 1, (
78
+ "apply_router_weight_on_input is only implemented for topk=1"
79
+ )
80
+ a1.mul_(topk_weights.to(a1.dtype))
81
+
82
+ num_tokens, hidden_dim = a1.size()
83
+ topk = topk_ids.size(1)
84
+
85
+ tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device)
86
+
87
+ num_local_experts = self.num_local_experts
88
+
89
+ if quant_config.quant_dtype is None:
90
+ b_type = a1.dtype
91
+ else:
92
+ b_type = quant_config.quant_dtype
93
+
94
+ b_a1 = torch.zeros(
95
+ (num_local_experts, self.max_num_tokens, hidden_dim),
96
+ dtype=b_type,
97
+ device=a1.device,
98
+ )
99
+
100
+ if quant_config.is_quantized:
101
+ scale_shape = quant_config.batched_scale_shape(
102
+ num_local_experts, self.max_num_tokens, hidden_dim
103
+ )
104
+
105
+ b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device)
106
+ else:
107
+ assert quant_config.a1_scale is None
108
+ b_a1_scale = None
109
+
110
+ first_expert = num_local_experts * self.rank
111
+ last_expert = first_expert + num_local_experts
112
+
113
+ a1_scale = normalize_scales_shape(quant_config.a1_scale)
114
+
115
+ for expert_id in range(first_expert, last_expert):
116
+ topks = torch.any(topk_ids == expert_id, dim=1).flatten()
117
+ rows = torch.count_nonzero(topks.flatten())
118
+ if rows == 0:
119
+ continue
120
+ idx = expert_id - first_expert
121
+ tokens_per_expert[idx] = rows
122
+ rhs = a1[: topks.numel()][topks]
123
+ if quant_config.quant_dtype is not None:
124
+ if a1_scale is not None:
125
+ if quant_config.is_per_act_token:
126
+ rhs_a1_scale = a1_scale[: topks.numel()][topks]
127
+ else:
128
+ rhs_a1_scale = a1_scale
129
+ else:
130
+ rhs_a1_scale = None
131
+ b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input(
132
+ rhs,
133
+ rhs_a1_scale,
134
+ quant_config.quant_dtype,
135
+ quant_config.per_act_token_quant,
136
+ quant_config.block_shape,
137
+ )
138
+ assert b_s is not None
139
+ if quant_config.is_per_act_token:
140
+ b_a1_scale[idx, :rows] = b_s[:rows]
141
+ else:
142
+ b_a1_scale[idx, : b_s.shape[0]] = b_s
143
+ else:
144
+ b_a1[idx, :rows, :] = rhs
145
+
146
+ assert b_a1_scale is None or b_a1_scale.ndim == 3
147
+
148
+ expert_tokens_meta = mk.ExpertTokensMetadata(
149
+ expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
150
+ )
151
+
152
+ return b_a1, b_a1_scale, expert_tokens_meta, None, None
153
+
154
+ def finalize(
155
+ self,
156
+ output: torch.Tensor,
157
+ fused_expert_output: torch.Tensor,
158
+ topk_weights: torch.Tensor,
159
+ topk_ids: torch.Tensor,
160
+ apply_router_weight_on_input: bool,
161
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
162
+ ) -> None:
163
+ if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
164
+ weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
165
+ weight_and_reduce_impl.apply(
166
+ output=output,
167
+ fused_expert_output=fused_expert_output,
168
+ topk_weights=topk_weights,
169
+ topk_ids=topk_ids,
170
+ apply_router_weight_on_input=apply_router_weight_on_input,
171
+ )