@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 781096e385
Fork 0

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 15, 22:15 UTC

Expert Parallelism (EP) Support for DeepSeek V2 (#12583)

Jongseok Park committed Feb 24, 2025 at 07:33 UTC 781096e385112f799c2b12a9aa2660b85efaf983
19 files changed +527 -59
benchmarks/kernels/benchmark_moe.py
+2 -1
@@ -468,7 +468,8 @@ def main(args: argparse.Namespace):
468 topk = config.num_experts_per_tok
469 intermediate_size = config.intermediate_size
470 shard_intermediate_size = 2 * intermediate_size // args.tp_size
471 - elif config.architectures[0] == "DeepseekV3ForCausalLM":
471 + elif (config.architectures[0] == "DeepseekV3ForCausalLM"
472 + or config.architectures[0] == "DeepseekV2ForCausalLM"):
473 E = config.n_routed_experts
474 topk = config.num_experts_per_tok
475 intermediate_size = config.moe_intermediate_size
tests/distributed/test_expert_parallel.py new
+227
@@ -0,0 +1,227 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from dataclasses import dataclass
4 +from typing import List, Literal, NamedTuple, Optional
5 +
6 +import pytest
7 +
8 +from vllm.config import TaskOption
9 +from vllm.logger import init_logger
10 +
11 +from ..utils import compare_two_settings, fork_new_process_for_each_test
12 +
13 +logger = init_logger("test_expert_parallel")
14 +
15 +
16 +class ParallelSetup(NamedTuple):
17 + tp_size: int
18 + eager_mode: bool
19 + chunked_prefill: bool
20 +
21 +
22 +class EPTestOptions(NamedTuple):
23 + trust_remote_code: bool
24 + tokenizer_mode: Optional[str]
25 + load_format: Optional[str] = None
26 + hf_overrides: Optional[str] = None
27 +
28 +
29 +@dataclass
30 +class EPTestSettings:
31 + parallel_setups: List[ParallelSetup]
32 + distributed_backends: List[str]
33 + task: TaskOption
34 + test_options: EPTestOptions
35 +
36 + @staticmethod
37 + def detailed(
38 + *,
39 + tp_base: int = 2,
40 + task: TaskOption = "auto",
41 + trust_remote_code: bool = False,
42 + tokenizer_mode: Optional[str] = None,
43 + load_format: Optional[str] = None,
44 + hf_overrides: Optional[str] = None,
45 + ):
46 + return EPTestSettings(
47 + parallel_setups=[
48 + ParallelSetup(tp_size=tp_base,
49 + eager_mode=False,
50 + chunked_prefill=False),
51 + ParallelSetup(tp_size=tp_base,
52 + eager_mode=False,
53 + chunked_prefill=True),
54 + ParallelSetup(tp_size=tp_base,
55 + eager_mode=True,
56 + chunked_prefill=False),
57 + ParallelSetup(tp_size=2 * tp_base,
58 + eager_mode=False,
59 + chunked_prefill=True),
60 + ParallelSetup(tp_size=2 * tp_base,
61 + eager_mode=True,
62 + chunked_prefill=False),
63 + ],
64 + distributed_backends=["mp", "ray"],
65 + task=task,
66 + test_options=EPTestOptions(trust_remote_code=trust_remote_code,
67 + tokenizer_mode=tokenizer_mode,
68 + load_format=load_format,
69 + hf_overrides=hf_overrides),
70 + )
71 +
72 + @staticmethod
73 + def fast(
74 + *,
75 + tp_base: int = 2,
76 + task: TaskOption = "auto",
77 + trust_remote_code: bool = False,
78 + tokenizer_mode: Optional[str] = None,
79 + load_format: Optional[str] = None,
80 + hf_overrides: Optional[str] = None,
81 + ):
82 + return EPTestSettings(
83 + parallel_setups=[
84 + ParallelSetup(tp_size=tp_base,
85 + eager_mode=True,
86 + chunked_prefill=False),
87 + ],
88 + distributed_backends=["mp"],
89 + task=task,
90 + test_options=EPTestOptions(trust_remote_code=trust_remote_code,
91 + tokenizer_mode=tokenizer_mode,
92 + load_format=load_format,
93 + hf_overrides=hf_overrides),
94 + )
95 +
96 + def iter_params(self, model_name: str):
97 + opts = self.test_options
98 +
99 + for parallel_setup in self.parallel_setups:
100 + for distributed_backend in self.distributed_backends:
101 + yield (model_name, parallel_setup, distributed_backend,
102 + self.task, opts)
103 +
104 +
105 +# NOTE: You can adjust tp_base locally to fit the model in GPU
106 +# The values displayed here are only a rough indicator of the size of the model
107 +
108 +# yapf: disable
109 +TEST_MODELS = {
110 + "deepseek-ai/DeepSeek-V2-Lite-Chat": EPTestSettings.fast(
111 + trust_remote_code=True),
112 + "mistralai/Mixtral-8x7B-Instruct-v0.1": EPTestSettings.fast(tp_base=4),
113 +}
114 +
115 +
116 +def _compare_tp(
117 + model_name: str,
118 + parallel_setup: ParallelSetup,
119 + distributed_backend: str,
120 + task: TaskOption,
121 + test_options: EPTestOptions,
122 + num_gpus_available: int,
123 + *,
124 + method: Literal["generate"],
125 +):
126 + (
127 + tp_size,
128 + eager_mode,
129 + chunked_prefill,
130 + ) = parallel_setup
131 + (
132 + trust_remote_code,
133 + tokenizer_mode,
134 + load_format,
135 + hf_overrides,
136 + ) = test_options
137 +
138 + if num_gpus_available < tp_size:
139 + pytest.skip(f"Need at least {tp_size} GPUs")
140 +
141 + common_args = [
142 + # use half precision for speed and memory savings in CI environment
143 + "--dtype",
144 + "float16",
145 + "--max-model-len",
146 + "2048",
147 + "--max-num-seqs",
148 + "8",
149 + "--load-format",
150 + "auto",
151 + ]
152 + if chunked_prefill:
153 + common_args.append("--enable-chunked-prefill")
154 + if eager_mode:
155 + common_args.append("--enforce-eager")
156 + if task != "auto":
157 + common_args.extend(["--task", task])
158 + if trust_remote_code:
159 + common_args.append("--trust-remote-code")
160 + if tokenizer_mode:
161 + common_args.extend(["--tokenizer-mode", tokenizer_mode])
162 + if load_format:
163 + common_args.extend(["--load-format", load_format])
164 + if hf_overrides:
165 + common_args.extend(["--hf-overrides", hf_overrides])
166 +
167 + ep_env = {
168 + "VLLM_TEST_ENABLE_EP": "1",
169 + }
170 +
171 + ep_args = [
172 + *common_args,
173 + "--tensor-parallel-size",
174 + str(tp_size),
175 + "--distributed-executor-backend",
176 + distributed_backend,
177 + ]
178 +
179 + # compare without expert parallelism
180 + tp_env = {
181 + "VLLM_TEST_ENABLE_EP": "0",
182 + }
183 +
184 + tp_args = [
185 + *common_args,
186 + "--tensor-parallel-size",
187 + str(tp_size),
188 + "--distributed-executor-backend",
189 + "mp",
190 + ]
191 +
192 + try:
193 + compare_two_settings(model_name,
194 + ep_args,
195 + tp_args,
196 + ep_env,
197 + tp_env,
198 + method=method,
199 + max_wait_seconds=360)
200 + except Exception:
201 + raise
202 +
203 +
204 +@pytest.mark.parametrize(
205 + ("model_name", "parallel_setup", "distributed_backend", "task",
206 + "test_options"),
207 + [
208 + params for model_name, settings in TEST_MODELS.items()
209 + for params in settings.iter_params(model_name)
210 + ],
211 +)
212 +@fork_new_process_for_each_test
213 +def test_ep(
214 + model_name: str,
215 + parallel_setup: ParallelSetup,
216 + distributed_backend: str,
217 + task: TaskOption,
218 + test_options: EPTestOptions,
219 + num_gpus_available,
220 +):
221 + _compare_tp(model_name,
222 + parallel_setup,
223 + distributed_backend,
224 + task,
225 + test_options,
226 + num_gpus_available,
227 + method="generate")
tests/kernels/test_awq_marlin.py
+2 -7
@@ -99,13 +99,8 @@ def test_fused_marlin_moe_awq(
99 num_bits=num_bits,
100 )
101
102 - torch_output = torch_moe(
103 - a,
104 - w_ref1.transpose(1, 2),
105 - w_ref2.transpose(1, 2),
106 - score,
107 - topk,
108 - )
102 + torch_output = torch_moe(a, w_ref1.transpose(1, 2), w_ref2.transpose(1, 2),
103 + score, topk, None)
104
105 assert compute_max_diff(marlin_output, torch_output) < 4e-2
106
tests/kernels/test_moe.py
+59 -6
@@ -26,6 +26,7 @@ from vllm.platforms import current_platform
26 from vllm.scalar_type import scalar_types
27
28 NUM_EXPERTS = [8, 64]
29 +EP_SIZE = [1, 4]
30 TOP_KS = [2, 6]
31
32
@@ -34,6 +35,7 @@ TOP_KS = [2, 6]
35 @pytest.mark.parametrize("k", [128, 511, 1024])
36 @pytest.mark.parametrize("e", NUM_EXPERTS)
37 @pytest.mark.parametrize("topk", TOP_KS)
38 +@pytest.mark.parametrize("ep_size", EP_SIZE)
39 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
40 def test_fused_moe(
41 m: int,
@@ -41,6 +43,7 @@ def test_fused_moe(
43 k: int,
44 e: int,
45 topk: int,
46 + ep_size: int,
47 dtype: torch.dtype,
48 ):
49 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
@@ -48,10 +51,38 @@ def test_fused_moe(
51 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
52
53 score = torch.randn((m, e), device="cuda", dtype=dtype)
51 - triton_output = fused_moe(a, w1, w2, score, topk, renormalize=False)
52 - torch_output = torch_moe(a, w1, w2, score, topk)
54 +
55 + if ep_size > 1:
56 + local_e = e // ep_size
57 + e_ids = torch.randint(0,
58 + e, (local_e, ),
59 + device="cuda",
60 + dtype=torch.int32)
61 + e_map = torch.full((e, ), -1, device="cuda", dtype=torch.int32)
62 + e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
63 + w1 = w1[e_ids]
64 + w2 = w2[e_ids]
65 + else:
66 + e_map = None
67 +
68 + triton_output = fused_moe(a,
69 + w1,
70 + w2,
71 + score,
72 + topk,
73 + global_num_experts=e,
74 + expert_map=e_map,
75 + renormalize=False)
76 + torch_output = torch_moe(a, w1, w2, score, topk, e_map)
77 torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
54 - iterative_output = iterative_moe(a, w1, w2, score, topk, renormalize=False)
78 + iterative_output = iterative_moe(a,
79 + w1,
80 + w2,
81 + score,
82 + topk,
83 + global_num_experts=e,
84 + expert_map=e_map,
85 + renormalize=False)
86 torch.testing.assert_close(iterative_output,
87 torch_output,
88 atol=2e-2,
@@ -63,13 +94,14 @@ def test_fused_moe(
94 @pytest.mark.parametrize("k", [128, 1024])
95 @pytest.mark.parametrize("e", NUM_EXPERTS)
96 @pytest.mark.parametrize("topk", TOP_KS)
97 +@pytest.mark.parametrize("ep_size", EP_SIZE)
98 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
99 @pytest.mark.parametrize("group_size", [64, 128])
100 @pytest.mark.parametrize("has_zp", [True, False])
101 @pytest.mark.parametrize("weight_bits", [4, 8])
102 def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
71 - dtype: torch.dtype, group_size: int, has_zp: bool,
72 - weight_bits: int):
103 + ep_size: int, dtype: torch.dtype, group_size: int,
104 + has_zp: bool, weight_bits: int):
105 print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
106 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
107 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
@@ -130,6 +162,25 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
162 if has_zp:
163 w_qzeros[expert_id] = qzeros
164
165 + if ep_size > 1:
166 + local_e = e // ep_size
167 + e_ids = torch.randint(0,
168 + e, (local_e, ),
169 + device="cuda",
170 + dtype=torch.int32)
171 + e_map = torch.full((e, ), -1, device="cuda", dtype=torch.int32)
172 + e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
173 + w1_ref = w1_ref[e_ids]
174 + w2_ref = w2_ref[e_ids]
175 + w1_qweight = w1_qweight[e_ids]
176 + w2_qweight = w2_qweight[e_ids]
177 + w1_scales = w1_scales[e_ids]
178 + w2_scales = w2_scales[e_ids]
179 + w1_qzeros = w1_qzeros[e_ids]
180 + w2_qzeros = w2_qzeros[e_ids]
181 + else:
182 + e_map = None
183 +
184 triton_output = fused_moe(a,
185 w1_qweight,
186 w2_qweight,
@@ -138,12 +189,14 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
189 renormalize=False,
190 use_int4_w4a16=weight_bits == 4,
191 use_int8_w8a16=weight_bits == 8,
192 + global_num_experts=e,
193 + expert_map=e_map,
194 w1_scale=w1_scales,
195 w2_scale=w2_scales,
196 w1_zp=w1_qzeros if has_zp else None,
197 w2_zp=w2_qzeros if has_zp else None,
198 block_shape=[0, group_size])
146 - torch_output = torch_moe(a, w1_ref, w2_ref, score, topk)
199 + torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map)
200 torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
201
202
tests/kernels/utils.py
+3 -1
@@ -1053,7 +1053,7 @@ def compute_max_diff(output, output_ref):
1053 torch.abs(output_ref))
1054
1055
1056 -def torch_moe(a, w1, w2, score, topk):
1056 +def torch_moe(a, w1, w2, score, topk, expert_map):
1057 B, D = a.shape
1058 a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
1059 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
@@ -1061,6 +1061,8 @@ def torch_moe(a, w1, w2, score, topk):
1061 topk_weight, topk_ids = torch.topk(score, topk)
1062 topk_weight = topk_weight.view(-1)
1063 topk_ids = topk_ids.view(-1)
1064 + if expert_map is not None:
1065 + topk_ids = expert_map[topk_ids]
1066 for i in range(w1.shape[0]):
1067 mask = topk_ids == i
1068 if mask.sum():
tests/utils.py
+3 -3
@@ -297,12 +297,12 @@ def _test_completion_close(
297 logprobs=5,
298 temperature=0.0)
299
300 - logporbs = completion.choices[0].logprobs.top_logprobs[0]
301 - logporbs = {k: round(v, 2) for k, v in logporbs.items()}
300 + logprobs = completion.choices[0].logprobs.top_logprobs[0]
301 + logprobs = {k: round(v, 2) for k, v in logprobs.items()}
302
303 results.append({
304 "test": "completion_close",
305 - "logprobs": logporbs,
305 + "logprobs": logprobs,
306 })
307
308 return results
vllm/config.py
+20
@@ -677,6 +677,23 @@ class ModelConfig:
677 "fallback to the eager mode.")
678 self.enforce_eager = True
679
680 + def _verify_with_expert_parallelism(self) -> None:
681 + num_expert_names = [
682 + "moe_num_experts", # Dbrx
683 + "num_experts", # Jamba
684 + "n_routed_experts", # DeepSeek
685 + "num_local_experts", # Mixtral
686 + ]
687 + num_experts = 0
688 + for name in num_expert_names:
689 + num_experts = getattr(self.hf_text_config, name, 0)
690 + if num_experts > 0:
691 + break
692 + if num_experts < 1:
693 + raise ValueError(
694 + "Number of experts in the model must be greater than 0 "
695 + "when expert parallelism is enabled.")
696 +
697 def verify_async_output_proc(self, parallel_config, speculative_config,
698 device_config) -> None:
699 if not self.use_async_output_proc:
@@ -730,6 +747,9 @@ class ModelConfig:
747 " must be divisible by tensor parallel size "
748 f"({tensor_parallel_size}).")
749
750 + if envs.VLLM_TEST_ENABLE_EP:
751 + self._verify_with_expert_parallelism()
752 +
753 pipeline_parallel_size = parallel_config.pipeline_parallel_size
754 if pipeline_parallel_size > 1:
755 architectures = getattr(self.hf_config, "architectures", [])
vllm/envs.py
+7
@@ -86,6 +86,7 @@ if TYPE_CHECKING:
86 VLLM_MLA_PERFORM_MATRIX_ABSORPTION: bool = True
87 VLLM_MLA_DISABLE_REQUANTIZATION: bool = False
88 VLLM_MLA_CUDA_MEM_ALIGN_KV_CACHE: bool = True
89 + VLLM_TEST_ENABLE_EP: bool = False
90 VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON: bool = False
91 VLLM_RAY_PER_WORKER_GPUS: float = 1.0
92 VLLM_RAY_BUNDLE_INDICES: str = ""
@@ -570,6 +571,12 @@ environment_variables: Dict[str, Callable[[], Any]] = {
571 lambda: bool(int(os.getenv("VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON", "0"))
572 ),
573
574 + # If set, vLLM will use the experimental expert parallel implementation on
575 + # the FusedMoE layer, using tensor parallelism size as expert parallelism
576 + # size.
577 + "VLLM_TEST_ENABLE_EP":
578 + lambda: bool(int(os.getenv("VLLM_TEST_ENABLE_EP", "0"))),
579 +
580 # Number of GPUs per worker in Ray, if it is set to be a fraction,
581 # it allows ray to schedule multiple actors on a single GPU,
582 # so that users can colocate other actors on the same GPUs as vLLM.
vllm/model_executor/layers/fused_moe/fused_moe.py
+96 -30
@@ -20,6 +20,18 @@ from vllm.utils import direct_register_custom_op
20 logger = init_logger(__name__)
21
22
23 +@triton.jit
24 +def write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N, offs_token,
25 + token_mask, BLOCK_SIZE_M, BLOCK_SIZE_N,
26 + compute_type):
27 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type)
28 + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
29 + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[
30 + None, :]
31 + c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
32 + tl.store(c_ptrs, accumulator, mask=c_mask)
33 +
34 +
35 @triton.jit
36 def fused_moe_kernel_gptq_awq(
37 # Pointers to matrices
@@ -120,17 +132,26 @@ def fused_moe_kernel_gptq_awq(
132 offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
133 token_mask = offs_token < num_valid_tokens
134
135 + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
136 + if off_experts == -1:
137 + # -----------------------------------------------------------
138 + # Write back zeros to the output when the expert is not
139 + # in the current expert parallel rank.
140 + write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N,
141 + offs_token, token_mask, BLOCK_SIZE_M,
142 + BLOCK_SIZE_N, compute_type)
143 + return
144 +
145 offs_bn = (pid_n * BLOCK_SIZE_N +
146 tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
147 offs_k = tl.arange(0, BLOCK_SIZE_K)
148 a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
149 offs_k[None, :] * stride_ak)
150
129 - off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
130 -
151 if use_int4_w4a16:
152 b_ptrs = b_ptr + off_experts * stride_be + \
133 - (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * stride_bn
153 + (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * \
154 + stride_bn
155 b_shifter = (offs_k[:, None] % 2) * 4
156 elif use_int8_w8a16:
157 b_ptrs = b_ptr + off_experts * stride_be + \
@@ -170,7 +191,8 @@ def fused_moe_kernel_gptq_awq(
191
192 b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + \
193 offs_bn[None, :] * stride_bsn + \
173 - ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * stride_bsk
194 + ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * \
195 + stride_bsk
196 b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other)
197 b_scale = b_scale.to(tl.float32)
198
@@ -319,13 +341,22 @@ def fused_moe_kernel(
341 offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
342 token_mask = offs_token < num_valid_tokens
343
344 + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
345 + if off_experts == -1:
346 + # -----------------------------------------------------------
347 + # Write back zeros to the output when the expert is not
348 + # in the current expert parallel rank.
349 + write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N,
350 + offs_token, token_mask, BLOCK_SIZE_M,
351 + BLOCK_SIZE_N, compute_type)
352 + return
353 +
354 offs_bn = (pid_n * BLOCK_SIZE_N +
355 tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
356 offs_k = tl.arange(0, BLOCK_SIZE_K)
357 a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
358 offs_k[None, :] * stride_ak)
359
328 - off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
360 b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk +
361 offs_bn[None, :] * stride_bn)
362 if use_int8_w8a16:
@@ -349,7 +380,6 @@ def fused_moe_kernel(
380 # of fp32 values for higher accuracy.
381 # `accumulator` will be converted back to fp16 after the loop.
382 accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
352 -
383 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
384 # Load the next block of A and B, generate a mask by checking the
385 # K dimension.
@@ -544,8 +574,11 @@ def moe_align_block_size_triton(
574
575
576 def moe_align_block_size(
547 - topk_ids: torch.Tensor, block_size: int,
548 - num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
577 + topk_ids: torch.Tensor,
578 + block_size: int,
579 + num_experts: int,
580 + expert_map: torch.Tensor = None
581 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
582 """
583 Aligns the token distribution across experts to be compatible with block
584 size for matrix multiplication.
@@ -555,6 +588,10 @@ def moe_align_block_size(
588 top-k expert indices for each token.
589 - block_size: The block size used in block matrix multiplication.
590 - num_experts: The total number of experts.
591 + - expert_map: A tensor of shape [num_experts] that maps the expert index
592 + from the global space to the local index space of the current
593 + expert parallel shard. If the expert is not in the current expert
594 + parallel shard, the mapping is set to -1.
595
596 Returns:
597 - sorted_token_ids: A tensor containing the sorted token indices according
@@ -589,7 +626,9 @@ def moe_align_block_size(
626 device=topk_ids.device)
627 sorted_ids.fill_(topk_ids.numel())
628 max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
592 - expert_ids = torch.empty((max_num_m_blocks, ),
629 + # Expert ids must be zeroed out to prevent index out of bounds error while
630 + # mapping global expert ids to local expert ids in expert parallelism.
631 + expert_ids = torch.zeros((max_num_m_blocks, ),
632 dtype=torch.int32,
633 device=topk_ids.device)
634 num_tokens_post_pad = torch.empty((1),
@@ -618,6 +657,9 @@ def moe_align_block_size(
657 else:
658 ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
659 expert_ids, num_tokens_post_pad)
660 + if expert_map is not None:
661 + expert_ids = expert_map[expert_ids]
662 +
663 return sorted_ids, expert_ids, num_tokens_post_pad
664
665
@@ -1001,6 +1043,8 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
1043 use_fp8_w8a8: bool = False,
1044 use_int8_w8a16: bool = False,
1045 use_int4_w4a16: bool = False,
1046 + global_num_experts: int = -1,
1047 + expert_map: Optional[torch.Tensor] = None,
1048 w1_scale: Optional[torch.Tensor] = None,
1049 w2_scale: Optional[torch.Tensor] = None,
1050 w1_zp: Optional[torch.Tensor] = None,
@@ -1009,8 +1053,9 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
1053 a2_scale: Optional[torch.Tensor] = None,
1054 block_shape: Optional[List[int]] = None) -> None:
1055 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
1012 - use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, w1_scale,
1013 - w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)
1056 + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16,
1057 + global_num_experts, expert_map, w1_scale, w2_scale,
1058 + w1_zp, w2_zp, a1_scale, a2_scale, block_shape)
1059
1060
1061 def inplace_fused_experts_fake(
@@ -1022,6 +1067,8 @@ def inplace_fused_experts_fake(
1067 use_fp8_w8a8: bool = False,
1068 use_int8_w8a16: bool = False,
1069 use_int4_w4a16: bool = False,
1070 + global_num_experts: int = -1,
1071 + expert_map: Optional[torch.Tensor] = None,
1072 w1_scale: Optional[torch.Tensor] = None,
1073 w2_scale: Optional[torch.Tensor] = None,
1074 w1_zp: Optional[torch.Tensor] = None,
@@ -1049,6 +1096,8 @@ def outplace_fused_experts(
1096 use_fp8_w8a8: bool = False,
1097 use_int8_w8a16: bool = False,
1098 use_int4_w4a16: bool = False,
1099 + global_num_experts: int = -1,
1100 + expert_map: Optional[torch.Tensor] = None,
1101 w1_scale: Optional[torch.Tensor] = None,
1102 w2_scale: Optional[torch.Tensor] = None,
1103 w1_zp: Optional[torch.Tensor] = None,
@@ -1058,8 +1107,9 @@ def outplace_fused_experts(
1107 block_shape: Optional[List[int]] = None) -> torch.Tensor:
1108 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
1109 False, use_fp8_w8a8, use_int8_w8a16,
1061 - use_int4_w4a16, w1_scale, w2_scale, w1_zp, w2_zp,
1062 - a1_scale, a2_scale, block_shape)
1110 + use_int4_w4a16, global_num_experts, expert_map,
1111 + w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
1112 + a2_scale, block_shape)
1113
1114
1115 def outplace_fused_experts_fake(
@@ -1071,6 +1121,8 @@ def outplace_fused_experts_fake(
1121 use_fp8_w8a8: bool = False,
1122 use_int8_w8a16: bool = False,
1123 use_int4_w4a16: bool = False,
1124 + global_num_experts: int = -1,
1125 + expert_map: Optional[torch.Tensor] = None,
1126 w1_scale: Optional[torch.Tensor] = None,
1127 w2_scale: Optional[torch.Tensor] = None,
1128 w1_zp: Optional[torch.Tensor] = None,
@@ -1098,26 +1150,27 @@ def fused_experts(hidden_states: torch.Tensor,
1150 use_fp8_w8a8: bool = False,
1151 use_int8_w8a16: bool = False,
1152 use_int4_w4a16: bool = False,
1153 + global_num_experts: int = -1,
1154 + expert_map: Optional[torch.Tensor] = None,
1155 w1_scale: Optional[torch.Tensor] = None,
1156 w2_scale: Optional[torch.Tensor] = None,
1157 w1_zp: Optional[torch.Tensor] = None,
1158 w2_zp: Optional[torch.Tensor] = None,
1159 a1_scale: Optional[torch.Tensor] = None,
1160 a2_scale: Optional[torch.Tensor] = None,
1107 - block_shape: Optional[List[int]] = None):
1161 + block_shape: Optional[List[int]] = None) -> torch.Tensor:
1162 +
1163 if inplace:
1109 - torch.ops.vllm.inplace_fused_experts(hidden_states, w1, w2,
1110 - topk_weights, topk_ids,
1111 - use_fp8_w8a8, use_int8_w8a16,
1112 - use_int4_w4a16, w1_scale,
1113 - w2_scale, w1_zp, w2_zp, a1_scale,
1114 - a2_scale, block_shape)
1164 + torch.ops.vllm.inplace_fused_experts(
1165 + hidden_states, w1, w2, topk_weights, topk_ids, use_fp8_w8a8,
1166 + use_int8_w8a16, use_int4_w4a16, global_num_experts, expert_map,
1167 + w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)
1168 return hidden_states
1169 else:
1170 return torch.ops.vllm.outplace_fused_experts(
1171 hidden_states, w1, w2, topk_weights, topk_ids, use_fp8_w8a8,
1119 - use_int8_w8a16, use_int4_w4a16, w1_scale, w2_scale, w1_zp, w2_zp,
1120 - a1_scale, a2_scale, block_shape)
1172 + use_int8_w8a16, use_int4_w4a16, global_num_experts, expert_map,
1173 + w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)
1174
1175
1176 def fused_experts_impl(hidden_states: torch.Tensor,
@@ -1129,6 +1182,8 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1182 use_fp8_w8a8: bool = False,
1183 use_int8_w8a16: bool = False,
1184 use_int4_w4a16: bool = False,
1185 + global_num_experts: int = -1,
1186 + expert_map: Optional[torch.Tensor] = None,
1187 w1_scale: Optional[torch.Tensor] = None,
1188 w2_scale: Optional[torch.Tensor] = None,
1189 w1_zp: Optional[torch.Tensor] = None,
@@ -1153,6 +1208,9 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1208
1209 num_tokens, _ = hidden_states.shape
1210 E, N, _ = w1.shape
1211 + if global_num_experts == -1:
1212 + global_num_experts = E
1213 + top_k_num = topk_ids.shape[1]
1214 # We execute the fused_moe kernel in chunks to circumvent this issue:
1215 # https://github.com/vllm-project/vllm/issues/5938
1216 CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
@@ -1166,20 +1224,20 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1224 try_get_optimal_moe_config,
1225 w1.shape,
1226 w2.shape,
1169 - topk_ids.shape[1],
1227 + top_k_num,
1228 config_dtype,
1229 block_shape=block_shape,
1230 )
1231
1232 config = get_config_func(M)
1233
1176 - intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
1234 + intermediate_cache1 = torch.empty((M, top_k_num, N),
1235 device=hidden_states.device,
1236 dtype=hidden_states.dtype)
1179 - intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
1237 + intermediate_cache2 = torch.empty((M * top_k_num, N // 2),
1238 device=hidden_states.device,
1239 dtype=hidden_states.dtype)
1182 - intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
1240 + intermediate_cache3 = torch.empty((M, top_k_num, w2.shape[1]),
1241 device=hidden_states.device,
1242 dtype=hidden_states.dtype)
1243
@@ -1221,7 +1279,8 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1279 curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1280
1281 sorted_token_ids, expert_ids, num_tokens_post_padded = (
1224 - moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
1282 + moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
1283 + global_num_experts, expert_map))
1284
1285 invoke_fused_moe_kernel(curr_hidden_states,
1286 w1,
@@ -1235,7 +1294,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1294 expert_ids,
1295 num_tokens_post_padded,
1296 False,
1238 - topk_ids.shape[1],
1297 + top_k_num,
1298 config,
1299 compute_type=compute_type,
1300 use_fp8_w8a8=use_fp8_w8a8,
@@ -1286,6 +1345,8 @@ def fused_moe(
1345 use_fp8_w8a8: bool = False,
1346 use_int8_w8a16: bool = False,
1347 use_int4_w4a16: bool = False,
1348 + global_num_experts: int = -1,
1349 + expert_map: Optional[torch.Tensor] = None,
1350 w1_scale: Optional[torch.Tensor] = None,
1351 w2_scale: Optional[torch.Tensor] = None,
1352 w1_zp: Optional[torch.Tensor] = None,
@@ -1320,6 +1381,11 @@ def fused_moe(
1381 - use_int4_w4a16 (bool): If True, use matmul of int4 weight and bf16/fp16
1382 activation to compute the inner products for w1 and w2.
1383 Defaults to False.
1384 + - global_num_experts (int): The total number of experts in the global
1385 + expert space.
1386 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
1387 + from the global expert space to the local expert space of the expert
1388 + parallel shard.
1389 - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
1390 w1.
1391 - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
@@ -1334,8 +1400,6 @@ def fused_moe(
1400 Returns:
1401 - torch.Tensor: The output tensor after applying the MoE layer.
1402 """
1337 - # Check constraints.
1338 - assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
1403
1404 if use_grouped_topk:
1405 assert num_expert_group is not None and topk_group is not None
@@ -1358,6 +1422,8 @@ def fused_moe(
1422 use_fp8_w8a8=use_fp8_w8a8,
1423 use_int8_w8a16=use_int8_w8a16,
1424 use_int4_w4a16=use_int4_w4a16,
1425 + global_num_experts=global_num_experts,
1426 + expert_map=expert_map,
1427 w1_scale=w1_scale,
1428 w2_scale=w2_scale,
1429 w1_zp=w1_zp,
vllm/model_executor/layers/fused_moe/layer.py
+65 -5
@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Tuple
6
7 import torch
8
9 +import vllm.envs as envs
10 from vllm.distributed import (get_tensor_model_parallel_rank,
11 get_tensor_model_parallel_world_size,
12 tensor_model_parallel_all_reduce)
@@ -55,6 +56,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):
56 use_grouped_topk: bool = False,
57 topk_group: Optional[int] = None,
58 num_expert_group: Optional[int] = None,
59 + global_num_experts: int = -1,
60 + expert_map: Optional[torch.Tensor] = None,
61 custom_routing_function: Optional[Callable] = None,
62 scoring_func: str = "softmax",
63 e_score_correction_bias: Optional[torch.Tensor] = None
@@ -113,6 +116,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
116 use_grouped_topk: bool = False,
117 topk_group: Optional[int] = None,
118 num_expert_group: Optional[int] = None,
119 + global_num_experts: int = -1,
120 + expert_map: Optional[torch.Tensor] = None,
121 custom_routing_function: Optional[Callable] = None,
122 scoring_func: str = "softmax",
123 e_score_correction_bias: Optional[torch.Tensor] = None
@@ -125,6 +130,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
130 use_grouped_topk=use_grouped_topk,
131 topk_group=topk_group,
132 num_expert_group=num_expert_group,
133 + global_num_experts=global_num_experts,
134 + expert_map=expert_map,
135 custom_routing_function=custom_routing_function,
136 scoring_func=scoring_func,
137 e_score_correction_bias=e_score_correction_bias)
@@ -139,6 +146,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
146 renormalize: bool,
147 topk_group: Optional[int] = None,
148 num_expert_group: Optional[int] = None,
149 + global_num_experts: int = -1,
150 + expert_map: Optional[torch.Tensor] = None,
151 custom_routing_function: Optional[Callable] = None,
152 scoring_func: str = "softmax",
153 e_score_correction_bias: Optional[torch.Tensor] = None
@@ -160,7 +169,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
169 w2=layer.w2_weight,
170 topk_weights=topk_weights,
171 topk_ids=topk_ids,
163 - inplace=True)
172 + inplace=True,
173 + global_num_experts=global_num_experts,
174 + expert_map=expert_map)
175
176 def forward_cpu(
177 self,
@@ -172,6 +183,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
183 renormalize: bool,
184 topk_group: Optional[int] = None,
185 num_expert_group: Optional[int] = None,
186 + global_num_experts: int = -1,
187 + expert_map: Optional[torch.Tensor] = None,
188 custom_routing_function: Optional[Callable] = None,
189 **kwargs,
190 ):
@@ -196,6 +209,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
209 renormalize: bool,
210 topk_group: Optional[int] = None,
211 num_expert_group: Optional[int] = None,
212 + global_num_experts: int = -1,
213 + expert_map: Optional[torch.Tensor] = None,
214 custom_routing_function: Optional[Callable] = None,
215 scoring_func: str = "softmax",
216 e_score_correction_bias: Optional[torch.Tensor] = None
@@ -215,6 +230,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
230 w2=layer.w2_weight,
231 topk=top_k,
232 gating_output=router_logits,
233 + global_num_experts=global_num_experts,
234 + expert_map=expert_map,
235 renormalize=renormalize)
236
237 forward_native = forward_cuda
@@ -255,6 +272,7 @@ class FusedMoE(torch.nn.Module):
272 topk_group: Optional[int] = None,
273 quant_config: Optional[QuantizationConfig] = None,
274 tp_size: Optional[int] = None,
275 + ep_size: Optional[int] = None,
276 prefix: str = "",
277 custom_routing_function: Optional[Callable] = None,
278 scoring_func: str = "softmax",
@@ -267,8 +285,13 @@ class FusedMoE(torch.nn.Module):
285
286 self.tp_size = (tp_size if tp_size is not None else
287 get_tensor_model_parallel_world_size())
288 + if envs.VLLM_TEST_ENABLE_EP:
289 + self.ep_size = self.tp_size
290 + self.tp_size = 1
291 + else:
292 + self.ep_size = 1
293 self.top_k = top_k
271 - self.num_experts = num_experts
294 + self.num_experts = num_experts # Global number of experts
295 assert intermediate_size % self.tp_size == 0
296 self.intermediate_size_per_partition = intermediate_size // self.tp_size
297 self.reduce_results = reduce_results
@@ -281,6 +304,26 @@ class FusedMoE(torch.nn.Module):
304 self.custom_routing_function = custom_routing_function
305 self.scoring_func = scoring_func
306 self.e_score_correction_bias = e_score_correction_bias
307 + self.expert_map = None
308 +
309 + if self.ep_size > 1:
310 + # Create a tensor of size num_experts filled with -1
311 + self.expert_map = torch.full((self.num_experts, ),
312 + -1,
313 + dtype=torch.int32)
314 + # Create a expert map for the local experts
315 + local_num_experts = num_experts // self.ep_size
316 + ep_rank = get_tensor_model_parallel_rank()
317 + if ep_rank < (self.ep_size - 1):
318 + # Each non-last rank gets local_num_experts experts.
319 + self.expert_map[ep_rank * local_num_experts:
320 + (ep_rank + 1) * local_num_experts] = \
321 + torch.arange(0, local_num_experts, dtype=torch.int32)
322 + else:
323 + # All remaining experts are assigned to the last rank.
324 + local_num_experts = num_experts - ep_rank * local_num_experts
325 + self.expert_map[-local_num_experts:] = \
326 + torch.arange(0, local_num_experts, dtype=torch.int32)
327
328 if self.scoring_func != "softmax" and not self.use_grouped_topk:
329 raise ValueError("Only softmax scoring function is supported for "
@@ -293,8 +336,11 @@ class FusedMoE(torch.nn.Module):
336 self.quant_method = quant_config.get_quant_method(self, prefix)
337 assert self.quant_method is not None
338
339 + local_num_experts = torch.sum(self.expert_map != -1) \
340 + if self.expert_map is not None else num_experts
341 +
342 moe_quant_params = {
297 - "num_experts": num_experts,
343 + "num_experts": local_num_experts,
344 "hidden_size": hidden_size,
345 "intermediate_size_per_partition":
346 self.intermediate_size_per_partition,
@@ -423,10 +469,22 @@ class FusedMoE(torch.nn.Module):
469 assert shard_id in ("w1", "w3")
470 expert_data.copy_(loaded_weight)
471
472 + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
473 + if self.expert_map is None:
474 + return expert_id
475 + return self.expert_map[expert_id].item()
476 +
477 def weight_loader(self, param: torch.nn.Parameter,
478 loaded_weight: torch.Tensor, weight_name: str,
479 shard_id: str, expert_id: int) -> None:
480
481 + expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
482 + if expert_id == -1:
483 + return
484 +
485 + # TP rank is set to 0 if EP is enabled
486 + tp_rank = 0 if self.ep_size > 1 else get_tensor_model_parallel_rank()
487 +
488 # compressed-tensors checkpoints with packed weights are stored flipped
489 # TODO (mgoin): check self.quant_method.quant_config.quant_format
490 # against known CompressionFormat enum values that have this quality
@@ -447,7 +505,6 @@ class FusedMoE(torch.nn.Module):
505 SHARD_ID_TO_SHARDED_DIM = {"w1": 0, "w2": 1, "w3": 0}
506
507 expert_data = param.data[expert_id]
450 - tp_rank = get_tensor_model_parallel_rank()
508
509 # is_transposed: if the dim to shard the weight
510 # should be flipped. Required by GPTQ, compressed-tensors
@@ -590,13 +647,16 @@ class FusedMoE(torch.nn.Module):
647 top_k=self.top_k,
648 renormalize=self.renormalize,
649 use_grouped_topk=self.use_grouped_topk,
650 + global_num_experts=self.num_experts,
651 + expert_map=self.expert_map,
652 topk_group=self.topk_group,
653 num_expert_group=self.num_expert_group,
654 custom_routing_function=self.custom_routing_function,
655 scoring_func=self.scoring_func,
656 e_score_correction_bias=self.e_score_correction_bias)
657
599 - if self.reduce_results and self.tp_size > 1:
658 + if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
659 + # Default set to False. (May have to add shared expert outputs.)
660 final_hidden_states = tensor_model_parallel_all_reduce(
661 final_hidden_states)
662
vllm/model_executor/layers/fused_moe/moe_torch_iterative.py
+8 -2
@@ -10,7 +10,9 @@ def fused_moe(
10 w2: torch.Tensor,
11 gating_output: torch.Tensor,
12 topk: int,
13 - renormalize: bool,
13 + global_num_experts: int,
14 + expert_map: torch.Tensor = None,
15 + renormalize: bool = False,
16 ) -> torch.Tensor:
17 """
18 Args:
@@ -18,6 +20,7 @@ def fused_moe(
20 w1: [num_experts, intermediate_size * 2, hidden_size]
21 w2: [num_experts, hidden_size, intermediate_size]
22 gating_output: [*, num_experts]
23 + expert_map: [num_experts]
24 """
25 orig_shape = hidden_states.shape
26 hidden_size = hidden_states.shape[-1]
@@ -27,13 +30,16 @@ def fused_moe(
30 dtype = hidden_states.dtype
31
32 hidden_states = hidden_states.view(num_tokens, hidden_size)
30 - gating_output = gating_output.view(num_tokens, num_experts)
33 + gating_output = gating_output.view(num_tokens, global_num_experts)
34 topk_weights = gating_output.softmax(dim=-1, dtype=torch.float)
35 topk_weights, selected_experts = topk_weights.topk(topk, dim=-1)
36 if renormalize:
37 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
38 topk_weights = topk_weights.to(dtype)
39
40 + if expert_map is not None:
41 + selected_experts = expert_map[selected_experts]
42 +
43 final_hidden_states = None
44 for expert_idx in range(num_experts):
45 expert_w1 = w1[expert_idx]
vllm/model_executor/layers/quantization/awq_marlin.py
+7
@@ -464,10 +464,17 @@ class AWQMoEMethod(FusedMoEMethodBase):
464 use_grouped_topk: bool = False,
465 topk_group: Optional[int] = None,
466 num_expert_group: Optional[int] = None,
467 + global_num_experts: int = -1,
468 + expert_map: Optional[torch.Tensor] = None,
469 custom_routing_function: Optional[Callable] = None,
470 scoring_func: str = "softmax",
471 e_score_correction_bias: Optional[torch.Tensor] = None,
472 ) -> torch.Tensor:
473 + if expert_map is not None:
474 + raise NotImplementedError(
475 + "Expert Parallelism is not supported for "
476 + "fused Marlin MoE method.")
477 +
478 topk_weights, topk_ids = FusedMoE.select_experts(
479 hidden_states=x,
480 router_logits=router_logits,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+10
@@ -214,6 +214,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
214 use_grouped_topk: bool = False,
215 topk_group: Optional[int] = None,
216 num_expert_group: Optional[int] = None,
217 + global_num_experts: int = -1,
218 + expert_map: Optional[torch.Tensor] = None,
219 custom_routing_function: Optional[Callable] = None,
220 scoring_func: str = "softmax",
221 e_score_correction_bias: Optional[torch.Tensor] = None,
@@ -239,6 +241,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
241 topk_ids=topk_ids,
242 inplace=True,
243 use_fp8_w8a8=True,
244 + global_num_experts=global_num_experts,
245 + expert_map=expert_map,
246 w1_scale=layer.w13_weight_scale,
247 w2_scale=layer.w2_weight_scale,
248 a1_scale=layer.w13_input_scale,
@@ -540,10 +544,16 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
544 use_grouped_topk: bool = False,
545 topk_group: Optional[int] = None,
546 num_expert_group: Optional[int] = None,
547 + global_num_experts: int = -1,
548 + expert_map: Optional[torch.Tensor] = None,
549 custom_routing_function: Optional[Callable] = None,
550 scoring_func: str = "softmax",
551 e_score_correction_bias: Optional[torch.Tensor] = None,
552 ) -> torch.Tensor:
553 + if expert_map is not None:
554 + raise NotImplementedError(
555 + "Expert Parallelism is not supported for "
556 + "fused Marlin MoE method.")
557
558 topk_weights, topk_ids = FusedMoE.select_experts(
559 hidden_states=x,
vllm/model_executor/layers/quantization/experts_int8.py
+4
@@ -108,6 +108,8 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
108 use_grouped_topk: bool = False,
109 topk_group: Optional[int] = None,
110 num_expert_group: Optional[int] = None,
111 + global_num_experts: int = -1,
112 + expert_map: Optional[torch.Tensor] = None,
113 custom_routing_function: Optional[Callable] = None,
114 scoring_func: str = "softmax",
115 e_score_correction_bias: Optional[torch.Tensor] = None,
@@ -133,6 +135,8 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
135 topk_ids=topk_ids,
136 inplace=True,
137 use_int8_w8a16=True,
138 + global_num_experts=global_num_experts,
139 + expert_map=expert_map,
140 w1_scale=layer.w13_scale,
141 w2_scale=layer.w2_scale)
142
vllm/model_executor/layers/quantization/fp8.py
+4
@@ -670,6 +670,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
670 use_grouped_topk: bool = False,
671 topk_group: Optional[int] = None,
672 num_expert_group: Optional[int] = None,
673 + global_num_experts: int = -1,
674 + expert_map: Optional[torch.Tensor] = None,
675 custom_routing_function: Optional[Callable] = None,
676 scoring_func: str = "softmax",
677 e_score_correction_bias: Optional[torch.Tensor] = None,
@@ -697,6 +699,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
699 topk_ids=topk_ids,
700 inplace=True,
701 use_fp8_w8a8=True,
702 + global_num_experts=global_num_experts,
703 + expert_map=expert_map,
704 w1_scale=(layer.w13_weight_scale_inv
705 if self.block_quant else layer.w13_weight_scale),
706 w2_scale=(layer.w2_weight_scale_inv
vllm/model_executor/layers/quantization/gptq_marlin.py
+2
@@ -585,6 +585,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
585 use_grouped_topk: bool = False,
586 topk_group: Optional[int] = None,
587 num_expert_group: Optional[int] = None,
588 + global_num_experts: int = -1,
589 + expert_map: Optional[torch.Tensor] = None,
590 custom_routing_function: Optional[Callable] = None,
591 scoring_func: str = "softmax",
592 e_score_correction_bias: Optional[torch.Tensor] = None,
vllm/model_executor/layers/quantization/moe_wna16.py
+4
@@ -288,6 +288,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
288 use_grouped_topk: bool = False,
289 topk_group: Optional[int] = None,
290 num_expert_group: Optional[int] = None,
291 + global_num_experts: int = -1,
292 + expert_map: Optional[torch.Tensor] = None,
293 custom_routing_function: Optional[Callable] = None,
294 scoring_func: str = "softmax",
295 e_score_correction_bias: Optional[torch.Tensor] = None,
@@ -317,6 +319,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
319 inplace=True,
320 use_int4_w4a16=weight_bits == 4,
321 use_int8_w8a16=weight_bits == 8,
322 + global_num_experts=global_num_experts,
323 + expert_map=expert_map,
324 w1_scale=layer.w13_scales,
325 w2_scale=layer.w2_scales,
326 w1_zp=layer.w13_qzeros if has_zp else None,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+4
@@ -198,6 +198,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
198 use_grouped_topk: bool = False,
199 topk_group: Optional[int] = None,
200 num_expert_group: Optional[int] = None,
201 + global_num_experts: int = -1,
202 + expert_map: Optional[torch.Tensor] = None,
203 custom_routing_function: Optional[Callable] = None,
204 scoring_func: str = "softmax",
205 e_score_correction_bias: Optional[torch.Tensor] = None,
@@ -223,6 +225,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
225 topk_ids=topk_ids,
226 inplace=True,
227 use_fp8_w8a8=True,
228 + global_num_experts=global_num_experts,
229 + expert_map=expert_map,
230 w1_scale=layer.w13_weight_scale,
231 w2_scale=layer.w2_weight_scale,
232 a1_scale=layer.w13_input_scale,
vllm/model_executor/models/deepseek_v2.py
-4
@@ -106,10 +106,6 @@ class DeepseekV2MoE(nn.Module):
106 self.routed_scaling_factor = config.routed_scaling_factor
107 self.n_shared_experts = config.n_shared_experts
108 self.routed_scaling_factor = config.routed_scaling_factor
109 - if self.tp_size > config.n_routed_experts:
110 - raise ValueError(
111 - f"Tensor parallel size {self.tp_size} is greater than "
112 - f"the number of experts {config.n_routed_experts}.")
109
110 if config.hidden_act != "silu":
111 raise ValueError(f"Unsupported activation: {config.hidden_act}. "