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
[FEAT][ROCm] Integrate Fused MoE Kernels from AITER (#14967)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
vllmellm committed
Mar 26, 2025 at 16:30 UTC
5ebf66748b8b67731972c389d879ca69c68dc2c4
9 files changed
+390
-65
tests/kernels/test_moe.py
+19
-6
@@ -3,7 +3,6 @@
3
4
Run `pytest tests/kernels/test_moe.py`.
5
"""
6
-
6
import pytest
7
import torch
8
from torch.nn import Parameter
@@ -216,11 +215,17 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
215
@pytest.mark.parametrize("dtype",
216
[torch.float32, torch.float16, torch.bfloat16])
217
@pytest.mark.parametrize("padding", [True, False])
218
+@pytest.mark.parametrize(
219
+ "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
220
@torch.inference_mode()
220
-def test_mixtral_moe(dtype: torch.dtype, padding: bool):
221
+def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
222
+ monkeypatch):
223
"""Make sure our Mixtral MoE implementation agrees with the one from
224
huggingface."""
225
226
+ if use_rocm_aiter:
227
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
228
+
229
# Instantiate our and huggingface's MoE blocks
230
config = MixtralConfig()
231
hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
@@ -268,10 +273,18 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool):
273
torch.bfloat16: 1e-2,
274
}
275
271
- torch.testing.assert_close(hf_states.flatten(0, 1),
272
- vllm_states,
273
- rtol=mixtral_moe_tol[dtype],
274
- atol=mixtral_moe_tol[dtype])
276
+ if use_rocm_aiter:
277
+ # The values of rtol and atol are set based on the tests in ROCM AITER package. # noqa: E501
278
+ # https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174 # noqa: E501
279
+ torch.testing.assert_close(hf_states.flatten(0, 1),
280
+ vllm_states,
281
+ rtol=0.01,
282
+ atol=100)
283
+ else:
284
+ torch.testing.assert_close(hf_states.flatten(0, 1),
285
+ vllm_states,
286
+ rtol=mixtral_moe_tol[dtype],
287
+ atol=mixtral_moe_tol[dtype])
288
289
290
@pytest.mark.parametrize("m", [1, 33, 64, 222])
tests/model_executor/test_enabled_custom_ops.py
+36
@@ -7,6 +7,10 @@ from vllm.model_executor.custom_op import CustomOp
7
from vllm.model_executor.layers.activation import (GeluAndMul,
8
ReLUSquaredActivation,
9
SiluAndMul)
10
+from vllm.model_executor.layers.fused_moe.fused_moe import (
11
+ dispatch_fused_experts_func, dispatch_topk_func,
12
+ torch_vllm_inplace_fused_experts, torch_vllm_outplace_fused_experts,
13
+ vllm_topk_softmax)
14
from vllm.model_executor.layers.layernorm import (
15
RMSNorm, dispatch_cuda_rmsnorm_func, fused_add_rms_norm, rms_norm,
16
rocm_aiter_fused_add_rms_norm, rocm_aiter_rms_norm)
@@ -92,6 +96,38 @@ def test_enabled_ops_invalid(env: str):
96
RMSNorm(1024).enabled()
97
98
99
+@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
100
+def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
101
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
102
+ topk_func = dispatch_topk_func()
103
+
104
+ if current_platform.is_rocm() and int(use_rocm_aiter):
105
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
106
+ rocm_aiter_topk_softmax)
107
+
108
+ assert topk_func == rocm_aiter_topk_softmax
109
+ else:
110
+ assert topk_func == vllm_topk_softmax
111
+
112
+
113
+@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
114
+@pytest.mark.parametrize("inplace", [True, False])
115
+def test_fused_experts_dispatch(use_rocm_aiter: str, inplace: bool,
116
+ monkeypatch):
117
+
118
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
119
+ fused_experts_func = dispatch_fused_experts_func(inplace)
120
+ if current_platform.is_rocm() and int(use_rocm_aiter):
121
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
122
+ rocm_aiter_fused_experts)
123
+
124
+ assert fused_experts_func == rocm_aiter_fused_experts
125
+ elif inplace:
126
+ assert fused_experts_func == torch_vllm_inplace_fused_experts
127
+ else:
128
+ assert fused_experts_func == torch_vllm_outplace_fused_experts
129
+
130
+
131
@pytest.mark.parametrize("add_residual", [True, False])
132
@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
133
@pytest.mark.parametrize("use_rocm_aiter_norm", ["0", "1"])
tests/models/decoder_only/language/test_mistral.py
+9
-32
@@ -174,15 +174,8 @@ SAMPLE_JSON_SCHEMA = {
174
@pytest.mark.parametrize("dtype", ["bfloat16"])
175
@pytest.mark.parametrize("max_tokens", [64])
176
@pytest.mark.parametrize("num_logprobs", [5])
177
-def test_models(
178
- hf_runner,
179
- vllm_runner,
180
- example_prompts,
181
- model: str,
182
- dtype: str,
183
- max_tokens: int,
184
- num_logprobs: int,
185
-) -> None:
177
+def test_models(hf_runner, vllm_runner, example_prompts, model: str,
178
+ dtype: str, max_tokens: int, num_logprobs: int) -> None:
179
# TODO(sang): Sliding window should be tested separately.
180
with hf_runner(model, dtype=dtype) as hf_model:
181
hf_outputs = hf_model.generate_greedy_logprobs_limit(
@@ -206,14 +199,8 @@ def test_models(
199
@pytest.mark.parametrize("dtype", ["bfloat16"])
200
@pytest.mark.parametrize("max_tokens", [64])
201
@pytest.mark.parametrize("num_logprobs", [5])
209
-def test_mistral_format(
210
- vllm_runner,
211
- example_prompts,
212
- model: str,
213
- dtype: str,
214
- max_tokens: int,
215
- num_logprobs: int,
216
-) -> None:
202
+def test_mistral_format(vllm_runner, example_prompts, model: str, dtype: str,
203
+ max_tokens: int, num_logprobs: int) -> None:
204
with vllm_runner(
205
model,
206
dtype=dtype,
@@ -244,11 +231,8 @@ def test_mistral_format(
231
232
@pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS)
233
@pytest.mark.parametrize("dtype", ["bfloat16"])
247
-def test_mistral_symbolic_languages(
248
- vllm_runner,
249
- model: str,
250
- dtype: str,
251
-) -> None:
234
+def test_mistral_symbolic_languages(vllm_runner, model: str,
235
+ dtype: str) -> None:
236
with vllm_runner(model,
237
dtype=dtype,
238
max_model_len=8192,
@@ -266,11 +250,7 @@ def test_mistral_symbolic_languages(
250
@pytest.mark.parametrize("dtype", ["bfloat16"])
251
@pytest.mark.parametrize("model",
252
MISTRAL_FORMAT_MODELS) # v1 can't do func calling
269
-def test_mistral_function_calling(
270
- vllm_runner,
271
- model: str,
272
- dtype: str,
273
-) -> None:
253
+def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None:
254
with vllm_runner(model,
255
dtype=dtype,
256
tokenizer_mode="mistral",
@@ -301,11 +281,8 @@ def test_mistral_function_calling(
281
@pytest.mark.parametrize("model", MODELS)
282
@pytest.mark.parametrize("guided_backend",
283
["outlines", "lm-format-enforcer", "xgrammar"])
304
-def test_mistral_guided_decoding(
305
- vllm_runner,
306
- model: str,
307
- guided_backend: str,
308
-) -> None:
284
+def test_mistral_guided_decoding(vllm_runner, model: str,
285
+ guided_backend: str) -> None:
286
with vllm_runner(model, dtype='bfloat16',
287
tokenizer_mode="mistral") as vllm_model:
288
tests/quantization/test_fp8.py
+20
-3
@@ -23,8 +23,14 @@ MODELS = [
23
reason="FP8 is not supported on this GPU type.")
24
@pytest.mark.parametrize("model_id", MODELS)
25
@pytest.mark.parametrize("force_marlin", [False, True])
26
+@pytest.mark.parametrize(
27
+ "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
28
def test_model_load_and_run(vllm_runner, model_id: str, force_marlin: bool,
27
- monkeypatch) -> None:
29
+ use_rocm_aiter: bool, monkeypatch) -> None:
30
+
31
+ if use_rocm_aiter:
32
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
33
+
34
if force_marlin:
35
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
36
@@ -47,7 +53,13 @@ KV_CACHE_MODELS = [
53
@pytest.mark.skipif(not is_quant_method_supported("fp8"),
54
reason="FP8 is not supported on this GPU type.")
55
@pytest.mark.parametrize("model_id", KV_CACHE_MODELS)
50
-def test_kv_cache_model_load_and_run(vllm_runner, model_id: str, monkeypatch):
56
+@pytest.mark.parametrize(
57
+ "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
58
+def test_kv_cache_model_load_and_run(vllm_runner, model_id: str,
59
+ use_rocm_aiter: bool, monkeypatch):
60
+ if use_rocm_aiter:
61
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
62
+
63
# vllm_runner.apply_model() relies on V0 internals.
64
monkeypatch.setenv("VLLM_USE_V1", "0")
65
with vllm_runner(model_id, kv_cache_dtype="fp8") as llm:
@@ -86,8 +98,13 @@ def test_kv_cache_model_load_and_run(vllm_runner, model_id: str, monkeypatch):
98
reason="FP8 is not supported on this GPU type.")
99
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
100
@pytest.mark.parametrize("force_marlin", [False, True])
101
+@pytest.mark.parametrize(
102
+ "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
103
def test_load_fp16_model(vllm_runner, kv_cache_dtype: str, force_marlin: bool,
90
- monkeypatch) -> None:
104
+ use_rocm_aiter: bool, monkeypatch) -> None:
105
+ if use_rocm_aiter:
106
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
107
+
108
# vllm_runner.apply_model() relies on V0 internals.
109
monkeypatch.setenv("VLLM_USE_V1", "0")
110
vllm/envs.py
+15
@@ -73,6 +73,8 @@ if TYPE_CHECKING:
73
VLLM_DISABLED_KERNELS: list[str] = []
74
VLLM_USE_V1: bool = True
75
VLLM_ROCM_USE_AITER: bool = False
76
+ VLLM_ROCM_USE_AITER_MOE: bool = True
77
+ VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE: bool = False
78
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
79
VLLM_ROCM_FP8_PADDING: bool = True
80
VLLM_ROCM_MOE_PADDING: bool = True
@@ -513,6 +515,19 @@ environment_variables: dict[str, Callable[[], Any]] = {
515
lambda: (os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in
516
("true", "1")),
517
518
+ # Whether to use aiter moe ops.
519
+ # By default is enabled.
520
+ "VLLM_ROCM_USE_AITER_MOE":
521
+ lambda: (os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in
522
+ ("true", "1")),
523
+
524
+ # Whether to use aiter block scaled moe kernel.
525
+ # By default this is disabled.
526
+ "VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE":
527
+ lambda:
528
+ (os.getenv("VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE", "false").lower() in
529
+ ("true", "1")),
530
+
531
# use aiter rms norm op if aiter ops are enabled.
532
"VLLM_ROCM_USE_AITER_RMSNORM":
533
lambda: (os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in
vllm/model_executor/layers/fused_moe/fused_moe.py
+69
-23
@@ -17,6 +17,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
17
from vllm.platforms import current_platform
18
from vllm.utils import direct_register_custom_op
19
20
+from .rocm_aiter_fused_moe import (is_rocm_aiter_moe_enabled,
21
+ rocm_aiter_fused_experts,
22
+ rocm_aiter_topk_softmax)
23
+
24
logger = init_logger(__name__)
25
26
@@ -1035,6 +1039,28 @@ def try_get_optimal_moe_config(
1039
return config
1040
1041
1042
+def vllm_topk_softmax(topk_weights: torch.Tensor, topk_indices: torch.Tensor,
1043
+ token_expert_indices: torch.Tensor,
1044
+ gating_output: torch.Tensor,
1045
+ renormalize: bool) -> tuple[torch.Tensor, ...]:
1046
+ ops.topk_softmax(
1047
+ topk_weights,
1048
+ topk_indices,
1049
+ token_expert_indices,
1050
+ gating_output,
1051
+ )
1052
+ if renormalize:
1053
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1054
+
1055
+ return topk_weights, topk_indices
1056
+
1057
+
1058
+def dispatch_topk_func() -> Callable[..., tuple[torch.Tensor, ...]]:
1059
+ if is_rocm_aiter_moe_enabled():
1060
+ return rocm_aiter_topk_softmax
1061
+ return vllm_topk_softmax
1062
+
1063
+
1064
def fused_topk(
1065
hidden_states: torch.Tensor,
1066
gating_output: torch.Tensor,
@@ -1059,17 +1085,14 @@ def fused_topk(
1085
dtype=torch.int32,
1086
device=hidden_states.device)
1087
1062
- ops.topk_softmax(
1063
- topk_weights,
1064
- topk_ids,
1065
- token_expert_indicies,
1066
- gating_output.float(), # TODO(woosuk): Optimize this.
1067
- )
1068
- del token_expert_indicies # Not used. Will be used in the future.
1088
+ gating_output_float = gating_output.float() # TODO(woosuk): Optimize this.
1089
1070
- if renormalize:
1071
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1090
+ topk_func = dispatch_topk_func()
1091
+ topk_weights, topk_ids = topk_func(topk_weights, topk_ids,
1092
+ token_expert_indicies,
1093
+ gating_output_float, renormalize)
1094
1095
+ del token_expert_indicies # Not used. Will be used in the future.
1096
return topk_weights, topk_ids
1097
1098
@@ -1259,6 +1282,24 @@ direct_register_custom_op(
1282
)
1283
1284
1285
+def torch_vllm_inplace_fused_experts(**kwargs) -> torch.Tensor:
1286
+ torch.ops.vllm.inplace_fused_experts(**kwargs)
1287
+ hidden_states = kwargs['hidden_states']
1288
+ return hidden_states
1289
+
1290
+
1291
+def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:
1292
+ return torch.ops.vllm.outplace_fused_experts(**kwargs)
1293
+
1294
+
1295
+def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
1296
+ if is_rocm_aiter_moe_enabled():
1297
+ return rocm_aiter_fused_experts
1298
+ if inplace:
1299
+ return torch_vllm_inplace_fused_experts
1300
+ return torch_vllm_outplace_fused_experts
1301
+
1302
+
1303
def fused_experts(hidden_states: torch.Tensor,
1304
w1: torch.Tensor,
1305
w2: torch.Tensor,
@@ -1278,20 +1319,25 @@ def fused_experts(hidden_states: torch.Tensor,
1319
a1_scale: Optional[torch.Tensor] = None,
1320
a2_scale: Optional[torch.Tensor] = None,
1321
block_shape: Optional[List[int]] = None) -> torch.Tensor:
1281
-
1282
- if inplace:
1283
- torch.ops.vllm.inplace_fused_experts(
1284
- hidden_states, w1, w2, topk_weights, topk_ids, activation,
1285
- use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,
1286
- expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1287
- block_shape)
1288
- return hidden_states
1289
- else:
1290
- return torch.ops.vllm.outplace_fused_experts(
1291
- hidden_states, w1, w2, topk_weights, topk_ids, activation,
1292
- use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,
1293
- expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1294
- block_shape)
1322
+ return dispatch_fused_experts_func(inplace)(
1323
+ hidden_states=hidden_states,
1324
+ w1=w1,
1325
+ w2=w2,
1326
+ topk_weights=topk_weights,
1327
+ topk_ids=topk_ids,
1328
+ activation=activation,
1329
+ use_fp8_w8a8=use_fp8_w8a8,
1330
+ use_int8_w8a16=use_int8_w8a16,
1331
+ use_int4_w4a16=use_int4_w4a16,
1332
+ global_num_experts=global_num_experts,
1333
+ expert_map=expert_map,
1334
+ w1_scale=w1_scale,
1335
+ w2_scale=w2_scale,
1336
+ w1_zp=w1_zp,
1337
+ w2_zp=w2_zp,
1338
+ a1_scale=a1_scale,
1339
+ a2_scale=a2_scale,
1340
+ block_shape=block_shape)
1341
1342
1343
def fused_experts_impl(hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/layer.py
+13
-1
@@ -8,7 +8,7 @@ import torch
8
import torch.nn.functional as F
9
from torch.nn.parameter import UninitializedParameter
10
11
-from vllm import envs
11
+import vllm.envs as envs
12
from vllm.config import get_current_vllm_config
13
from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
14
get_tensor_model_parallel_world_size,
@@ -16,6 +16,8 @@ from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
16
from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.logger import init_logger
18
from vllm.model_executor.custom_op import CustomOp
19
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
20
+ is_rocm_aiter_moe_enabled, shuffle_weights)
21
from vllm.model_executor.layers.quantization.base_config import (
22
QuantizationConfig, QuantizeMethodBase)
23
from vllm.model_executor.utils import set_weight_attrs
@@ -118,6 +120,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
120
layer.w2_weight.data),
121
requires_grad=False)
122
123
+ if is_rocm_aiter_moe_enabled():
124
+ # reshaping weights is required for aiter moe kernel.
125
+ shuffled_w13, shuffled_w2 = shuffle_weights(
126
+ layer.w13_weight.data, layer.w2_weight.data)
127
+
128
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13,
129
+ requires_grad=False)
130
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2,
131
+ requires_grad=False)
132
+
133
if current_platform.is_cpu():
134
if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
135
import intel_extension_for_pytorch as ipex
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
new
+157
@@ -0,0 +1,157 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import List, Optional
3
+
4
+import torch
5
+
6
+import vllm.envs as envs
7
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
8
+ per_token_group_quant_fp8)
9
+from vllm.platforms import current_platform
10
+
11
+
12
+def is_rocm_aiter_moe_enabled() -> bool:
13
+ return current_platform.is_rocm() \
14
+ and envs.VLLM_ROCM_USE_AITER_MOE \
15
+ and envs.VLLM_ROCM_USE_AITER \
16
+
17
+
18
+def is_rocm_aiter_block_scaled_moe_enabled() -> bool:
19
+ return is_rocm_aiter_moe_enabled() and \
20
+ envs.VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE
21
+
22
+
23
+def rocm_aiter_fused_experts(
24
+ *,
25
+ hidden_states: torch.Tensor,
26
+ w1: torch.Tensor,
27
+ w2: torch.Tensor,
28
+ topk_weights: torch.Tensor,
29
+ topk_ids: torch.Tensor,
30
+ use_fp8_w8a8: bool = False,
31
+ w1_scale: Optional[torch.Tensor] = None,
32
+ w2_scale: Optional[torch.Tensor] = None,
33
+ block_shape: Optional[List[int]] = None,
34
+ expert_mask: Optional[torch.Tensor] = None,
35
+ **kwagrs # Ignore additional keyword arguments
36
+) -> torch.Tensor:
37
+
38
+ import aiter as rocm_aiter
39
+ import aiter.fused_moe_bf16_asm as rocm_aiter_asm_fmoe
40
+
41
+ if envs.VLLM_ROCM_USE_AITER_FP8_BLOCK_SCALED_MOE and use_fp8_w8a8:
42
+ assert w1_scale is not None
43
+ assert w2_scale is not None
44
+
45
+ local_E = E = w1.shape[0]
46
+ if expert_mask is not None:
47
+ E = expert_mask.numel()
48
+
49
+ topk = topk_ids.shape[1]
50
+ model_dim = w1.shape[-1]
51
+ dtype = hidden_states.dtype
52
+ # The default block sizes are 128 in AITER.
53
+ if block_shape is None:
54
+ block_shape = [128, 128]
55
+
56
+ scale_blk_k = block_shape[1]
57
+
58
+ (
59
+ sorted_token_ids,
60
+ sorted_weight_buf,
61
+ sorted_expert_ids,
62
+ num_valid_ids,
63
+ out_asm,
64
+ ) = rocm_aiter_asm_fmoe.moe_sorting_ck(topk_ids,
65
+ topk_weights,
66
+ E,
67
+ model_dim,
68
+ dtype,
69
+ expert_mask=expert_mask)
70
+
71
+ a1, a1_scale = per_token_group_quant_fp8(hidden_states, scale_blk_k)
72
+ rocm_aiter.fmoe_fp8_blockscale_g1u1(
73
+ out_asm,
74
+ a1,
75
+ w1,
76
+ w2,
77
+ sorted_token_ids,
78
+ sorted_weight_buf,
79
+ sorted_expert_ids,
80
+ num_valid_ids,
81
+ topk,
82
+ w1_scale.view(local_E, -1),
83
+ w2_scale.view(local_E, -1),
84
+ a1_scale.t().contiguous(),
85
+ block_shape[0],
86
+ block_shape[1],
87
+ None,
88
+ )
89
+ return out_asm
90
+
91
+ elif use_fp8_w8a8:
92
+ return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
93
+ w1=w1,
94
+ w2=w2,
95
+ topk_weight=topk_weights,
96
+ topk_ids=topk_ids,
97
+ fc1_scale=w1_scale,
98
+ fc2_scale=w2_scale,
99
+ fc1_smooth_scale=None,
100
+ fc2_smooth_scale=None,
101
+ a16=False)
102
+
103
+ return rocm_aiter.ck_moe(hidden_states=hidden_states,
104
+ w1=w1,
105
+ w2=w2,
106
+ topk_weights=topk_weights,
107
+ topk_ids=topk_ids)
108
+
109
+
110
+def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
111
+ topk_indices: torch.Tensor,
112
+ token_expert_indices: torch.Tensor,
113
+ gating_output: torch.Tensor,
114
+ renormalize: bool) -> tuple[torch.Tensor, ...]:
115
+ import aiter as rocm_aiter
116
+ rocm_aiter.topk_softmax(topk_weights, topk_indices, token_expert_indices,
117
+ gating_output, renormalize)
118
+
119
+ return topk_weights, topk_indices
120
+
121
+
122
+def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
123
+ """
124
+ Applies shuffle_weight function from AITER to each
125
+ input tensor and returns them.
126
+
127
+ Args:
128
+ *tensors: Variable number of torch.Tensor objects.
129
+
130
+ Returns:
131
+ A tuple of shuffled tensors.
132
+ """
133
+ from aiter.ops.shuffle import shuffle_weight
134
+
135
+ return tuple(shuffle_weight(tensor) for tensor in tensors)
136
+
137
+
138
+def expand_weights(*tensors: torch.Tensor,
139
+ expansion_dims: list[int]) -> tuple[torch.Tensor, ...]:
140
+ """
141
+ Expands the dimensions of input tensors.
142
+
143
+ Args:
144
+ *tensors: A variable number of torch.Tensor objects.
145
+ expansion_dims: A list of expansion dimensions
146
+ corresponding to each tensor.
147
+
148
+ Returns:
149
+ A tuple of tensors with expanded dimensions.
150
+ """
151
+
152
+ assert len(tensors) == len(expansion_dims), \
153
+ "Number of tensors must match the number of expansion dimensions."
154
+
155
+ return tuple(
156
+ tensor.unsqueeze(-1).unsqueeze(-1).expand((-1, dim, -1))
157
+ for tensor, dim in zip(tensors, expansion_dims))
vllm/model_executor/layers/quantization/fp8.py
+52
@@ -13,6 +13,9 @@ from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
15
FusedMoeWeightScaleSupported)
16
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
17
+ expand_weights, is_rocm_aiter_block_scaled_moe_enabled,
18
+ is_rocm_aiter_moe_enabled, shuffle_weights)
19
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
20
UnquantizedLinearMethod)
21
from vllm.model_executor.layers.quantization.base_config import (
@@ -554,6 +557,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
557
layer.w2_weight = Parameter(w2_weight, requires_grad=False)
558
layer.w2_weight_scale_inv = Parameter(w2_weight_scale_inv,
559
requires_grad=False)
560
+ if is_rocm_aiter_block_scaled_moe_enabled():
561
+ # reshaping weights is required for aiter moe kernel.
562
+ shuffled_w13, shuffled_w2 = shuffle_weights(
563
+ layer.w13_weight.data, layer.w2_weight.data)
564
+
565
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13,
566
+ requires_grad=False)
567
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2,
568
+ requires_grad=False)
569
return
570
571
# If checkpoint is fp16, quantize in place.
@@ -581,6 +593,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):
593
requires_grad=False)
594
layer.w2_weight = torch.nn.Parameter(w2_weight,
595
requires_grad=False)
596
+ if is_rocm_aiter_moe_enabled():
597
+ # reshaping weights is required for aiter moe kernel.
598
+ w13_scales, w2_scales = expand_weights(
599
+ layer.w13_weight_scale.data,
600
+ layer.w2_weight_scale.data,
601
+ expansion_dims=[
602
+ layer.w13_weight.shape[1], layer.w2_weight.shape[1]
603
+ ])
604
+ layer.w13_weight_scale = torch.nn.Parameter(
605
+ w13_scales.contiguous(), requires_grad=False)
606
+ layer.w2_weight_scale = torch.nn.Parameter(
607
+ w2_scales.contiguous(), requires_grad=False)
608
+
609
+ shuffled_w13, shuffled_w2 = shuffle_weights(
610
+ layer.w13_weight, layer.w2_weight)
611
+
612
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13,
613
+ requires_grad=False)
614
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2,
615
+ requires_grad=False)
616
return
617
618
# If checkpoint is fp8, we need to handle that the
@@ -648,6 +680,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):
680
dq_weight, max_w13_scales[expert_id])
681
start += shard_size
682
683
+ if is_rocm_aiter_moe_enabled():
684
+ # reshaping weights is required for aiter moe kernel.
685
+ expansion_dims = [
686
+ layer.w13_weight.shape[1], layer.w2_weight.shape[1]
687
+ ]
688
+ max_w13_scales, w2_scales = expand_weights(
689
+ max_w13_scales,
690
+ layer.w2_weight_scale.data,
691
+ expansion_dims=expansion_dims)
692
+ layer.w2_weight_scale = torch.nn.Parameter(
693
+ w2_scales.contiguous(), requires_grad=False)
694
+
695
+ shuffled_w13, shuffled_w2 = shuffle_weights(
696
+ layer.w13_weight, layer.w2_weight)
697
+
698
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13,
699
+ requires_grad=False)
700
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2,
701
+ requires_grad=False)
702
+
703
layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
704
requires_grad=False)
705
return