@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 015fab8c2f
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, 04:45 UTC

[Kernels][Bugfix] Use torch op for all kernels in FusedMoE forward. Add additional testing for cudagraphs. (#19717)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Jun 25, 2025 at 02:22 UTC 015fab8c2fa4db8776f7e91abd50371911673d88
14 files changed +381 -240
tests/kernels/moe/test_cutlass_moe.py
+10 -1
@@ -29,7 +29,10 @@ MNK_FACTORS = [
29 (224, 1024, 1536),
30 (224, 3072, 1024),
31 (224, 3072, 1536),
32 - (1024 * 128, 1024, 1024),
32 + (32768, 1024, 1024),
33 + # These sizes trigger wrong answers.
34 + #(7232, 2048, 5120),
35 + #(40000, 2048, 5120),
36 ]
37
38 vllm_config = VllmConfig(parallel_config=ParallelConfig(
@@ -232,8 +235,10 @@ def test_cutlass_moe_8_bit_no_graph(
235 topk: int,
236 per_act_token: bool,
237 per_out_ch: bool,
238 + monkeypatch,
239 ):
240 current_platform.seed_everything(7)
241 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
242 with set_current_vllm_config(vllm_config):
243 mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
244 per_out_ch)
@@ -274,8 +279,10 @@ def test_cutlass_moe_8_bit_cuda_graph(
279 topk: int,
280 per_act_token: bool,
281 per_out_ch: bool,
282 + monkeypatch,
283 ):
284 current_platform.seed_everything(7)
285 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
286 with set_current_vllm_config(vllm_config):
287 dtype = torch.half
288
@@ -329,8 +336,10 @@ def test_cutlass_moe_8_bit_EP(
336 per_act_token: bool,
337 per_out_channel: bool,
338 ep_size: int,
339 + monkeypatch,
340 ):
341 current_platform.seed_everything(7)
342 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
343 with set_current_vllm_config(vllm_config):
344 mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
345 per_out_channel)
tests/kernels/moe/test_moe.py
+202 -90
@@ -4,6 +4,9 @@
4
5 Run `pytest tests/kernels/test_moe.py`.
6 """
7 +import functools
8 +from typing import Callable, Optional, Union
9 +
10 import pytest
11 import torch
12 from torch.nn import Parameter
@@ -14,6 +17,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
17 import vllm.model_executor.layers.fused_moe # noqa
18 from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
19 from vllm.config import VllmConfig, set_current_vllm_config
20 +from vllm.forward_context import set_forward_context
21 from vllm.model_executor.layers.fused_moe import fused_moe
22 from vllm.model_executor.layers.fused_moe.fused_moe import (
23 fused_topk, modular_triton_fused_moe)
@@ -40,7 +44,76 @@ vllm_config.scheduler_config.max_num_seqs = 128
44 vllm_config.scheduler_config.max_model_len = 8192
45
46
43 -@pytest.mark.parametrize("m", [1, 33, 64, 222, 1024 * 128])
47 +def run_moe_test(
48 + baseline: Union[Callable, torch.Tensor],
49 + moe_fn: Callable,
50 + a: torch.Tensor,
51 + w1: torch.Tensor,
52 + w2: torch.Tensor,
53 + score: torch.Tensor,
54 + topk: int,
55 + global_num_experts: int = -1,
56 + expert_map: Optional[torch.Tensor] = None,
57 + padding: bool = False,
58 + use_compile: bool = False,
59 + use_cudagraph: bool = False,
60 + atol: float = 2e-2,
61 + rtol: float = 0,
62 +) -> torch.Tensor:
63 + if isinstance(baseline, torch.Tensor):
64 + baseline_output = baseline
65 + else:
66 + baseline_output = baseline(a,
67 + w1,
68 + w2,
69 + score,
70 + topk,
71 + global_num_experts=global_num_experts,
72 + expert_map=expert_map)
73 +
74 + # Pad the weight if moe padding is enabled
75 + if padding:
76 + w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
77 + w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
78 +
79 + if use_compile:
80 + moe_fn = torch.compile(moe_fn, backend="inductor", fullgraph=True)
81 + torch._dynamo.mark_dynamic(a, 0)
82 + torch._dynamo.mark_dynamic(score, 0)
83 +
84 + test_output = moe_fn(a,
85 + w1,
86 + w2,
87 + score,
88 + topk,
89 + global_num_experts=global_num_experts,
90 + expert_map=expert_map)
91 +
92 + if use_cudagraph:
93 + test_output.fill_(0)
94 + stream = torch.cuda.Stream()
95 + graph = torch.cuda.CUDAGraph()
96 + with torch.cuda.graph(graph, stream=stream):
97 + test_output = moe_fn(a,
98 + w1,
99 + w2,
100 + score,
101 + topk,
102 + global_num_experts=global_num_experts,
103 + expert_map=expert_map)
104 + torch.cuda.synchronize()
105 + graph.replay()
106 + torch.cuda.synchronize()
107 +
108 + torch.testing.assert_close(test_output,
109 + baseline_output,
110 + atol=atol,
111 + rtol=rtol)
112 +
113 + return baseline_output
114 +
115 +
116 +@pytest.mark.parametrize("m", [1, 33, 64, 222, 32768, 40000])
117 @pytest.mark.parametrize("n", [128, 1024, 2048])
118 @pytest.mark.parametrize("k", [128, 511, 1024])
119 @pytest.mark.parametrize("e", NUM_EXPERTS)
@@ -48,6 +121,7 @@ vllm_config.scheduler_config.max_model_len = 8192
121 @pytest.mark.parametrize("ep_size", EP_SIZE)
122 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
123 @pytest.mark.parametrize("padding", [True, False])
124 +@pytest.mark.parametrize("chunk_size", [8192])
125 def test_fused_moe(
126 m: int,
127 n: int,
@@ -57,7 +131,17 @@ def test_fused_moe(
131 ep_size: int,
132 dtype: torch.dtype,
133 padding: bool,
134 + chunk_size: int,
135 + monkeypatch,
136 ):
137 + current_platform.seed_everything(7)
138 +
139 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
140 +
141 + #
142 + # Setup test data
143 + #
144 +
145 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
146 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
147 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
@@ -77,58 +161,70 @@ def test_fused_moe(
161 else:
162 e_map = None
163
80 - m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=False,
81 - use_int8_w8a8=False,
82 - use_int8_w8a16=False,
83 - use_int4_w4a16=False,
84 - per_channel_quant=False,
85 - block_shape=None)
86 -
87 - with set_current_vllm_config(vllm_config):
88 - torch_output = torch_moe(a, w1, w2, score, topk, e_map)
89 - iterative_output = iterative_moe(a,
90 - w1,
91 - w2,
92 - score,
93 - topk,
94 - global_num_experts=e,
95 - expert_map=e_map,
96 - renormalize=False)
97 -
98 - # Pad the weight if moe padding is enabled
99 - if padding:
100 - w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
101 - torch.cuda.empty_cache()
102 - w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
103 - torch.cuda.empty_cache()
164 + #
165 + # Setup test functions
166 + #
167 +
168 + m_fused_moe_fn = modular_triton_fused_moe(use_fp8_w8a8=False,
169 + use_int8_w8a8=False,
170 + use_int8_w8a16=False,
171 + use_int4_w4a16=False,
172 + per_channel_quant=False,
173 + block_shape=None)
174 +
175 + def m_fused_moe(
176 + a: torch.Tensor,
177 + w1: torch.Tensor,
178 + w2: torch.Tensor,
179 + score: torch.Tensor,
180 + topk: int,
181 + global_num_experts: int = -1,
182 + expert_map: Optional[torch.Tensor] = None,
183 + ) -> torch.Tensor:
184 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
185 + return m_fused_moe_fn(a,
186 + w1,
187 + w2,
188 + topk_weights,
189 + topk_ids,
190 + global_num_experts=global_num_experts,
191 + expert_map=expert_map)
192 +
193 + fused_moe_fn = functools.partial(fused_moe, renormalize=False)
194 +
195 + #
196 + # Run tests
197 + #
198 + runner = functools.partial(
199 + run_moe_test,
200 + a=a,
201 + w1=w1,
202 + w2=w2,
203 + score=score,
204 + topk=topk,
205 + global_num_experts=e,
206 + expert_map=e_map,
207 + padding=padding,
208 + )
209
105 - triton_output = fused_moe(a,
106 - w1,
107 - w2,
108 - score,
109 - topk,
110 - global_num_experts=e,
111 - expert_map=e_map,
112 - renormalize=False)
210 + # Note: for now use_compile will error out if the problem size is
211 + # large enough to trigger chunking. I'm leaving the flag and
212 + # setup code in case we are able to revisit this later.
213 + use_compile = False
214
114 - topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
115 - m_triton_output = m_fused_moe(a,
116 - w1,
117 - w2,
118 - topk_weights,
119 - topk_ids,
120 - global_num_experts=e,
121 - expert_map=e_map)
215 + use_cudagraph = (n >= 1024 and k >= 1024
216 + and current_platform.is_cuda_alike())
217
123 - torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
124 - torch.testing.assert_close(m_triton_output,
125 - torch_output,
126 - atol=2e-2,
127 - rtol=0)
128 - torch.testing.assert_close(iterative_output,
129 - torch_output,
130 - atol=2e-2,
131 - rtol=0)
218 + with set_current_vllm_config(vllm_config):
219 + baseline_output = runner(torch_moe, iterative_moe)
220 + runner(baseline_output,
221 + fused_moe_fn,
222 + use_compile=use_compile,
223 + use_cudagraph=use_cudagraph)
224 + runner(baseline_output,
225 + m_fused_moe,
226 + use_compile=use_compile,
227 + use_cudagraph=use_cudagraph)
228
229
230 @pytest.mark.parametrize("m", [1, 32, 222])
@@ -238,7 +334,12 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
334 w1_zp=w1_qzeros if has_zp else None,
335 w2_zp=w2_qzeros if has_zp else None,
336 block_shape=[0, group_size])
241 - torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map)
337 + torch_output = torch_moe(a,
338 + w1_ref,
339 + w2_ref,
340 + score,
341 + topk,
342 + expert_map=e_map)
343
344 torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
345
@@ -265,45 +366,51 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
366 pytest.skip("AITER ROCm test skip for float32")
367
368 # Instantiate our and huggingface's MoE blocks
268 - config = MixtralConfig()
269 - hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
270 - vllm_moe = MixtralMoE(
271 - num_experts=config.num_local_experts,
272 - top_k=config.num_experts_per_tok,
273 - hidden_size=config.hidden_size,
274 - intermediate_size=config.intermediate_size,
275 - params_dtype=dtype,
276 - tp_size=1,
277 - dp_size=1,
278 - ).cuda()
279 -
280 - # Load the weights
281 - vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
282 - for i in range(config.num_local_experts):
283 - weights = (hf_moe.experts[i].w1.weight.data,
284 - hf_moe.experts[i].w3.weight.data)
285 - vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
286 - vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
287 -
288 - # Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
289 - hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
290 - # vLLM uses 1D query [num_tokens, hidden_dim]
291 - vllm_inputs = hf_inputs.flatten(0, 1)
369 + vllm_config.compilation_config.static_forward_context = dict()
370 + with (set_current_vllm_config(vllm_config),
371 + set_forward_context(None, vllm_config)):
372 + config = MixtralConfig()
373 + hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
374 + vllm_moe = MixtralMoE(
375 + num_experts=config.num_local_experts,
376 + top_k=config.num_experts_per_tok,
377 + hidden_size=config.hidden_size,
378 + intermediate_size=config.intermediate_size,
379 + params_dtype=dtype,
380 + tp_size=1,
381 + dp_size=1,
382 + ).cuda()
383 +
384 + # Load the weights
385 + vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
386 + for i in range(config.num_local_experts):
387 + weights = (hf_moe.experts[i].w1.weight.data,
388 + hf_moe.experts[i].w3.weight.data)
389 + vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
390 + vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
391 +
392 + # Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
393 + hf_inputs = torch.randn(
394 + (1, 64, config.hidden_size)).to(dtype).to("cuda")
395 + # vLLM uses 1D query [num_tokens, hidden_dim]
396 + vllm_inputs = hf_inputs.flatten(0, 1)
397
293 - # Pad the weight if moe padding is enabled
294 - if padding:
295 - vllm_moe.experts.w13_weight = Parameter(F.pad(
296 - vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[..., 0:-128],
297 - requires_grad=False)
298 - torch.cuda.empty_cache()
299 - vllm_moe.experts.w2_weight = Parameter(F.pad(
300 - vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
301 - requires_grad=False)
302 - torch.cuda.empty_cache()
303 -
304 - # Run forward passes for both MoE blocks
305 - hf_states, _ = hf_moe.forward(hf_inputs)
306 - vllm_states = vllm_moe.forward(vllm_inputs)
398 + # Pad the weight if moe padding is enabled
399 + if padding:
400 + vllm_moe.experts.w13_weight = Parameter(F.pad(
401 + vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[...,
402 + 0:-128],
403 + requires_grad=False)
404 + torch.cuda.empty_cache()
405 + vllm_moe.experts.w2_weight = Parameter(F.pad(
406 + vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[...,
407 + 0:-128],
408 + requires_grad=False)
409 + torch.cuda.empty_cache()
410 +
411 + # Run forward passes for both MoE blocks
412 + hf_states, _ = hf_moe.forward(hf_inputs)
413 + vllm_states = vllm_moe.forward(vllm_inputs)
414
415 mixtral_moe_tol = {
416 torch.float32: 1e-3,
@@ -546,7 +653,12 @@ def test_fused_marlin_moe(
653 topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
654
655 with set_current_vllm_config(vllm_config):
549 - torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
656 + torch_output = torch_moe(a,
657 + w_ref1,
658 + w_ref2,
659 + score,
660 + topk,
661 + expert_map=e_map)
662
663 marlin_output = torch.ops.vllm.fused_marlin_moe(
664 a,
tests/kernels/moe/test_nvfp4_moe.py
+1 -1
@@ -136,7 +136,7 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
136 device=w2.device,
137 block_size=quant_blocksize)
138
139 - torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk, None)
139 + torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk)
140
141 torch.testing.assert_close(torch_output,
142 cutlass_output,
tests/kernels/moe/test_pplx_cutlass_moe.py
+3 -19
@@ -6,9 +6,9 @@ from typing import Optional
6 import pytest
7 import torch
8
9 +from tests.kernels.utils import torch_experts
10 from vllm import _custom_ops as ops
11 from vllm.config import VllmConfig, set_current_vllm_config
11 -from vllm.model_executor.layers.activation import SiluAndMul
12 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
13 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -164,22 +164,6 @@ vllm_config.scheduler_config.max_num_seqs = 128
164 vllm_config.scheduler_config.max_model_len = 8192
165
166
167 -def torch_moe2(a, w1, w2, topk_weight, topk_ids):
168 - M, K = a.shape
169 - topk = topk_ids.shape[1]
170 - a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
171 - out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
172 - num_experts = w1.shape[0]
173 - for i in range(num_experts):
174 - mask = (topk_ids == i).view(-1)
175 - if mask.sum():
176 - out[mask] = SiluAndMul()(
177 - a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
178 -
179 - return (out.view(M, -1, w2.shape[1]) *
180 - topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
181 -
182 -
167 def _pplx_moe(
168 pgi: ProcessGroupInfo,
169 dp_size: int,
@@ -210,8 +194,8 @@ def _pplx_moe(
194 group_name = cpu_group.group_name
195
196 with set_current_vllm_config(vllm_config):
213 - torch_output = torch_moe2(a_full, w1_full, w2_full, topk_weights,
214 - topk_ids)
197 + torch_output = torch_experts(a_full, w1_full, w2_full, topk_weights,
198 + topk_ids)
199 pplx_output = pplx_cutlass_moe(pgi, dp_size, a, w1, w2, w1_scale,
200 w2_scale, topk_weights, topk_ids,
201 a1_scale, out_dtype, per_act_token,
tests/kernels/moe/test_pplx_moe.py
+10 -27
@@ -18,8 +18,8 @@ try:
18 except ImportError:
19 has_pplx = False
20
21 +from tests.kernels.utils import torch_experts
22 from vllm.config import VllmConfig, set_current_vllm_config
22 -from vllm.model_executor.layers.activation import SiluAndMul
23 from vllm.model_executor.layers.fused_moe import override_config
24 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
25 BatchedExperts, BatchedPrepareAndFinalize, BatchedTritonExperts)
@@ -163,29 +163,6 @@ def batched_moe(
163 return fused_experts(a, w1, w2, topk_weight, topk_ids, num_experts)
164
165
166 -# Note: same as torch_moe but with fused_topk factored out.
167 -def torch_moe2(
168 - a: torch.Tensor,
169 - w1: torch.Tensor,
170 - w2: torch.Tensor,
171 - topk_weight: torch.Tensor,
172 - topk_ids: torch.Tensor,
173 -) -> torch.Tensor:
174 - M, K = a.shape
175 - topk = topk_ids.shape[1]
176 - a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
177 - out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
178 - num_experts = w1.shape[0]
179 - for i in range(num_experts):
180 - mask = (topk_ids == i).view(-1)
181 - if mask.sum():
182 - out[mask] = SiluAndMul()(
183 - a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
184 -
185 - return (out.view(M, -1, w2.shape[1]) *
186 - topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
187 -
188 -
166 @pytest.mark.parametrize("m", [1, 33, 64, 222])
167 @pytest.mark.parametrize("n", [128, 1024, 2048])
168 @pytest.mark.parametrize("k", [128, 512, 1024])
@@ -209,7 +186,7 @@ def test_fused_moe_batched_experts(
186
187 with set_current_vllm_config(vllm_config):
188 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
212 - baseline_output = torch_moe2(a, w1, w2, topk_weight, topk_ids)
189 + baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
190 torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids)
191 batched_output = batched_moe(a, w1, w2, topk_weight, topk_ids)
192
@@ -409,7 +386,7 @@ def pplx_moe(
386 w2: torch.Tensor,
387 topk_weight: torch.Tensor,
388 topk_ids: torch.Tensor,
412 - use_compile: bool = True,
389 + use_compile: bool = False,
390 use_cudagraphs: bool = True,
391 ) -> torch.Tensor:
392 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
@@ -470,10 +447,16 @@ def pplx_moe(
447 w1_chunk = chunk_by_rank(w1, rank, world_size).to(device)
448 w2_chunk = chunk_by_rank(w2, rank, world_size).to(device)
449
450 + # Note: for now use_compile will error out if the problem size is
451 + # large enough to trigger chunking. I'm leaving the flag and
452 + # setup code in case we are able to revisit this later.
453 if use_compile:
454 _fused_experts = torch.compile(fused_experts,
455 backend='inductor',
456 fullgraph=True)
457 + torch._dynamo.mark_dynamic(a_chunk, 0)
458 + torch._dynamo.mark_dynamic(chunk_topk_weight, 0)
459 + torch._dynamo.mark_dynamic(chunk_topk_ids, 0)
460 else:
461 _fused_experts = fused_experts
462
@@ -576,7 +559,7 @@ def _pplx_moe(
559
560 with set_current_vllm_config(vllm_config), override_config(moe_config):
561 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
579 - torch_output = torch_moe2(a, w1, w2, topk_weight, topk_ids)
562 + torch_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
563 pplx_output = pplx_moe(group_name, pgi.rank, pgi.world_size, dp_size,
564 a, w1, w2, topk_weight, topk_ids)
565 # TODO (bnell): fix + re-enable
tests/kernels/quantization/test_block_fp8.py
+42 -7
@@ -403,19 +403,24 @@ def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,
403 itertools.product(M_moe_dg, N_moe, K_moe, E, TOP_KS, SEEDS))
404 @pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
405 @torch.inference_mode()
406 -def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
407 -
408 - block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
409 - block_size = [block_m, block_m]
410 - dtype = torch.bfloat16
411 -
406 +def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
407 + monkeypatch):
408 if topk > E:
409 pytest.skip(f"Skipping test: topk={topk} > E={E}")
410
411 if not _valid_deep_gemm_shape(M, N, K):
412 pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}")
413
414 + chunk_size = 1024
415 +
416 torch.manual_seed(seed)
417 +
418 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
419 +
420 + block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
421 + block_size = [block_m, block_m]
422 + dtype = torch.bfloat16
423 +
424 fp8_info = torch.finfo(torch.float8_e4m3fn)
425 fp8_max, fp8_min = fp8_info.max, fp8_info.min
426
@@ -451,6 +456,14 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
456 w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
457 w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
458
459 + # Note: for now use_compile will error out if the problem size is
460 + # large enough to trigger chunking. I'm leaving the flag and
461 + # setup code in case we are able to revisit this later.
462 + use_compile = False
463 +
464 + use_cudagraph = (chunk_size < M and N >= 1024 and K >= 1024
465 + and current_platform.is_cuda_alike())
466 +
467 # Set the context to avoid lots of warning spam.
468 with set_current_vllm_config(vllm_config):
469 if M >= 128:
@@ -463,7 +476,29 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
476 topk_weights, topk_ids, token_expert_indices = fused_topk(
477 a, score.float(), topk, False)
478
466 - out = deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids)
479 + if use_compile:
480 + deep_gemm_moe_fp8_fn = torch.compile(deep_gemm_moe_fp8,
481 + backend="inductor",
482 + fullgraph=True)
483 + torch._dynamo.mark_dynamic(a, 0)
484 + torch._dynamo.mark_dynamic(topk_weights, 0)
485 + torch._dynamo.mark_dynamic(topk_ids, 0)
486 + else:
487 + deep_gemm_moe_fp8_fn = deep_gemm_moe_fp8
488 +
489 + out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
490 + topk_ids)
491 +
492 + if use_cudagraph:
493 + out.fill_(0)
494 + stream = torch.cuda.Stream()
495 + graph = torch.cuda.CUDAGraph()
496 + with torch.cuda.graph(graph, stream=stream):
497 + out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
498 + topk_ids)
499 + torch.cuda.synchronize()
500 + graph.replay()
501 + torch.cuda.synchronize()
502
503 #print(f"{out.sum()=}")
504 #print(f"{ref_out.sum()=}")
tests/kernels/utils.py
+25 -3
@@ -1054,12 +1054,21 @@ def compute_max_diff(output, output_ref):
1054 torch.abs(output_ref))
1055
1056
1057 -def torch_moe(a, w1, w2, score, topk, expert_map):
1057 +def torch_experts(a: torch.Tensor,
1058 + w1: torch.Tensor,
1059 + w2: torch.Tensor,
1060 + topk_weight: torch.Tensor,
1061 + topk_ids: torch.Tensor,
1062 + global_num_experts: int = -1,
1063 + expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
1064 + assert (global_num_experts == -1
1065 + or (global_num_experts == w1.shape[0] and expert_map is None)
1066 + or (expert_map is not None
1067 + and global_num_experts == expert_map.shape[0]))
1068 + topk = topk_ids.shape[1]
1069 B, D = a.shape
1070 a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
1071 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
1061 - score = torch.softmax(score, dim=-1, dtype=torch.float32)
1062 - topk_weight, topk_ids = torch.topk(score, topk)
1072 topk_weight = topk_weight.view(-1)
1073 topk_ids = topk_ids.view(-1)
1074 if expert_map is not None:
@@ -1073,6 +1082,19 @@ def torch_moe(a, w1, w2, score, topk, expert_map):
1082 topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
1083
1084
1085 +def torch_moe(a: torch.Tensor,
1086 + w1: torch.Tensor,
1087 + w2: torch.Tensor,
1088 + score: torch.Tensor,
1089 + topk: int,
1090 + global_num_experts: int = -1,
1091 + expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
1092 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
1093 + topk_weight, topk_ids = torch.topk(score, topk)
1094 + return torch_experts(a, w1, w2, topk_weight, topk_ids, global_num_experts,
1095 + expert_map)
1096 +
1097 +
1098 def torch_moe_single(a, w, score, topk):
1099 B, D = a.shape
1100 a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
vllm/envs.py
+1
@@ -981,6 +981,7 @@ def compute_hash() -> str:
981 "VLLM_DP_RANK",
982 "VLLM_DP_SIZE",
983 "VLLM_USE_STANDALONE_COMPILE",
984 + "VLLM_FUSED_MOE_CHUNK_SIZE",
985 ]
986 for key in environment_variables_to_hash:
987 if key in environment_variables:
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+35 -35
@@ -41,24 +41,24 @@ def run_cutlass_moe_fp8(
41 assert w1.dtype == torch.float8_e4m3fn
42 assert w2.dtype == torch.float8_e4m3fn
43 if expert_num_tokens is None:
44 - assert a1q.shape[1] == w1.shape[2], "Hidden size mismatch w1"
44 + assert a1q.size(1) == w1.size(2), "Hidden size mismatch w1"
45 else:
46 - assert a1q.shape[2] == w1.shape[2], "Hidden size mismatch w1"
47 - assert w1.shape[1] == w2.shape[2] * 2, "Hidden size mismatch w2"
48 - assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
49 - 1] == w1.shape[1], "W1 scale shape mismatch"
50 - assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
51 - 1] == w2.shape[1], "W2 scale shape mismatch"
52 - assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
53 - assert a1q_scale is None or a1q_scale.dim(
54 - ) == 0 or a1q_scale.shape[0] == 1 or a1q_scale.shape[0] == a1q.shape[
55 - 0], "Input scale shape mismatch"
56 - assert w1.shape[0] == w2.shape[0], "Weights expert number mismatch"
57 - assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
58 - assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
59 - assert a2_scale is None or a2_scale.dim(
60 - ) == 0 or a2_scale.shape[0] == 1 or a2_scale.shape[0] == a1q.shape[
61 - 0], "Intermediate scale shape mismatch"
46 + assert a1q.size(2) == w1.size(2), "Hidden size mismatch w1"
47 + assert w1.size(1) == w2.size(2) * 2, "Hidden size mismatch w2"
48 + assert w1_scale.dim() == 1 or w1_scale.size(
49 + 1) == 1 or w1_scale.shape[1] == w1.size(1), "W1 scale shape mismatch"
50 + assert w2_scale.dim() == 1 or w2_scale.size(
51 + 1) == 1 or w2_scale.shape[1] == w2.size(1), "W2 scale shape mismatch"
52 + assert w1.size(0) == w2.size(0), "Expert number mismatch"
53 + assert a1q_scale is None or a1q_scale.dim() == 0 or a1q_scale.size(
54 + 0) == 1 or a1q_scale.size(
55 + 0) == a1q.shape[0], "Input scale shape mismatch"
56 + assert w1.size(0) == w2.size(0), "Weights expert number mismatch"
57 + assert w1.size(0) == w1_scale.size(0), "w1 scales expert number mismatch"
58 + assert w1.size(0) == w2_scale.size(0), "w2 scales expert number mismatch"
59 + assert a2_scale is None or a2_scale.dim() == 0 or a2_scale.size(
60 + 0) == 1 or a2_scale.size(
61 + 0) == a1q.shape[0], "Intermediate scale shape mismatch"
62 assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
63 if expert_map is not None:
64 assert expert_num_tokens is None
@@ -75,12 +75,12 @@ def run_cutlass_moe_fp8(
75 # their tokens are already contiguous for each expert as a result of
76 # the dispatch function.
77
78 - M = a1q.shape[0] # non batched expert M
79 - padded_M = a1q.shape[1] # batched expert M
78 + M = a1q.size(0) # non batched expert M
79 + padded_M = a1q.size(1) # batched expert M
80 _, K, N = w2.shape
81 device = a1q.device
82
83 - assert w1.shape[2] == K
83 + assert w1.size(2) == K
84 assert global_num_experts != -1
85 assert a1q_scale is not None
86
@@ -91,8 +91,8 @@ def run_cutlass_moe_fp8(
91 else:
92 local_topk_ids = topk_ids
93
94 - topk = local_topk_ids.shape[1]
95 - local_E = w1.shape[0]
94 + topk = local_topk_ids.size(1)
95 + local_E = w1.size(0)
96
97 if use_batched_format:
98 assert expert_num_tokens is not None
@@ -111,10 +111,10 @@ def run_cutlass_moe_fp8(
111 problem_sizes2, expert_num_tokens,
112 local_E, padded_M, N, K)
113
114 - w1_scale = w1_scale.reshape(w1_scale.shape[0], -1)
115 - w2_scale = w2_scale.reshape(w2_scale.shape[0], -1)
116 - a1q = a1q.reshape(-1, a1q.shape[2])
117 - a1q_scale = a1q_scale.reshape(-1, a1q_scale.shape[2]).contiguous()
114 + w1_scale = w1_scale.reshape(w1_scale.size(0), -1)
115 + w2_scale = w2_scale.reshape(w2_scale.size(0), -1)
116 + a1q = a1q.reshape(-1, a1q.size(2))
117 + a1q_scale = a1q_scale.reshape(-1, a1q_scale.size(2)).contiguous()
118
119 else:
120 expert_offsets = torch.empty((global_num_experts + 1),
@@ -151,19 +151,19 @@ def run_cutlass_moe_fp8(
151 a1q_scale = a1q_scale[a_map] if per_act_token else a1q_scale
152 expert_offsets = expert_offsets[:-1]
153
154 - ab_strides1 = torch.full((w1.shape[0], ),
154 + ab_strides1 = torch.full((w1.size(0), ),
155 K,
156 device=device,
157 dtype=torch.int64)
158 - c_strides1 = torch.full((w1.shape[0], ),
158 + c_strides1 = torch.full((w1.size(0), ),
159 2 * N,
160 device=device,
161 dtype=torch.int64)
162 - ab_strides2 = torch.full((w1.shape[0], ),
162 + ab_strides2 = torch.full((w1.size(0), ),
163 N,
164 device=device,
165 dtype=torch.int64)
166 - c_strides2 = torch.full((w1.shape[0], ),
166 + c_strides2 = torch.full((w1.size(0), ),
167 K,
168 device=device,
169 dtype=torch.int64)
@@ -237,7 +237,7 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
237 workspace2: tuple[int, ...] = ()
238 output: tuple[int, ...] = ()
239 if self.use_batched_format:
240 - padded_M = aq.shape[1]
240 + padded_M = aq.size(1)
241 workspace1 = (self.max_experts_per_worker, padded_M, max(N, K))
242 workspace2 = (self.max_experts_per_worker, padded_M, (N // 2))
243 output = (self.max_experts_per_worker, padded_M, K)
@@ -332,7 +332,7 @@ def cutlass_moe_fp8(
332 """
333 per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
334 a2_scale.numel() != 1 if a2_scale is not None else False)
335 - per_out_ch = w1_scale.numel() != w1_q.shape[0]
335 + per_out_ch = w1_scale.numel() != w1_q.size(0)
336
337 out_dtype = a.dtype
338
@@ -425,11 +425,11 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
425 assert (m == m_a), "input shape mismatch"
426 assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1"
427 assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
428 - assert (topk_weights.shape[0] == m and topk_ids.shape[0]
428 + assert (topk_weights.size(0) == m and topk_ids.size(0)
429 == m), ("topk must be provided for each row of a")
430
431 out_dtype = a.dtype
432 - num_topk = topk_ids.shape[1]
432 + num_topk = topk_ids.size(1)
433
434 expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
435 blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
@@ -463,7 +463,7 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
463 out_dtype, device)
464 del rep_a_fp4, rep_a_blockscale
465 # hidden size dimension is split to one halfpytho sized tensor.
466 - intermediate = torch.empty((m * num_topk, w1_fp4.shape[1] // 2),
466 + intermediate = torch.empty((m * num_topk, w1_fp4.size(1) // 2),
467 device=device,
468 dtype=out_dtype)
469
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+1 -1
@@ -48,7 +48,7 @@ def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
48 M = hidden_states.size(0)
49 _, K, N = w2.size()
50 if not _valid_deep_gemm_shape(M, N, K):
51 - logger.debug("DeepGemm disabled: unalinged problem size.")
51 + logger.debug("DeepGemm disabled: unaligned problem size.")
52 return False
53
54 if (w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn):
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+1 -1
@@ -25,7 +25,7 @@ def dequant_fp8(expert_x_fp8: torch.Tensor,
25 expert_x_fp32 = expert_x_fp8.to(torch.float32).view(
26 num_experts, -1, DEEPEP_QUANT_BLOCK_SIZE)
27 expert_x_scales = expert_x_scales.view(num_experts, -1, 1)
28 - return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.shape)
28 + return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.size())
29
30
31 class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
vllm/model_executor/layers/fused_moe/fused_moe.py
+42 -42
@@ -488,10 +488,10 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
488
489 if use_fp8_w8a8 or use_int8_w8a8:
490 assert B_scale is not None
491 - assert (block_shape is None or triton.cdiv(B.shape[-2], block_shape[0])
492 - == B_scale.shape[-2])
493 - assert (block_shape is None or triton.cdiv(B.shape[-1], block_shape[1])
494 - == B_scale.shape[-1])
491 + assert (block_shape is None
492 + or triton.cdiv(B.size(-2), block_shape[0]) == B_scale.size(-2))
493 + assert (block_shape is None
494 + or triton.cdiv(B.size(-1), block_shape[1]) == B_scale.size(-1))
495
496 elif use_int8_w8a16 or use_int4_w4a16:
497 assert B_scale is not None
@@ -500,19 +500,19 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
500 assert A_scale is None
501 assert B_scale is None
502
503 - M = A.shape[0]
503 + M = A.size(0)
504 num_tokens = M * top_k
505
506 - EM = sorted_token_ids.shape[0]
507 - if A.shape[0] < config["BLOCK_SIZE_M"]:
506 + EM = sorted_token_ids.size(0)
507 + if A.size(0) < config["BLOCK_SIZE_M"]:
508 # optimize for small batch_size.
509 # We assume that top_ids of each token is unique, so
510 # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
511 # and we can skip some invalid blocks.
512 - EM = min(sorted_token_ids.shape[0],
513 - A.shape[0] * top_k * config['BLOCK_SIZE_M'])
512 + EM = min(sorted_token_ids.size(0),
513 + A.size(0) * top_k * config['BLOCK_SIZE_M'])
514 grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(
515 - B.shape[1], META['BLOCK_SIZE_N']), )
515 + B.size(1), META['BLOCK_SIZE_N']), )
516
517 if (use_int8_w8a16 or use_int4_w4a16) and \
518 block_shape is not None and block_shape[1] > 0:
@@ -522,16 +522,16 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
522 use_moe_wna16_cuda = should_moe_wna16_use_cuda(
523 num_valid_tokens=num_tokens,
524 group_size=block_shape[1],
525 - num_experts=B.shape[0],
525 + num_experts=B.size(0),
526 bit=4 if use_int4_w4a16 else 8)
527 config = config.copy()
528 config.update(
529 get_moe_wna16_block_config(config=config,
530 use_moe_wna16_cuda=use_moe_wna16_cuda,
531 num_valid_tokens=num_tokens,
532 - size_k=A.shape[1],
533 - size_n=B.shape[1],
534 - num_experts=B.shape[1],
532 + size_k=A.size(1),
533 + size_n=B.size(1),
534 + num_experts=B.size(1),
535 group_size=block_shape[1],
536 real_top_k=top_k,
537 block_size_m=config["BLOCK_SIZE_M"]))
@@ -556,8 +556,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
556 sorted_token_ids,
557 expert_ids,
558 num_tokens_post_padded,
559 - B.shape[1],
560 - A.shape[1],
559 + B.size(1),
560 + A.size(1),
561 EM,
562 num_tokens,
563 A.stride(0),
@@ -573,7 +573,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
573 B_zp.stride(0) if B_zp is not None else 0,
574 B_zp.stride(2) if B_zp is not None else 0,
575 B_zp.stride(1) if B_zp is not None else 0,
576 - block_k_diviable=A.shape[1] % config["BLOCK_SIZE_K"] == 0,
576 + block_k_diviable=A.size(1) % config["BLOCK_SIZE_K"] == 0,
577 group_size=block_shape[1],
578 MUL_ROUTED_WEIGHT=mul_routed_weight,
579 top_k=top_k,
@@ -599,8 +599,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
599 sorted_token_ids,
600 expert_ids,
601 num_tokens_post_padded,
602 - B.shape[1],
603 - B.shape[2],
602 + B.size(1),
603 + B.size(2),
604 EM,
605 num_tokens,
606 A.stride(0),
@@ -818,7 +818,7 @@ def try_get_optimal_moe_config(
818 M: int,
819 is_marlin: bool = False,
820 block_shape: Optional[list[int]] = None,
821 -):
821 +) -> dict[str, int]:
822 from vllm.model_executor.layers.fused_moe import get_config
823 override_config = get_config()
824 if override_config:
@@ -873,10 +873,10 @@ def fused_topk(
873 renormalize: bool,
874 indices_type: Optional[torch.dtype] = None,
875 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
876 - assert hidden_states.shape[0] == gating_output.shape[0], (
876 + assert hidden_states.size(0) == gating_output.size(0), (
877 "Number of tokens mismatch")
878
879 - M, _ = hidden_states.shape
879 + M, _ = hidden_states.size()
880
881 topk_weights = torch.empty(M,
882 topk,
@@ -915,7 +915,7 @@ def grouped_topk(
915 e_score_correction_bias: Optional[torch.Tensor] = None
916 ) -> tuple[torch.Tensor, torch.Tensor]:
917
918 - assert hidden_states.shape[0] == gating_output.shape[0], (
918 + assert hidden_states.size(0) == gating_output.size(0), (
919 "Number of tokens mismatch")
920
921 if scoring_func == "softmax":
@@ -925,7 +925,7 @@ def grouped_topk(
925 else:
926 raise ValueError(f"Unsupported scoring function: {scoring_func}")
927
928 - num_token = scores.shape[0]
928 + num_token = scores.size(0)
929 if e_score_correction_bias is not None:
930 # Store original scores before applying correction bias. We use biased
931 # scores for expert selection but original scores for routing weights
@@ -942,7 +942,7 @@ def grouped_topk(
942 group_mask.scatter_(1, group_idx, 1) # [n, n_group]
943 score_mask = group_mask.unsqueeze(-1).expand(
944 num_token, num_expert_group,
945 - scores.shape[-1] // num_expert_group).reshape(num_token, -1) # [n, e]
945 + scores.size(-1) // num_expert_group).reshape(num_token, -1) # [n, e]
946 tmp_scores = scores.masked_fill(~score_mask.bool(),
947 float("-inf")) # [n, e]
948
@@ -1162,7 +1162,7 @@ def fused_experts(hidden_states: torch.Tensor,
1162 allow_deep_gemm: bool = False) -> torch.Tensor:
1163 # For now, disable DeepGemm for small N (<= 512) until better
1164 # permute/unpermute ops are available.
1165 - N = w1.shape[1]
1165 + N = w1.size(1)
1166 if (allow_deep_gemm and use_fp8_w8a8 and N > 512
1167 and _valid_deep_gemm(hidden_states, w1, w2)):
1168 assert apply_router_weight_on_input is False
@@ -1233,13 +1233,13 @@ def fused_experts_impl(
1233 ) -> torch.Tensor:
1234 # Check constraints.
1235 if use_int4_w4a16:
1236 - assert hidden_states.shape[1] // 2 == w1.shape[
1237 - 2], "Hidden size mismatch"
1236 + assert hidden_states.size(1) // 2 == w1.size(2), (
1237 + "Hidden size mismatch")
1238 else:
1239 - assert hidden_states.shape[1] == w1.shape[2], (
1240 - f"Hidden size mismatch {hidden_states.shape[1]} != {w1.shape[2]}")
1239 + assert hidden_states.size(1) == w1.size(2), (
1240 + f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}")
1241
1242 - assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1242 + assert topk_weights.size() == topk_ids.size(), "topk shape mismatch"
1243 assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
1244 assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
1245 assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
@@ -1247,12 +1247,12 @@ def fused_experts_impl(
1247 torch.float32, torch.float16, torch.bfloat16
1248 ]
1249
1250 - num_tokens = hidden_states.shape[0]
1251 - E, N, _ = w1.shape
1252 - K = w2.shape[1]
1250 + num_tokens = hidden_states.size(0)
1251 + E, N, _ = w1.size()
1252 + K = w2.size(1)
1253 if global_num_experts == -1:
1254 global_num_experts = E
1255 - top_k_num = topk_ids.shape[1]
1255 + top_k_num = topk_ids.size(1)
1256 # We execute the fused_moe kernel in chunks to circumvent this issue:
1257 # https://github.com/vllm-project/vllm/issues/5938
1258 CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
@@ -1269,8 +1269,8 @@ def fused_experts_impl(
1269
1270 get_config_func = functools.partial(
1271 try_get_optimal_moe_config,
1272 - w1.shape,
1273 - w2.shape,
1272 + w1.size(),
1273 + w2.size(),
1274 top_k_num,
1275 config_dtype,
1276 block_shape=block_shape,
@@ -1310,7 +1310,7 @@ def fused_experts_impl(
1310 min((chunk + 1) * CHUNK_SIZE,
1311 num_tokens))
1312 curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
1313 - tokens_in_chunk, _ = curr_hidden_states.shape
1313 + tokens_in_chunk, _ = curr_hidden_states.size()
1314
1315 if tokens_in_chunk == 0:
1316 break
@@ -1322,7 +1322,7 @@ def fused_experts_impl(
1322 # do not need to be adjusted.
1323 intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
1324 intermediate_cache2 = intermediate_cache2[:tokens_in_chunk *
1325 - topk_ids.shape[1]]
1325 + topk_ids.size(1)]
1326 intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
1327 config = get_config_func(tokens_in_chunk)
1328
@@ -1398,7 +1398,7 @@ def fused_experts_impl(
1398 per_channel_quant=per_channel_quant,
1399 block_shape=block_shape)
1400
1401 - ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
1401 + ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.size()),
1402 out_hidden_states[begin_chunk_idx:end_chunk_idx])
1403
1404 return out_hidden_states
@@ -1611,8 +1611,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1611 dtype=hidden_states.dtype)
1612
1613 config = try_get_optimal_moe_config(
1614 - w1.shape,
1615 - w2.shape,
1614 + w1.size(),
1615 + w2.size(),
1616 top_k_num,
1617 config_dtype,
1618 num_tokens,
vllm/model_executor/layers/fused_moe/layer.py
+7 -12
@@ -861,13 +861,11 @@ class FusedMoE(torch.nn.Module):
861 self.global_num_experts = num_experts
862
863 # For smuggling this layer into the fused moe custom op
864 - self.use_direct_call = self.dp_size == 1
865 - if not self.use_direct_call:
866 - compilation_config = vllm_config.compilation_config
867 - if prefix in compilation_config.static_forward_context:
868 - raise ValueError("Duplicate layer name: {}".format(prefix))
869 - compilation_config.static_forward_context[prefix] = self
870 - self.layer_name = prefix
864 + compilation_config = vllm_config.compilation_config
865 + if prefix in compilation_config.static_forward_context:
866 + raise ValueError("Duplicate layer name: {}".format(prefix))
867 + compilation_config.static_forward_context[prefix] = self
868 + self.layer_name = prefix
869
870 # Determine expert maps
871 if self.use_ep:
@@ -1361,11 +1359,8 @@ class FusedMoE(torch.nn.Module):
1359
1360 def forward(self, hidden_states: torch.Tensor,
1361 router_logits: torch.Tensor):
1364 - if self.use_direct_call:
1365 - return self.forward_impl(hidden_states, router_logits)
1366 - else:
1367 - return torch.ops.vllm.moe_forward(hidden_states, router_logits,
1368 - self.layer_name)
1362 + return torch.ops.vllm.moe_forward(hidden_states, router_logits,
1363 + self.layer_name)
1364
1365 def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
1366 full_router_logits: torch.Tensor):
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+1 -1
@@ -69,7 +69,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
69 a1 = a1 * rank_topk_weights.to(a1.dtype)
70
71 repeat_cols = 4
72 - repeat_rows = 1 if self.per_act_token else a1.shape[0]
72 + repeat_rows = 1 if self.per_act_token else a1.size(0)
73 a1q, a1q_scale = moe_kernel_quantize_input(
74 a1, (None if self.per_act_token else a1_scale), self.quant_dtype,
75 self.per_act_token, self.block_shape)