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, 08:30 UTC
[FEAT] [ROCm] Upgrade AITER Fused MoE kernels. (#18271)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
May 27, 2025 at 14:14 UTC
d260f799a95d1c49f1da17309dc42df1381f809b
4 files changed
+133
-317
vllm/model_executor/layers/fused_moe/layer.py
+2
-4
@@ -419,10 +419,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
419
shuffle_weights)
420
421
if self.rocm_aiter_moe_enabled:
422
- # use 2stage ck moe layout
423
- shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
424
- layer.w2_weight.data,
425
- layout=(32, 32))
422
+ shuffled_w13, shuffled_w2 = shuffle_weights(
423
+ layer.w13_weight.data, layer.w2_weight.data)
424
425
layer.w13_weight.data = shuffled_w13
426
layer.w2_weight.data = shuffled_w2
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+123
-278
@@ -1,4 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+from enum import IntEnum
3
from functools import cache
4
from typing import Optional
5
@@ -9,6 +10,28 @@ from vllm.platforms import current_platform
10
from vllm.utils import direct_register_custom_op
11
12
13
+class QuantMethod(IntEnum):
14
+ # This allows interfacing with AITER QuantType Enum
15
+ # without importing the QuantType from AITER globally.
16
+
17
+ # Note that these quantization methods are
18
+ # supported in AITER package. However,
19
+ # not all are used in this module.
20
+
21
+ NO = 0 # a16w16
22
+ PER_TENSOR = 1 # w8a8 (pre_Tensor)
23
+ PER_TOKEN = 2 # w8a8/w8a4 (per_Token)
24
+ BLOCK_1X128 = 3 # block quantized w8a8 (per_1x128)
25
+ BLOCK_128x128 = 4 # block quantized w8a8 (per_128x128)
26
+
27
+
28
+class ActivationMethod(IntEnum):
29
+ # This allows interfacing with AITER ActivationType enum
30
+ # without importing the ActivationType enum from AITER globally.
31
+ SILU = 0
32
+ GELU = 1
33
+
34
+
35
@cache
36
def is_rocm_aiter_moe_enabled() -> bool:
37
return current_platform.is_rocm() \
@@ -29,13 +52,12 @@ def rocm_aiter_asm_moe_tkw1_impl(
52
a16: bool = False,
53
per_tensor_quant_scale: Optional[torch.Tensor] = None,
54
expert_mask: Optional[torch.Tensor] = None,
32
- activation_str: str = "silu") -> torch.Tensor:
55
+ activation_method: int = ActivationMethod.SILU.value) -> torch.Tensor:
56
57
from aiter import ActivationType
58
from aiter.fused_moe_bf16_asm import asm_moe_tkw1
59
37
- activation = \
38
- ActivationType.Gelu if activation_str == "gelu" else ActivationType.Silu
60
+ activation = ActivationType(activation_method)
61
62
return asm_moe_tkw1(hidden_states,
63
w1,
@@ -65,163 +87,7 @@ def rocm_aiter_asm_moe_tkw1_fake(
87
a16: bool = False,
88
per_tensor_quant_scale: Optional[torch.Tensor] = None,
89
expert_mask: Optional[torch.Tensor] = None,
68
- activation_str: str = "silu") -> torch.Tensor:
69
- return torch.empty_like(hidden_states)
70
-
71
-
72
-def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(
73
- topk_ids: torch.Tensor,
74
- topk_weights: torch.Tensor,
75
- hidden_states_dtype: torch.dtype,
76
- expert_mask: torch.Tensor,
77
- a1: torch.Tensor,
78
- w1: torch.Tensor,
79
- w2: torch.Tensor,
80
- w1_scale: torch.Tensor,
81
- w2_scale: torch.Tensor,
82
- a1_scale: torch.Tensor,
83
- block_shape: list[int],
84
- smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:
85
- from aiter import fmoe_fp8_blockscale_g1u1
86
- from aiter.fused_moe_bf16_asm import moe_sorting_ck
87
-
88
- topk = topk_ids.shape[1]
89
- model_dim = w1.shape[-1]
90
- local_E = E = w1.shape[0]
91
- if expert_mask is not None:
92
- E = expert_mask.numel()
93
-
94
- (
95
- sorted_token_ids,
96
- sorted_weight_buf,
97
- sorted_expert_ids,
98
- num_valid_ids,
99
- out_asm,
100
- ) = moe_sorting_ck(topk_ids,
101
- topk_weights,
102
- E,
103
- model_dim,
104
- hidden_states_dtype,
105
- expert_mask=expert_mask)
106
-
107
- fmoe_fp8_blockscale_g1u1(out_asm, a1, w1, w2, sorted_token_ids,
108
- sorted_weight_buf, sorted_expert_ids,
109
- num_valid_ids, topk,
110
- a1_scale.t().contiguous(),
111
- w1_scale.view(local_E, -1),
112
- w2_scale.view(local_E,
113
- -1), *block_shape, smooth_scale)
114
-
115
- return out_asm
116
-
117
-
118
-def rocm_aiter_fmoe_fp8_blockscale_g1u1_fake(
119
- topk_ids: torch.Tensor,
120
- topk_weights: torch.Tensor,
121
- hidden_states_dtype: torch.dtype,
122
- expert_mask: torch.Tensor,
123
- a1: torch.Tensor,
124
- w1: torch.Tensor,
125
- w2: torch.Tensor,
126
- w1_scale: torch.Tensor,
127
- w2_scale: torch.Tensor,
128
- a1_scale: torch.Tensor,
129
- block_shape: list[int],
130
- smooth_scale: Optional[torch.Tensor] = None) -> torch.Tensor:
131
-
132
- return torch.empty_like(a1, dtype=hidden_states_dtype)
133
-
134
-
135
-def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,
136
- w1: torch.Tensor,
137
- w2: torch.Tensor,
138
- topk_weights: torch.Tensor,
139
- topk_ids: torch.Tensor,
140
- fc1_scale: Optional[torch.Tensor] = None,
141
- fc2_scale: Optional[torch.Tensor] = None,
142
- fc1_smooth_scale: Optional[torch.Tensor] = None,
143
- fc2_smooth_scale: Optional[torch.Tensor] = None,
144
- a16: bool = False,
145
- activation: str = "silu") -> torch.Tensor:
146
- import aiter.fused_moe_bf16_asm as rocm_aiter_asm_fmoe
147
- from aiter import ActivationType
148
-
149
- assert activation in ["silu", "gelu"], "The given activation:" \
150
- f" {activation}" \
151
- " is not supported in" \
152
- " AITER."
153
- if activation == "silu":
154
- aiter_activation = ActivationType.Silu
155
- else:
156
- aiter_activation = ActivationType.Gelu
157
-
158
- return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
159
- w1=w1,
160
- w2=w2,
161
- topk_weight=topk_weights,
162
- topk_ids=topk_ids,
163
- fc1_scale=fc1_scale,
164
- fc2_scale=fc2_scale,
165
- fc1_smooth_scale=fc1_smooth_scale,
166
- fc2_smooth_scale=fc2_smooth_scale,
167
- a16=a16,
168
- activation=aiter_activation)
169
-
170
-
171
-def rocm_aiter_asm_moe_fake(hidden_states: torch.Tensor,
172
- w1: torch.Tensor,
173
- w2: torch.Tensor,
174
- topk_weights: torch.Tensor,
175
- topk_ids: torch.Tensor,
176
- fc1_scale: Optional[torch.Tensor] = None,
177
- fc2_scale: Optional[torch.Tensor] = None,
178
- fc1_smooth_scale: Optional[torch.Tensor] = None,
179
- fc2_smooth_scale: Optional[torch.Tensor] = None,
180
- a16: bool = False,
181
- activation: str = "silu") -> torch.Tensor:
182
- return torch.empty_like(hidden_states)
183
-
184
-
185
-def rocm_aiter_ck_moe_2stages_impl(
186
- hidden_states: torch.Tensor,
187
- w1: torch.Tensor,
188
- w2: torch.Tensor,
189
- topk_weights: torch.Tensor,
190
- topk_ids: torch.Tensor,
191
- fc1_scale: Optional[torch.Tensor] = None,
192
- fc2_scale: Optional[torch.Tensor] = None,
193
- a1_scale: Optional[torch.Tensor] = None,
194
- a2_scale: Optional[torch.Tensor] = None,
195
- block_size: Optional[list[int]] = None,
196
- expert_mask: Optional[torch.Tensor] = None,
197
-) -> torch.Tensor:
198
- from aiter.fused_moe_bf16_asm import ck_moe_2stages
199
- return ck_moe_2stages(a1=hidden_states,
200
- w1=w1,
201
- w2=w2,
202
- topk_weight=topk_weights,
203
- topk_ids=topk_ids,
204
- fc1_scale=fc1_scale,
205
- fc2_scale=fc2_scale,
206
- a1_scale=a1_scale,
207
- a2_scale=a2_scale,
208
- block_size=block_size,
209
- expert_mask=expert_mask)
210
-
211
-
212
-def rocm_aiter_ck_moe_2stages_fake(
213
- hidden_states: torch.Tensor,
214
- w1: torch.Tensor,
215
- w2: torch.Tensor,
216
- topk_weights: torch.Tensor,
217
- topk_ids: torch.Tensor,
218
- fc1_scale: Optional[torch.Tensor] = None,
219
- fc2_scale: Optional[torch.Tensor] = None,
220
- a1_scale: Optional[torch.Tensor] = None,
221
- a2_scale: Optional[torch.Tensor] = None,
222
- block_size: Optional[list[int]] = None,
223
- expert_mask: Optional[torch.Tensor] = None,
224
-) -> torch.Tensor:
90
+ activation_method: int = ActivationMethod.SILU.value) -> torch.Tensor:
91
return torch.empty_like(hidden_states)
92
93
@@ -274,6 +140,50 @@ def rocm_aiter_biased_grouped_topk_fake(
140
pass
141
142
143
+def rocm_aiter_fused_moe_impl(
144
+ hidden_states: torch.Tensor,
145
+ w1: torch.Tensor,
146
+ w2: torch.Tensor,
147
+ topk_weight: torch.Tensor,
148
+ topk_ids: torch.Tensor,
149
+ expert_mask: Optional[torch.Tensor] = None,
150
+ activation_method: int = ActivationMethod.SILU.value,
151
+ quant_method: int = QuantMethod.NO.value,
152
+ doweight_stage1: bool = False,
153
+ w1_scale: Optional[torch.Tensor] = None,
154
+ w2_scale: Optional[torch.Tensor] = None,
155
+ a1_scale: Optional[torch.Tensor] = None,
156
+ a2_scale: Optional[torch.Tensor] = None,
157
+) -> torch.Tensor:
158
+ from aiter import ActivationType, QuantType
159
+ from aiter.fused_moe import fused_moe
160
+
161
+ activation = ActivationType(activation_method)
162
+ quant_type = QuantType(quant_method)
163
+
164
+ return fused_moe(hidden_states, w1, w2, topk_weight, topk_ids, expert_mask,
165
+ activation, quant_type, doweight_stage1, w1_scale,
166
+ w2_scale, a1_scale, a2_scale)
167
+
168
+
169
+def rocm_aiter_fused_moe_fake(
170
+ hidden_states: torch.Tensor,
171
+ w1: torch.Tensor,
172
+ w2: torch.Tensor,
173
+ topk_weight: torch.Tensor,
174
+ topk_ids: torch.Tensor,
175
+ expert_mask: Optional[torch.Tensor] = None,
176
+ activation_method: int = ActivationMethod.SILU.value,
177
+ quant_method: int = QuantMethod.NO.value,
178
+ doweight_stage1: bool = False,
179
+ w1_scale: Optional[torch.Tensor] = None,
180
+ w2_scale: Optional[torch.Tensor] = None,
181
+ a1_scale: Optional[torch.Tensor] = None,
182
+ a2_scale: Optional[torch.Tensor] = None,
183
+) -> torch.Tensor:
184
+ return torch.empty_like(hidden_states)
185
+
186
+
187
if current_platform.is_rocm():
188
189
direct_register_custom_op(
@@ -285,26 +195,10 @@ if current_platform.is_rocm():
195
)
196
197
direct_register_custom_op(
288
- op_name="rocm_aiter_fmoe_fp8_blockscale_g1u1",
289
- op_func=rocm_aiter_fmoe_fp8_blockscale_g1u1_impl,
290
- mutates_args=[],
291
- fake_impl=rocm_aiter_fmoe_fp8_blockscale_g1u1_fake,
292
- dispatch_key=current_platform.dispatch_key,
293
- )
294
-
295
- direct_register_custom_op(
296
- op_name="rocm_aiter_asm_moe",
297
- op_func=rocm_aiter_asm_moe_impl,
298
- mutates_args=[],
299
- fake_impl=rocm_aiter_asm_moe_fake,
300
- dispatch_key=current_platform.dispatch_key,
301
- )
302
-
303
- direct_register_custom_op(
304
- op_name="rocm_aiter_ck_moe_2stages",
305
- op_func=rocm_aiter_ck_moe_2stages_impl,
198
+ op_name="rocm_aiter_fused_moe",
199
+ op_func=rocm_aiter_fused_moe_impl,
200
mutates_args=[],
307
- fake_impl=rocm_aiter_ck_moe_2stages_fake,
201
+ fake_impl=rocm_aiter_fused_moe_fake,
202
dispatch_key=current_platform.dispatch_key,
203
)
204
@@ -373,32 +267,14 @@ def rocm_aiter_fused_experts(
267
a2_scale: Optional[torch.Tensor] = None,
268
block_shape: Optional[list[int]] = None) -> torch.Tensor:
269
376
- from vllm.model_executor.layers.quantization.utils.fp8_utils import (
377
- per_token_group_quant_fp8)
378
-
270
+ activation_method = (ActivationMethod.SILU
271
+ if activation == "silu" else ActivationMethod.GELU)
272
# All AITER Fused MoE kernels are expecting the following datatypes
273
topk_weights = topk_weights.to(torch.float32)
274
topk_ids = topk_ids.to(torch.int32)
275
383
- # w8a8 block-scaled
384
- if block_shape is not None and use_fp8_w8a8:
385
- assert not apply_router_weight_on_input, (
386
- "apply_router_weight_on_input is not supported for block scaled moe"
387
- )
388
- assert w1_scale is not None
389
- assert w2_scale is not None
390
-
391
- # The default block sizes are 128 in AITER.
392
- block_shape = [128, 128] if block_shape is None else block_shape
393
-
394
- a1, a1_scale = per_token_group_quant_fp8(hidden_states, block_shape[1])
395
-
396
- return torch.ops.vllm.rocm_aiter_fmoe_fp8_blockscale_g1u1(
397
- topk_ids, topk_weights, hidden_states.dtype, None, a1, w1, w2,
398
- w1_scale, w2_scale, a1_scale, block_shape, None)
399
-
276
# w8a8 per-channel quantization
401
- elif per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
277
+ if per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
278
# AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
279
# This applies topk_weights on the GEMM output of the first FC layer
280
# rather than the second FC.
@@ -421,60 +297,44 @@ def rocm_aiter_fused_experts(
297
a16=False,
298
per_tensor_quant_scale=None,
299
expert_mask=None,
424
- activation_str=activation)
425
-
426
- # w8a8 per-tensor activation per-tensor weight
427
- elif use_fp8_w8a8:
428
- assert not apply_router_weight_on_input, (
429
- "apply_router_weight_on_input is not supported for fp8_w8a8")
430
-
431
- # - faster static per-tensor-activation static per-tensor-weight
432
- # fp8 quantization w8a8
433
- if a1_scale is not None and a2_scale is not None:
434
- return torch.ops.vllm.rocm_aiter_ck_moe_2stages(
435
- hidden_states=hidden_states,
436
- w1=w1,
437
- w2=w2,
438
- topk_weights=topk_weights,
439
- topk_ids=topk_ids,
440
- fc1_scale=w1_scale,
441
- fc2_scale=w2_scale,
442
- a1_scale=a1_scale,
443
- a2_scale=a2_scale)
444
-
445
- # - fallback static per-tensor-activation static per-tensor-weight
446
- # fp8 quantization w8a8
447
- # - dynamic per-tensor activation static per-tensor-weight
448
- # fp8 quantization w8a8
449
- return torch.ops.vllm.rocm_aiter_asm_moe(hidden_states=hidden_states,
450
- w1=w1,
451
- w2=w2,
452
- topk_weights=topk_weights,
453
- topk_ids=topk_ids,
454
- fc1_scale=w1_scale,
455
- fc2_scale=w2_scale,
456
- fc1_smooth_scale=None,
457
- fc2_smooth_scale=None,
458
- a16=False,
459
- activation=activation)
460
- if apply_router_weight_on_input:
461
- assert (topk_weights.dim() == 2
462
- ), "`topk_weights` should be in shape (num_tokens, topk)"
463
- _, topk = topk_weights.shape
464
- assert (
465
- topk == 1
466
- ), "Only support topk=1 when `apply_router_weight_on_input` is True"
467
-
468
- hidden_states = hidden_states * topk_weights.to(hidden_states.dtype)
469
- topk_ids = topk_ids.to(torch.int32)
470
- topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)
300
+ activation_method=activation_method)
301
472
- return torch.ops.vllm.rocm_aiter_ck_moe_2stages(
473
- hidden_states=hidden_states,
474
- w1=w1,
475
- w2=w2,
476
- topk_weights=topk_weights,
477
- topk_ids=topk_ids)
302
+ else:
303
+ quant_method = QuantMethod.NO.value
304
+
305
+ # w8a8 block-scaled
306
+ if block_shape is not None and use_fp8_w8a8:
307
+ assert not apply_router_weight_on_input, (
308
+ "apply_router_weight_on_input is\
309
+ not supported for block scaled moe")
310
+ assert w1_scale is not None
311
+ assert w2_scale is not None
312
+ quant_method = QuantMethod.BLOCK_128x128.value
313
+ elif use_fp8_w8a8:
314
+ # Currently only per tensor quantization method is enabled.
315
+ quant_method = QuantMethod.PER_TENSOR.value
316
+
317
+ if apply_router_weight_on_input:
318
+ assert (topk_weights.dim() == 2
319
+ ), "`topk_weights` should be in shape (num_tokens, topk)"
320
+ _, topk = topk_weights.shape
321
+ assert (
322
+ topk == 1
323
+ ), "Only support topk=1 when `apply_router_weight_on_input` is True"
324
+
325
+ return torch.ops.vllm.rocm_aiter_fused_moe(
326
+ hidden_states,
327
+ w1,
328
+ w2,
329
+ topk_weights,
330
+ topk_ids,
331
+ quant_method=quant_method,
332
+ activation_method=activation_method,
333
+ w1_scale=w1_scale,
334
+ w2_scale=w2_scale,
335
+ a1_scale=a1_scale,
336
+ a2_scale=a2_scale,
337
+ doweight_stage1=apply_router_weight_on_input)
338
339
340
def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
@@ -488,14 +348,21 @@ def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,
348
return topk_weights, topk_indices
349
350
491
-def shuffle_weights(*tensors: torch.Tensor,
492
- layout: tuple[int, int]) -> tuple[torch.Tensor, ...]:
351
+def shuffle_weights(
352
+ *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
353
+) -> tuple[torch.Tensor, ...]:
354
"""
355
Applies shuffle_weight function from AITER to each
356
input tensor and returns them.
357
+
358
+ Rearranges (shuffles) the input tensor/s
359
+ into a specified block layout for optimized computation.
360
361
Args:
498
- *tensors: Variable number of torch.Tensor objects.
362
+ *tensors: Variable number of torch.Tensor objects.
363
+ layout: A pair of integers specifying the
364
+ block sizes used to divide the tensors during shuffling.
365
+ Default is (16, 16).
366
367
Returns:
368
A Tuple of shuffled tensors.
@@ -503,25 +370,3 @@ def shuffle_weights(*tensors: torch.Tensor,
370
from aiter.ops.shuffle import shuffle_weight
371
372
return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
506
-
507
-
508
-def expand_weights(*tensors: torch.Tensor,
509
- expansion_dims: list[int]) -> tuple[torch.Tensor, ...]:
510
- """
511
- Expands the dimensions of input tensors.
512
-
513
- Args:
514
- *tensors: A variable number of torch.Tensor objects.
515
- expansion_dims: A list of expansion dimensions
516
- corresponding to each tensor.
517
-
518
- Returns:
519
- A Tuple of tensors with expanded dimensions.
520
- """
521
-
522
- assert len(tensors) == len(expansion_dims), \
523
- "Number of tensors must match the number of expansion dimensions."
524
-
525
- return tuple(
526
- tensor.unsqueeze(-1).unsqueeze(-1).expand((-1, dim, -1))
527
- for tensor, dim in zip(tensors, expansion_dims))
\ No newline at end of file
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+2
-3
@@ -286,9 +286,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
286
rocm_aiter_fused_experts, shuffle_weights)
287
288
# reshaping weights is required for aiter moe kernel.
289
- shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,
290
- layer.w2_weight.data,
291
- layout=(16, 16))
289
+ shuffled_w13, shuffled_w2 = shuffle_weights(
290
+ layer.w13_weight.data, layer.w2_weight.data)
291
292
layer.w13_weight = torch.nn.Parameter(shuffled_w13,
293
requires_grad=False)
vllm/model_executor/layers/quantization/fp8.py
+6
-32
@@ -595,7 +595,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
595
def process_weights_after_loading(self, layer: Module) -> None:
596
# Lazy import to avoid importing triton too early.
597
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
598
- expand_weights, is_rocm_aiter_moe_enabled, shuffle_weights)
598
+ is_rocm_aiter_moe_enabled, shuffle_weights)
599
600
self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
601
@@ -627,9 +627,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
627
if self.rocm_aiter_moe_enabled:
628
# reshaping weights is required for aiter moe kernel.
629
shuffled_w13, shuffled_w2 = shuffle_weights(
630
- layer.w13_weight.data,
631
- layer.w2_weight.data,
632
- layout=(16, 16))
630
+ layer.w13_weight.data, layer.w2_weight.data)
631
632
layer.w13_weight = torch.nn.Parameter(shuffled_w13,
633
requires_grad=False)
@@ -675,20 +673,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
673
requires_grad=False)
674
if self.rocm_aiter_moe_enabled:
675
# reshaping weights is required for aiter moe kernel.
678
- w13_scales, w2_scales = expand_weights(
679
- layer.w13_weight_scale.data,
680
- layer.w2_weight_scale.data,
681
- expansion_dims=[
682
- layer.w13_weight.shape[1], layer.w2_weight.shape[1]
683
- ])
684
- layer.w13_weight_scale = torch.nn.Parameter(
685
- w13_scales.contiguous(), requires_grad=False)
686
- layer.w2_weight_scale = torch.nn.Parameter(
687
- w2_scales.contiguous(), requires_grad=False)
688
-
689
- shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,
690
- layer.w2_weight,
691
- layout=(16, 16))
676
+ shuffled_w13, shuffled_w2 = shuffle_weights(
677
+ layer.w13_weight, layer.w2_weight)
678
679
layer.w13_weight = torch.nn.Parameter(shuffled_w13,
680
requires_grad=False)
@@ -760,20 +746,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
746
start += shard_size
747
748
if self.rocm_aiter_moe_enabled:
763
- # reshaping weights is required for aiter moe kernel.
764
- expansion_dims = [
765
- layer.w13_weight.shape[1], layer.w2_weight.shape[1]
766
- ]
767
- max_w13_scales, w2_scales = expand_weights(
768
- max_w13_scales,
769
- layer.w2_weight_scale.data,
770
- expansion_dims=expansion_dims)
771
- layer.w2_weight_scale = torch.nn.Parameter(
772
- w2_scales.contiguous(), requires_grad=False)
773
-
774
- shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,
775
- layer.w2_weight,
776
- layout=(32, 32))
749
+ shuffled_w13, shuffled_w2 = shuffle_weights(
750
+ layer.w13_weight, layer.w2_weight)
751
752
layer.w13_weight = torch.nn.Parameter(shuffled_w13,
753
requires_grad=False)