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

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

[MoE Refactor][7/N] AITER MK (#31102)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>

Robert Shaw committed Dec 22, 2025 at 18:42 UTC b57b967386e1962955c93b7cb39828448b68789b

4 files changed +144 -66

vllm/model_executor/layers/fused_moe/fused_moe.py

+5 -1

@@ -2132,6 +2132,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

2132 torch.float16,

2133 torch.bfloat16,

2134 torch.float8_e4m3fn,

2135 + torch.float8_e4m3fnuz,

2136 ]

2137

2138 E, num_tokens, N, K, top_k_num = self.moe_problem_size(

@@ -2156,7 +2157,10 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

2157 compute_type = tl.float16

2158 elif hidden_states.dtype == torch.float32:

2159 compute_type = tl.float32

2159 - elif hidden_states.dtype == torch.float8_e4m3fn:

2160 + elif (

2161 + hidden_states.dtype == torch.float8_e4m3fn

2162 + or hidden_states.dtype == torch.float8_e4m3fnuz

2163 + ):

2164 compute_type = tl.bfloat16

2165 else:

2166 raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

vllm/model_executor/layers/fused_moe/prepare_finalize.py

@@ -13,6 +13,10 @@ from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

15 class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

16 + def __init__(self, defer_input_quant: bool = False) -> None:

17 + super().__init__()

18 + self.defer_input_quant = defer_input_quant

19 +

20 @property

21 def activation_format(self) -> mk.FusedMoEActivationFormat:

22 return mk.FusedMoEActivationFormat.Standard

@@ -48,6 +52,11 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

52 # Note: do not use inplace for shared experts overlap

53 a1 = a1 * topk_weights.to(a1.dtype)

55 + # Defer input quant to moe kernel for backends (e.g. AITER, FI)

56 + # which use a single kernel call for quant + experts.

57 + if self.defer_input_quant:

58 + return a1, None, None, None, None

59 +

60 a1q, a1q_scale = moe_kernel_quantize_input(

61 a1,

62 quant_config.a1_scale,

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+79

@@ -5,11 +5,15 @@ from functools import lru_cache

6 import torch

8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.model_executor.layers.fused_moe.config import (

11 FUSED_MOE_UNQUANTIZED_CONFIG,

12 FusedMoEQuantConfig,

13 )

14 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

15 + TopKWeightAndReduceNoOP,

16 +)

19 class QuantMethod(IntEnum):

@@ -263,3 +267,78 @@ def rocm_aiter_fused_experts(

267 a2_scale=quant_config.a2_scale,

268 doweight_stage1=apply_router_weight_on_input,

269 )

270 +

271 +

272 +class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

273 + def __init__(self, quant_config):

274 + super().__init__(quant_config)

275 +

276 + @property

277 + def activation_formats(

278 + self,

279 + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

280 + return (

281 + mk.FusedMoEActivationFormat.Standard,

282 + mk.FusedMoEActivationFormat.Standard,

283 + )

284 +

285 + def supports_expert_map(self):

286 + return True

287 +

288 + def supports_chunking(self):

289 + return False

290 +

291 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

292 + return TopKWeightAndReduceNoOP()

293 +

294 + def workspace_shapes(

295 + self,

296 + M: int,

297 + N: int,

298 + K: int,

299 + topk: int,

300 + global_num_experts: int,

301 + local_num_experts: int,

302 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

303 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

304 + # Workspaces are managed internally by AITER.

305 + workspace1 = (0,)

306 + workspace2 = (0,)

307 + output = (M, K)

308 + return (workspace1, workspace2, output)

309 +

310 + def apply(

311 + self,

312 + output: torch.Tensor,

313 + hidden_states: torch.Tensor,

314 + w1: torch.Tensor,

315 + w2: torch.Tensor,

316 + topk_weights: torch.Tensor,

317 + topk_ids: torch.Tensor,

318 + activation: str,

319 + global_num_experts: int,

320 + expert_map: torch.Tensor | None,

321 + a1q_scale: torch.Tensor | None,

322 + a2_scale: torch.Tensor | None,

323 + workspace13: torch.Tensor,

324 + workspace2: torch.Tensor,

325 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

326 + apply_router_weight_on_input: bool,

327 + ):

328 + assert a1q_scale is None

329 + assert a2_scale is None

330 + assert expert_tokens_meta is None

331 +

332 + result = rocm_aiter_fused_experts(

333 + hidden_states=hidden_states,

334 + w1=w1,

335 + w2=w2,

336 + topk_weights=topk_weights,

337 + topk_ids=topk_ids,

338 + activation=activation,

339 + apply_router_weight_on_input=apply_router_weight_on_input,

340 + expert_map=expert_map,

341 + quant_config=self.quant_config,

342 + )

343 + assert result.shape == output.shape

344 + output.copy_(result)

vllm/model_executor/layers/quantization/fp8.py

+51 -65

@@ -117,6 +117,7 @@ class Fp8MoeBackend(Enum):

117 DEEPGEMM = 3

118 MARLIN = 4

119 TRITON = 5

120 + AITER = 6

121

122

123 def get_fp8_moe_backend(

@@ -189,6 +190,10 @@ def get_fp8_moe_backend(

190 logger.info_once("Using DeepGEMM backend for FP8 MoE", scope="local")

191 return Fp8MoeBackend.DEEPGEMM

192

193 + if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MOE:

194 + logger.info_once("Using ROCm AITER backend for FP8 MoE", scope="local")

195 + return Fp8MoeBackend.AITER

196 +

197 # default to Triton

198 logger.info_once("Using Triton backend for FP8 MoE")

199 return Fp8MoeBackend.TRITON

@@ -888,16 +893,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

893 layer.w13_input_scale = None

894 layer.w2_input_scale = None

895

891 - self.rocm_aiter_moe_enabled = False

892 -

896 def process_weights_after_loading(self, layer: Module) -> None:

897 if getattr(layer, "_already_called_process_weights_after_loading", False):

898 return

899

897 - # Lazy import to avoid importing triton too early.

898 -

899 - self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()

900 -

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

901 if self.block_quant:

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

@@ -932,7 +931,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

931 replace_parameter(layer, "w13_weight_scale_inv", w13_weight_scale_inv)

932 replace_parameter(layer, "w2_weight", w2_weight)

933 replace_parameter(layer, "w2_weight_scale_inv", w2_weight_scale_inv)

935 - if self.rocm_aiter_moe_enabled:

934 + if self.fp8_backend == Fp8MoeBackend.AITER:

935 # reshaping weights is required for aiter moe kernel.

936 shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

937 layer.w13_weight.data, layer.w2_weight.data

@@ -1026,7 +1025,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1025 )

1026 start += shard_size

1027

1029 - if self.rocm_aiter_moe_enabled:

1028 + if self.fp8_backend == Fp8MoeBackend.AITER:

1029 shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1030 layer.w13_weight, layer.w2_weight

1031 )

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

1071 self.moe_quant_config = config

1072

1073 self.kernel = mk.FusedMoEModularKernel(

1074 + # TODO(rob): we can use the generic MoEPrepareAndFinalizeNoEP

1075 + # with the changes to defer input quantization

1076 FlashInferAllGatherMoEPrepareAndFinalize(

1077 use_dp=(self.moe.dp_size > 1),

1078 use_deepseek_fp8_block_scale=self.block_quant,

@@ -1093,6 +1094,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1094 Fp8MoeBackend.DEEPGEMM,

1095 Fp8MoeBackend.TRITON,

1096 Fp8MoeBackend.MARLIN,

1097 + Fp8MoeBackend.AITER,

1098 ]:

1099 from vllm.model_executor.layers.fused_moe import (

1100 TritonOrDeepGemmExperts,

@@ -1103,24 +1105,33 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1105 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

1106 MoEPrepareAndFinalizeNoEP,

1107 )

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

1109 + AiterExperts,

1110 + )

1111

1112 config = self.get_fused_moe_quant_config(layer)

1113 assert config is not None

1114 self.moe_quant_config = config

1110 - use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN

1111 - allow_deep_gemm = self.fp8_backend == Fp8MoeBackend.DEEPGEMM

1112 - moe_kernel = (

1113 - MarlinExperts(quant_config=self.moe_quant_config)

1114 - if use_marlin

1115 - else TritonOrDeepGemmExperts(

1116 - quant_config=self.moe_quant_config,

1117 - allow_deep_gemm=allow_deep_gemm,

1118 - )

1119 - )

1115

1121 - self.kernel = mk.FusedMoEModularKernel(

1122 - MoEPrepareAndFinalizeNoEP(), moe_kernel

1123 - )

1116 + if self.fp8_backend == Fp8MoeBackend.AITER:

1117 + self.kernel = mk.FusedMoEModularKernel(

1118 + # TODO: make defer_input_quant an attr of the AiterExperts

1119 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

1120 + AiterExperts(quant_config=self.moe_quant_config),

1121 + )

1122 + elif self.fp8_backend == Fp8MoeBackend.MARLIN:

1123 + self.kernel = mk.FusedMoEModularKernel(

1124 + MoEPrepareAndFinalizeNoEP(),

1125 + MarlinExperts(quant_config=self.moe_quant_config),

1126 + )

1127 + else:

1128 + self.kernel = mk.FusedMoEModularKernel(

1129 + MoEPrepareAndFinalizeNoEP(),

1130 + TritonOrDeepGemmExperts(

1131 + quant_config=self.moe_quant_config,

1132 + allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),

1133 + ),

1134 + )

1135 self.use_inplace = True

1136

1137 def maybe_make_prepare_finalize(

@@ -1128,7 +1139,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1139 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1140 ) -> mk.FusedMoEPrepareAndFinalize | None:

1141 if (

1131 - self.rocm_aiter_moe_enabled

1142 + self.fp8_backend == Fp8MoeBackend.AITER

1143 or self.fp8_backend == Fp8MoeBackend.MARLIN

1144 or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1145 ):

@@ -1161,11 +1172,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1172 TritonOrDeepGemmExperts,

1173 )

1174

1164 - assert (

1165 - self.fp8_backend != Fp8MoeBackend.MARLIN

1166 - ) and not self.rocm_aiter_moe_enabled, (

1167 - "Marlin and ROCm AITER are not supported with all2all yet."

1168 - )

1175 + if self.fp8_backend in [Fp8MoeBackend.MARLIN, Fp8MoeBackend.AITER]:

1176 + raise NotImplementedError(

1177 + "Marlin and ROCm AITER are not supported with all2all yet."

1178 + )

1179

1180 assert self.moe_quant_config is not None

1181

@@ -1313,37 +1323,18 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1323 hidden_states=x,

1324 router_logits=router_logits,

1325 )

1316 -

1317 - if self.rocm_aiter_moe_enabled:

1318 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501

1319 - rocm_aiter_fused_experts,

1320 - )

1321 -

1322 - # TODO(rob): convert this to MK.

1323 - result = rocm_aiter_fused_experts(

1324 - x,

1325 - layer.w13_weight,

1326 - layer.w2_weight,

1327 - topk_weights=topk_weights,

1328 - topk_ids=topk_ids,

1329 - activation=layer.activation,

1330 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1331 - expert_map=layer.expert_map,

1332 - quant_config=self.moe_quant_config,

1333 - )

1334 - else:

1335 - result = self.kernel(

1336 - x,

1337 - layer.w13_weight,

1338 - layer.w2_weight,

1339 - topk_weights,

1340 - topk_ids,

1341 - inplace=self.use_inplace,

1342 - activation=layer.activation,

1343 - global_num_experts=layer.global_num_experts,

1344 - expert_map=layer.expert_map,

1345 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1346 - )

1326 + result = self.kernel(

1327 + x,

1328 + layer.w13_weight,

1329 + layer.w2_weight,

1330 + topk_weights,

1331 + topk_ids,

1332 + inplace=self.use_inplace,

1333 + activation=layer.activation,

1334 + global_num_experts=layer.global_num_experts,

1335 + expert_map=layer.expert_map,

1336 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

1337 + )

1338

1339 return result

1340

@@ -1456,15 +1447,10 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1447 layer.w13_input_scale = None

1448 layer.w2_input_scale = None

1449

1459 - self.rocm_aiter_moe_enabled = False

1460 -

1450 def process_weights_after_loading(self, layer: Module) -> None:

1451 if getattr(layer, "_already_called_process_weights_after_loading", False):

1452 return

1453

1465 - # Lazy import to avoid importing triton too early.

1466 - self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()

1467 -

1454 # If checkpoint is fp16, quantize in place.

1455 fp8_dtype = current_platform.fp8_dtype()

1456 w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)

@@ -1481,7 +1467,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1467 replace_parameter(layer, "w2_weight", w2_weight)

1468

1469 # Reshuffle weights for AITER if needed.

1484 - if self.rocm_aiter_moe_enabled:

1470 + if self.fp8_backend == Fp8MoeBackend.AITER:

1471 shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1472 layer.w13_weight, layer.w2_weight

1473 )

@@ -1489,7 +1475,7 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1475 replace_parameter(layer, "w2_weight", shuffled_w2)

1476

1477 # Rushuffle weights for MARLIN if needed.

1492 - if self.fp8_backend == Fp8MoeBackend.MARLIN:

1478 + elif self.fp8_backend == Fp8MoeBackend.MARLIN:

1479 prepare_moe_fp8_layer_for_marlin(

1480 layer, False, input_dtype=self.marlin_input_dtype

1481 )