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, 17:15 UTC
[ROCm][Kernel] MoE weights padding (#14454)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: charlifu <charlifu@amd.com>
Gregory Shtrasberg committed
Mar 24, 2025 at 19:45 UTC
f533b5837fa67f53957a12387a01067f9edef0d8
5 files changed
+65
-16
tests/kernels/test_moe.py
+35
-10
@@ -3,8 +3,11 @@
3
4
Run `pytest tests/kernels/test_moe.py`.
5
"""
6
+
7
import pytest
8
import torch
9
+from torch.nn import Parameter
10
+from torch.nn import functional as F
11
from transformers import MixtralConfig
12
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
13
@@ -37,6 +40,7 @@ TOP_KS = [2, 6]
40
@pytest.mark.parametrize("topk", TOP_KS)
41
@pytest.mark.parametrize("ep_size", EP_SIZE)
42
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
43
+@pytest.mark.parametrize("padding", [True, False])
44
def test_fused_moe(
45
m: int,
46
n: int,
@@ -45,6 +49,7 @@ def test_fused_moe(
49
topk: int,
50
ep_size: int,
51
dtype: torch.dtype,
52
+ padding: bool,
53
):
54
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
55
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
@@ -65,16 +70,7 @@ def test_fused_moe(
70
else:
71
e_map = None
72
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)
73
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)
74
iterative_output = iterative_moe(a,
75
w1,
76
w2,
@@ -83,6 +79,23 @@ def test_fused_moe(
79
global_num_experts=e,
80
expert_map=e_map,
81
renormalize=False)
82
+
83
+ # Pad the weight if moe padding is enabled
84
+ if padding:
85
+ w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
86
+ torch.cuda.empty_cache()
87
+ w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
88
+ torch.cuda.empty_cache()
89
+
90
+ triton_output = fused_moe(a,
91
+ w1,
92
+ w2,
93
+ score,
94
+ topk,
95
+ global_num_experts=e,
96
+ expert_map=e_map,
97
+ renormalize=False)
98
+ torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
99
torch.testing.assert_close(iterative_output,
100
torch_output,
101
atol=2e-2,
@@ -202,8 +215,9 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
215
216
@pytest.mark.parametrize("dtype",
217
[torch.float32, torch.float16, torch.bfloat16])
218
+@pytest.mark.parametrize("padding", [True, False])
219
@torch.inference_mode()
206
-def test_mixtral_moe(dtype: torch.dtype):
220
+def test_mixtral_moe(dtype: torch.dtype, padding: bool):
221
"""Make sure our Mixtral MoE implementation agrees with the one from
222
huggingface."""
223
@@ -233,6 +247,17 @@ def test_mixtral_moe(dtype: torch.dtype):
247
# vLLM uses 1D query [num_tokens, hidden_dim]
248
vllm_inputs = hf_inputs.flatten(0, 1)
249
250
+ # Pad the weight if moe padding is enabled
251
+ if padding:
252
+ vllm_moe.experts.w13_weight = Parameter(F.pad(
253
+ vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[..., 0:-128],
254
+ requires_grad=False)
255
+ torch.cuda.empty_cache()
256
+ vllm_moe.experts.w2_weight = Parameter(F.pad(
257
+ vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
258
+ requires_grad=False)
259
+ torch.cuda.empty_cache()
260
+
261
# Run forward passes for both MoE blocks
262
hf_states, _ = hf_moe.forward(hf_inputs)
263
vllm_states = vllm_moe.forward(vllm_inputs)
vllm/envs.py
+5
@@ -75,6 +75,7 @@ if TYPE_CHECKING:
75
VLLM_ROCM_USE_AITER: bool = False
76
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
77
VLLM_ROCM_FP8_PADDING: bool = True
78
+ VLLM_ROCM_MOE_PADDING: bool = True
79
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
80
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
81
VLLM_DISABLE_COMPILE_CACHE: bool = False
@@ -520,6 +521,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
521
"VLLM_ROCM_FP8_PADDING":
522
lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))),
523
524
+ # Pad the weights for the moe kernel
525
+ "VLLM_ROCM_MOE_PADDING":
526
+ lambda: bool(int(os.getenv("VLLM_ROCM_MOE_PADDING", "1"))),
527
+
528
# Divisor for dynamic query scale factor calculation for FP8 KV Cache
529
"Q_SCALE_CONSTANT":
530
lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),
vllm/model_executor/layers/fused_moe/fused_moe.py
+3
-3
@@ -800,7 +800,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
800
expert_ids,
801
num_tokens_post_padded,
802
B.shape[1],
803
- A.shape[1],
803
+ B.shape[2],
804
EM,
805
topk_ids.numel(),
806
A.stride(0),
@@ -1322,8 +1322,8 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1322
1323
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1324
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
1325
- assert w1.is_contiguous(), "Expert weights1 must be contiguous"
1326
- assert w2.is_contiguous(), "Expert weights2 must be contiguous"
1325
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
1326
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
1327
assert hidden_states.dtype in [
1328
torch.float32, torch.float16, torch.bfloat16
1329
]
vllm/model_executor/layers/fused_moe/layer.py
+19
@@ -5,6 +5,7 @@ from enum import Enum
5
from typing import Callable, List, Optional, Tuple
6
7
import torch
8
+import torch.nn.functional as F
9
from torch.nn.parameter import UninitializedParameter
10
11
from vllm import envs
@@ -96,9 +97,27 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
97
layer.register_parameter("w2_weight", w2_weight)
98
set_weight_attrs(w2_weight, extra_weight_attrs)
99
100
+ def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
101
+ # Pad the weight tensor. This is an optimization on ROCm platform, which
102
+ # can benefit from tensors located far enough from one another in memory
103
+ if (envs.VLLM_ROCM_MOE_PADDING and current_platform.is_rocm()
104
+ and weight.stride(-1) == 1
105
+ and (weight.stride(-2) * weight.element_size()) % 512 == 0):
106
+ num_pad = 256 // weight.element_size()
107
+ weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
108
+ torch.cuda.empty_cache()
109
+ return weight
110
+
111
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
112
super().process_weights_after_loading(layer)
113
114
+ layer.w13_weight = torch.nn.Parameter(self._maybe_pad_weight(
115
+ layer.w13_weight.data),
116
+ requires_grad=False)
117
+ layer.w2_weight = torch.nn.Parameter(self._maybe_pad_weight(
118
+ layer.w2_weight.data),
119
+ requires_grad=False)
120
+
121
if current_platform.is_cpu():
122
if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
123
import intel_extension_for_pytorch as ipex
vllm/model_executor/layers/quantization/fp8.py
+3
-3
@@ -255,7 +255,7 @@ class Fp8LinearMethod(LinearMethodBase):
255
else:
256
layer.register_parameter("input_scale", None)
257
258
- def add_padding_to_weight(self, weight: torch.Tensor) -> torch.Tensor:
258
+ def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
259
# Pad the weight tensor. This is an optimization on ROCm platform, which
260
# can benefit from tensors located far enough from one another in memory
261
if (envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm()
@@ -279,7 +279,7 @@ class Fp8LinearMethod(LinearMethodBase):
279
weight = layer.weight.data
280
weight_scale_inv = layer.weight_scale_inv.data
281
282
- weight = self.add_padding_to_weight(weight)
282
+ weight = self._maybe_pad_weight(weight)
283
284
# Torch.compile cannot use Parameter subclasses.
285
layer.weight = Parameter(weight, requires_grad=False)
@@ -343,7 +343,7 @@ class Fp8LinearMethod(LinearMethodBase):
343
logical_widths=layer.logical_widths,
344
)
345
346
- weight = self.add_padding_to_weight(weight)
346
+ weight = self._maybe_pad_weight(weight)
347
# Update layer with new values.
348
layer.weight = Parameter(weight.t(), requires_grad=False)
349
layer.weight_scale = Parameter(weight_scale, requires_grad=False)