@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 612c2edb4f

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, 11:00 UTC

[FEAT] [ROCm]: Add AITER CK 2 Stages MoE support (#17110)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

TJian committed May 14, 2025 at 18:03 UTC 612c2edb4f940b26550ca64576995957e4118a12

7 files changed +201 -112

tests/kernels/moe/test_moe.py

@@ -224,9 +224,16 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,

224 """Make sure our Mixtral MoE implementation agrees with the one from

225 huggingface."""

226

227 + # clear the cache before every test

228 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

229 + is_rocm_aiter_moe_enabled)

230 + is_rocm_aiter_moe_enabled.cache_clear()

231 if use_rocm_aiter:

232 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

233

234 + if dtype == torch.float32:

235 + pytest.skip("AITER ROCm test skip for float32")

236 +

237 # Instantiate our and huggingface's MoE blocks

238 config = MixtralConfig()

239 hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")

tests/model_executor/test_enabled_custom_ops.py

+2 -22

@@ -8,10 +8,8 @@ from vllm.model_executor.custom_op import CustomOp

8 from vllm.model_executor.layers.activation import (GeluAndMul,

9 ReLUSquaredActivation,

10 SiluAndMul)

11 -from vllm.model_executor.layers.fused_moe.fused_moe import (

12 - dispatch_fused_experts_func, dispatch_topk_func,

13 - torch_vllm_inplace_fused_experts, torch_vllm_outplace_fused_experts,

14 - vllm_topk_softmax)

11 +from vllm.model_executor.layers.fused_moe.fused_moe import (dispatch_topk_func,

12 + vllm_topk_softmax)

13 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

14 is_rocm_aiter_moe_enabled)

15 from vllm.model_executor.layers.layernorm import (

@@ -142,24 +140,6 @@ def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):

140 assert topk_func == vllm_topk_softmax

141

142

145 -@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])

146 -@pytest.mark.parametrize("inplace", [True, False])

147 -def test_fused_experts_dispatch(use_rocm_aiter: str, inplace: bool,

148 - monkeypatch):

149 -

150 - monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)

151 - is_rocm_aiter_moe_enabled.cache_clear()

152 - fused_experts_func = dispatch_fused_experts_func(inplace)

153 - if current_platform.is_rocm() and int(use_rocm_aiter):

154 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

155 - rocm_aiter_fused_experts)

156 - assert fused_experts_func == rocm_aiter_fused_experts

157 - elif inplace:

158 - assert fused_experts_func == torch_vllm_inplace_fused_experts

159 - else:

160 - assert fused_experts_func == torch_vllm_outplace_fused_experts

161 -

162 -

143 @pytest.mark.parametrize("add_residual", [True, False])

144 @pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])

145 @pytest.mark.parametrize("use_rocm_aiter_norm", ["0", "1"])

vllm/model_executor/layers/fused_moe/fused_moe.py

-3

@@ -1100,9 +1100,6 @@ def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:

1100

1101

1102 def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:

1103 - if is_rocm_aiter_moe_enabled():

1104 - from .rocm_aiter_fused_moe import rocm_aiter_fused_experts

1105 - return rocm_aiter_fused_experts

1103 if inplace:

1104 return torch_vllm_inplace_fused_experts

1105 return torch_vllm_outplace_fused_experts

vllm/model_executor/layers/fused_moe/layer.py

+27 -5

@@ -84,6 +84,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):

84 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

85 """MoE method without quantization."""

87 + def __init__(self):

88 + super().__init__()

89 +

90 + self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()

91 + if self.rocm_aiter_moe_enabled:

92 + from .rocm_aiter_fused_moe import rocm_aiter_fused_experts

93 + self.rocm_aiter_fused_experts = rocm_aiter_fused_experts

94 + else:

95 + self.rocm_aiter_fused_experts = None # type: ignore

96 +

97 def create_weights(self, layer: torch.nn.Module, num_experts: int,

98 hidden_size: int, intermediate_size_per_partition: int,

99 params_dtype: torch.dtype, **extra_weight_attrs):

@@ -126,11 +136,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

136 layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)

137 # Lazy import to avoid importing triton.

138 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

129 - is_rocm_aiter_moe_enabled, shuffle_weights)

130 - if is_rocm_aiter_moe_enabled():

131 - # reshaping weights is required for aiter moe kernel.

132 - shuffled_w13, shuffled_w2 = shuffle_weights(

133 - layer.w13_weight.data, layer.w2_weight.data)

139 + shuffle_weights)

140 +

141 + if self.rocm_aiter_moe_enabled:

142 + # use 2stage ck moe layout

143 + shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,

144 + layer.w2_weight.data,

145 + layout=(32, 32))

146

147 layer.w13_weight.data = shuffled_w13

148 layer.w2_weight.data = shuffled_w2

@@ -211,6 +223,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

223 scoring_func=scoring_func,

224 e_score_correction_bias=e_score_correction_bias)

225

226 + if self.rocm_aiter_moe_enabled:

227 + return self.rocm_aiter_fused_experts(

228 + hidden_states=x,

229 + w1=layer.w13_weight,

230 + w2=layer.w2_weight,

231 + topk_weights=topk_weights,

232 + topk_ids=topk_ids,

233 + activation=activation,

234 + apply_router_weight_on_input=apply_router_weight_on_input)

235 +

236 return fused_experts(

237 hidden_states=x,

238 w1=layer.w13_weight,

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+106 -67

@@ -20,7 +20,7 @@ def rocm_aiter_asm_moe_tkw1_impl(

20 hidden_states: torch.Tensor,

21 w1: torch.Tensor,

22 w2: torch.Tensor,

23 - topk_weight: torch.Tensor,

23 + topk_weights: torch.Tensor,

24 topk_ids: torch.Tensor,

25 fc1_scale: Optional[torch.Tensor] = None,

26 fc2_scale: Optional[torch.Tensor] = None,

@@ -40,7 +40,7 @@ def rocm_aiter_asm_moe_tkw1_impl(

40 return asm_moe_tkw1(hidden_states,

41 w1,

42 w2,

43 - topk_weight,

43 + topk_weights,

44 topk_ids,

45 fc1_scale=fc1_scale,

46 fc2_scale=fc2_scale,

@@ -56,7 +56,7 @@ def rocm_aiter_asm_moe_tkw1_fake(

56 hidden_states: torch.Tensor,

57 w1: torch.Tensor,

58 w2: torch.Tensor,

59 - topk_weight: torch.Tensor,

59 + topk_weights: torch.Tensor,

60 topk_ids: torch.Tensor,

61 fc1_scale: Optional[torch.Tensor] = None,

62 fc2_scale: Optional[torch.Tensor] = None,

@@ -69,23 +69,6 @@ def rocm_aiter_asm_moe_tkw1_fake(

69 return torch.empty_like(hidden_states)

72 -def rocm_aiter_ck_moe_impl(hidden_states: torch.Tensor, w1: torch.Tensor,

73 - w2: torch.Tensor, topk_weights: torch.Tensor,

74 - topk_ids: torch.Tensor) -> torch.Tensor:

75 - from aiter import ck_moe

76 - return ck_moe(hidden_states=hidden_states,

77 - w1=w1,

78 - w2=w2,

79 - topk_weights=topk_weights,

80 - topk_ids=topk_ids)

81 -

82 -

83 -def rocm_aiter_ck_moe_fake(hidden_states: torch.Tensor, w1: torch.Tensor,

84 - w2: torch.Tensor, topk_weights: torch.Tensor,

85 - topk_ids: torch.Tensor) -> torch.Tensor:

86 - return torch.empty_like(hidden_states)

87 -

88 -

72 def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(

73 topk_ids: torch.Tensor,

74 topk_weights: torch.Tensor,

@@ -152,7 +135,7 @@ def rocm_aiter_fmoe_fp8_blockscale_g1u1_fake(

135 def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,

136 w1: torch.Tensor,

137 w2: torch.Tensor,

155 - topk_weight: 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,

@@ -175,7 +158,7 @@ def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,

158 return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,

159 w1=w1,

160 w2=w2,

178 - topk_weight=topk_weight,

161 + topk_weight=topk_weights,

162 topk_ids=topk_ids,

163 fc1_scale=fc1_scale,

164 fc2_scale=fc2_scale,

@@ -188,7 +171,7 @@ def rocm_aiter_asm_moe_impl(hidden_states: torch.Tensor,

171 def rocm_aiter_asm_moe_fake(hidden_states: torch.Tensor,

172 w1: torch.Tensor,

173 w2: torch.Tensor,

191 - topk_weight: 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,

@@ -199,6 +182,49 @@ def rocm_aiter_asm_moe_fake(hidden_states: 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:

225 + return torch.empty_like(hidden_states)

226 +

227 +

228 def rocm_aiter_topk_softmax_impl(topk_weights: torch.Tensor,

229 topk_indices: torch.Tensor,

230 token_expert_indices: torch.Tensor,

@@ -258,14 +284,6 @@ if current_platform.is_rocm():

284 dispatch_key=current_platform.dispatch_key,

285 )

286

261 - direct_register_custom_op(

262 - op_name="rocm_aiter_ck_moe",

263 - op_func=rocm_aiter_ck_moe_impl,

264 - mutates_args=[],

265 - fake_impl=rocm_aiter_ck_moe_fake,

266 - dispatch_key=current_platform.dispatch_key,

267 - )

268 -

287 direct_register_custom_op(

288 op_name="rocm_aiter_fmoe_fp8_blockscale_g1u1",

289 op_func=rocm_aiter_fmoe_fp8_blockscale_g1u1_impl,

@@ -282,6 +300,14 @@ if current_platform.is_rocm():

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,

306 + mutates_args=[],

307 + fake_impl=rocm_aiter_ck_moe_2stages_fake,

308 + dispatch_key=current_platform.dispatch_key,

309 + )

310 +

311 direct_register_custom_op(

312 op_name="rocm_aiter_topk_softmax",

313 op_func=rocm_aiter_topk_softmax_impl,

@@ -331,29 +357,21 @@ def rocm_aiter_biased_group_topk(

357 return topk_weights, topk_ids

358

359

334 -def rocm_aiter_fused_experts(hidden_states: torch.Tensor,

335 - w1: torch.Tensor,

336 - w2: torch.Tensor,

337 - topk_weights: torch.Tensor,

338 - topk_ids: torch.Tensor,

339 - inplace: bool = False,

340 - activation: str = "silu",

341 - apply_router_weight_on_input: bool = False,

342 - use_fp8_w8a8: bool = False,

343 - use_int8_w8a8: bool = False,

344 - use_int8_w8a16: bool = False,

345 - use_int4_w4a16: bool = False,

346 - per_channel_quant: bool = False,

347 - global_num_experts: int = -1,

348 - expert_map: Optional[torch.Tensor] = None,

349 - w1_scale: Optional[torch.Tensor] = None,

350 - w2_scale: Optional[torch.Tensor] = None,

351 - w1_zp: Optional[torch.Tensor] = None,

352 - w2_zp: Optional[torch.Tensor] = None,

353 - a1_scale: Optional[torch.Tensor] = None,

354 - a2_scale: Optional[torch.Tensor] = None,

355 - block_shape: Optional[list[int]] = None,

356 - allow_deep_gemm: bool = False) -> torch.Tensor:

360 +def rocm_aiter_fused_experts(

361 + hidden_states: torch.Tensor,

362 + w1: torch.Tensor,

363 + w2: torch.Tensor,

364 + topk_weights: torch.Tensor,

365 + topk_ids: torch.Tensor,

366 + activation: str = "silu",

367 + apply_router_weight_on_input: bool = False,

368 + use_fp8_w8a8: bool = False,

369 + per_channel_quant: bool = False,

370 + w1_scale: Optional[torch.Tensor] = None,

371 + w2_scale: Optional[torch.Tensor] = None,

372 + a1_scale: Optional[torch.Tensor] = None,

373 + a2_scale: Optional[torch.Tensor] = None,

374 + block_shape: Optional[list[int]] = None) -> torch.Tensor:

375

376 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

377 per_token_group_quant_fp8)

@@ -376,8 +394,8 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,

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(

379 - topk_ids, topk_weights, hidden_states.dtype, expert_map, a1, w1,

380 - w2, w1_scale, w2_scale, a1_scale, block_shape, None)

397 + topk_ids, topk_weights, hidden_states.dtype, None, a1, w1, w2,

398 + w1_scale, w2_scale, a1_scale, block_shape, None)

399

400 # w8a8 per-channel quantization

401 elif per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:

@@ -402,17 +420,36 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,

420 fc2_smooth_scale=None,

421 a16=False,

422 per_tensor_quant_scale=None,

405 - expert_mask=expert_map,

423 + 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,

415 - topk_weight=topk_weights,

452 + topk_weights=topk_weights,

453 topk_ids=topk_ids,

454 fc1_scale=w1_scale,

455 fc2_scale=w2_scale,

@@ -432,12 +469,12 @@ def rocm_aiter_fused_experts(hidden_states: torch.Tensor,

469 topk_ids = topk_ids.to(torch.int32)

470 topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)

471

435 - # w16a16 fallback to rocm_aiter_ck_moe w16a16

436 - return torch.ops.vllm.rocm_aiter_ck_moe(hidden_states=hidden_states,

437 - w1=w1,

438 - w2=w2,

439 - topk_weights=topk_weights,

440 - topk_ids=topk_ids)

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)

478

479

480 def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,

@@ -451,7 +488,8 @@ def rocm_aiter_topk_softmax(topk_weights: torch.Tensor,

488 return topk_weights, topk_indices

489

490

454 -def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:

491 +def shuffle_weights(*tensors: torch.Tensor,

492 + layout: tuple[int, int]) -> tuple[torch.Tensor, ...]:

493 """

494 Applies shuffle_weight function from AITER to each

495 input tensor and returns them.

@@ -463,7 +501,8 @@ def shuffle_weights(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:

501 A Tuple of shuffled tensors.

502 """

503 from aiter.ops.shuffle import shuffle_weight

466 - return tuple(shuffle_weight(tensor) for tensor in tensors)

504 +

505 + return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)

506

507

508 def expand_weights(*tensors: torch.Tensor,

@@ -485,4 +524,4 @@ def expand_weights(*tensors: torch.Tensor,

524

525 return tuple(

526 tensor.unsqueeze(-1).unsqueeze(-1).expand((-1, dim, -1))

488 - for tensor, dim in zip(tensors, expansion_dims))

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

+25 -7

@@ -125,6 +125,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

125 # Disable marlin for rocm

126 if current_platform.is_rocm():

127 self.use_marlin = False

128 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

129 + is_rocm_aiter_moe_enabled)

130 +

131 + self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()

132

133 def create_weights(self, layer: torch.nn.Module, num_experts: int,

134 hidden_size: int, intermediate_size_per_partition: int,

@@ -276,24 +280,22 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

280 layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,

281 requires_grad=False)

282

279 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

280 - is_rocm_aiter_moe_enabled)

281 -

283 # Property to determine if AITER is used

283 - if is_rocm_aiter_moe_enabled():

284 + if self.rocm_aiter_moe_enabled:

285 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501

286 rocm_aiter_fused_experts, shuffle_weights)

287

288 # reshaping weights is required for aiter moe kernel.

288 - shuffled_w13, shuffled_w2 = shuffle_weights(

289 - layer.w13_weight.data, layer.w2_weight.data)

289 + shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight.data,

290 + layer.w2_weight.data,

291 + layout=(16, 16))

292

293 layer.w13_weight = torch.nn.Parameter(shuffled_w13,

294 requires_grad=False)

295 layer.w2_weight = torch.nn.Parameter(shuffled_w2,

296 requires_grad=False)

297

296 - self.fused_experts_func = rocm_aiter_fused_experts

298 + self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts

299 else:

300 from vllm.model_executor.layers.fused_moe import fused_experts

301 self.fused_experts_func = fused_experts

@@ -335,6 +337,22 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

337 scoring_func=scoring_func,

338 e_score_correction_bias=e_score_correction_bias)

339

340 + if self.rocm_aiter_moe_enabled:

341 + return self.rocm_aiter_fused_experts_func(

342 + hidden_states=x,

343 + w1=layer.w13_weight,

344 + w2=layer.w2_weight,

345 + topk_weights=topk_weights,

346 + topk_ids=topk_ids,

347 + activation=activation,

348 + apply_router_weight_on_input=apply_router_weight_on_input,

349 + use_fp8_w8a8=True,

350 + per_channel_quant=self.weight_quant.strategy ==

351 + QuantizationStrategy.CHANNEL,

352 + w1_scale=layer.w13_weight_scale,

353 + w2_scale=layer.w2_weight_scale,

354 + a1_scale=layer.w13_input_scale,

355 + a2_scale=layer.w2_input_scale)

356 if self.use_marlin:

357 assert activation == "silu", (

358 f"{activation} not supported for Marlin MoE.")

vllm/model_executor/layers/quantization/fp8.py

+34 -8

@@ -591,6 +591,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

591 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

592 expand_weights, is_rocm_aiter_moe_enabled, shuffle_weights)

593

594 + self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()

595 +

596 # TODO (rob): refactor block quant into separate class.

597 if self.block_quant:

598 assert self.quant_config.activation_scheme == "dynamic"

@@ -616,10 +618,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):

618 layer.w2_weight = Parameter(w2_weight, requires_grad=False)

619 layer.w2_weight_scale_inv = Parameter(w2_weight_scale_inv,

620 requires_grad=False)

619 - if is_rocm_aiter_moe_enabled():

621 + if self.rocm_aiter_moe_enabled:

622 # reshaping weights is required for aiter moe kernel.

623 shuffled_w13, shuffled_w2 = shuffle_weights(

622 - layer.w13_weight.data, layer.w2_weight.data)

624 + layer.w13_weight.data,

625 + layer.w2_weight.data,

626 + layout=(16, 16))

627

628 layer.w13_weight = torch.nn.Parameter(shuffled_w13,

629 requires_grad=False)

@@ -663,7 +667,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

667 requires_grad=False)

668 layer.w2_weight = torch.nn.Parameter(w2_weight,

669 requires_grad=False)

666 - if is_rocm_aiter_moe_enabled():

670 + if self.rocm_aiter_moe_enabled:

671 # reshaping weights is required for aiter moe kernel.

672 w13_scales, w2_scales = expand_weights(

673 layer.w13_weight_scale.data,

@@ -676,8 +680,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

680 layer.w2_weight_scale = torch.nn.Parameter(

681 w2_scales.contiguous(), requires_grad=False)

682

679 - shuffled_w13, shuffled_w2 = shuffle_weights(

680 - layer.w13_weight, layer.w2_weight)

683 + shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,

684 + layer.w2_weight,

685 + layout=(16, 16))

686

687 layer.w13_weight = torch.nn.Parameter(shuffled_w13,

688 requires_grad=False)

@@ -748,7 +753,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

753 dq_weight, max_w13_scales[expert_id])

754 start += shard_size

755

751 - if is_rocm_aiter_moe_enabled():

756 + if self.rocm_aiter_moe_enabled:

757 # reshaping weights is required for aiter moe kernel.

758 expansion_dims = [

759 layer.w13_weight.shape[1], layer.w2_weight.shape[1]

@@ -760,8 +765,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

765 layer.w2_weight_scale = torch.nn.Parameter(

766 w2_scales.contiguous(), requires_grad=False)

767

763 - shuffled_w13, shuffled_w2 = shuffle_weights(

764 - layer.w13_weight, layer.w2_weight)

768 + shuffled_w13, shuffled_w2 = shuffle_weights(layer.w13_weight,

769 + layer.w2_weight,

770 + layout=(32, 32))

771

772 layer.w13_weight = torch.nn.Parameter(shuffled_w13,

773 requires_grad=False)

@@ -796,6 +802,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

802 activation: str = "silu",

803 ) -> torch.Tensor:

804 from vllm.model_executor.layers.fused_moe import fused_experts

805 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

806 + rocm_aiter_fused_experts)

807

808 topk_weights, topk_ids = FusedMoE.select_experts(

809 hidden_states=x,

@@ -810,6 +818,24 @@ class Fp8MoEMethod(FusedMoEMethodBase):

818 e_score_correction_bias=e_score_correction_bias,

819 )

820

821 + if self.rocm_aiter_moe_enabled:

822 + return rocm_aiter_fused_experts(

823 + x,

824 + layer.w13_weight,

825 + layer.w2_weight,

826 + topk_weights=topk_weights,

827 + topk_ids=topk_ids,

828 + activation=activation,

829 + use_fp8_w8a8=True,

830 + apply_router_weight_on_input=apply_router_weight_on_input,

831 + w1_scale=(layer.w13_weight_scale_inv

832 + if self.block_quant else layer.w13_weight_scale),

833 + w2_scale=(layer.w2_weight_scale_inv

834 + if self.block_quant else layer.w2_weight_scale),

835 + a1_scale=layer.w13_input_scale,

836 + a2_scale=layer.w2_input_scale,

837 + block_shape=self.quant_config.weight_block_size)

838 +

839 if self.use_marlin:

840 assert activation == "silu", (

841 f"{activation} not supported for Marlin MoE.")