@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1e9500410a
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 4, 07:30 UTC

[ROCm][Quantization][2/N] Refactor quark_moe w4a8 w/ oracle (#39136)

Signed-off-by: Bowen Bao <bowenbao@amd.com>

Bowen Bao committed May 4, 2026 at 19:50 UTC 1e9500410a21782847ae86561b4de7f3aa69f0bc
6 files changed +875 -347
tests/kernels/moe/test_ocp_mx_moe.py
+388 -8
@@ -28,6 +28,25 @@ HOPPER_MXFP4_BF16_AVAILABLE = (
28 and has_flashinfer()
29 )
30
31 +# ROCm platform and dependencies
32 +ROCM_AVAILABLE = current_platform.is_rocm()
33 +ROCM_TRITON_KERNELS_AVAILABLE = False
34 +ROCM_AITER_AVAILABLE = False
35 +ROCM_GFX950 = False
36 +
37 +if ROCM_AVAILABLE:
38 + from vllm._aiter_ops import rocm_aiter_ops
39 + from vllm.platforms.rocm import on_gfx950
40 + from vllm.utils.import_utils import has_triton_kernels
41 +
42 + ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
43 + ROCM_GFX950 = on_gfx950()
44 + ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled()
45 +
46 + if ROCM_AITER_AVAILABLE:
47 + from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp
48 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
49 +
50 if TRTLLM_GEN_MXFP4_AVAILABLE:
51 from flashinfer import (
52 fp4_quantize,
@@ -111,6 +130,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
130
131 def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = None):
132 # Note we add an extra bias of 1 to the linear layer
133 + # Uses chunked layout: first half is gate, second half is up
134 x_glu, x_linear = torch.chunk(x, 2, dim=-1)
135 if limit is not None:
136 x_glu = x_glu.clamp(max=limit)
@@ -119,6 +139,16 @@ def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = Non
139 return out_glu * (x_linear + beta)
140
141
142 +def swigluoai(x, alpha: float = 1.702, limit: float = 7.0):
143 + # OAI swiglu uses interleaved layout: gate/up alternating
144 + # See SwigluOAIAndMul in vllm/model_executor/layers/activation.py
145 + gate, up = x[..., ::2], x[..., 1::2]
146 + gate = gate.clamp(max=limit)
147 + up = up.clamp(min=-limit, max=limit)
148 + glu = gate * torch.sigmoid(gate * alpha)
149 + return (up + 1) * glu
150 +
151 +
152 fp4_lookup_table = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6]
153
154
@@ -168,8 +198,20 @@ def reference_moe(
198 beta,
199 limit,
200 act_type,
171 - is_gated,
201 + activation: str = "swiglu",
202 + use_interleaved_layout: bool = False,
203 ):
204 + """
205 + Reference MoE implementation for accuracy testing.
206 +
207 + Args:
208 + activation: One of "swiglu", "silu", "relu2". Controls the activation
209 + function used after the first MLP.
210 + use_interleaved_layout: If True, uses interleaved gate/up layout
211 + (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI.
212 + If False, uses chunked layout (gate, up = chunk(x, 2)) as used
213 + by standard swiglu/silu.
214 + """
215 # renormalize routing
216 experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True)
217 expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
@@ -179,12 +221,21 @@ def reference_moe(
221 mlp1_weight = w13[expert_indices, ...]
222 mlp1_bias = bias13[expert_indices, ...]
223 t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias
182 - if is_gated:
183 - t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
184 - else:
224 +
225 + # Apply activation
226 + if activation in ("swiglu", "silu"):
227 + if use_interleaved_layout:
228 + # SWIGLUOAI: interleaved gate/up layout
229 + t = swigluoai(t, alpha=alpha, limit=limit)
230 + else:
231 + # Standard swiglu/silu: chunked layout
232 + t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
233 + elif activation == "relu2":
234 # RELU2_NO_MUL: relu(x)^2
235 t = torch.relu(t)
236 t = t * t
237 + else:
238 + raise ValueError(f"Unknown activation: {activation}")
239
240 if act_type == "mxfp8":
241 t_quantized, t_scale = mxfp8_quantize(
@@ -585,7 +636,8 @@ def test_trtllm_gen_mxfp4_fused_moe(
636 beta,
637 limit,
638 act_type,
588 - is_gated=True,
639 + activation="swiglu",
640 + use_interleaved_layout=False,
641 )
642 ref_result[start_idx:end_idx].copy_(chunk_result)
643
@@ -722,7 +774,8 @@ def test_flashinfer_cutlass_mxfp4_fused_moe(
774 beta,
775 limit,
776 "bf16",
725 - is_gated=True,
777 + activation="swiglu",
778 + use_interleaved_layout=False,
779 )
780
781 from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
@@ -908,7 +961,8 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(
961 beta,
962 limit,
963 "mxfp8",
911 - is_gated=True,
964 + activation="swiglu",
965 + use_interleaved_layout=False,
966 )
967
968 # Prepare inputs for FlashInfer CUTLASS fused MoE
@@ -1080,7 +1134,8 @@ def test_trtllm_gen_mxfp8_block_scale_moe(
1134 beta=0.0,
1135 limit=None,
1136 act_type="mxfp8",
1083 - is_gated=is_gated,
1137 + activation="swiglu" if is_gated else "relu2",
1138 + use_interleaved_layout=False,
1139 )
1140
1141 # Shuffle weights/scales with the same indexed layout used by TRTLLM kernels.
@@ -1150,3 +1205,328 @@ def test_trtllm_gen_mxfp8_block_scale_moe(
1205
1206 # Block-scale MXFP8 kernels are approximate; require majority close.
1207 check_accuracy(ref, out, atol=0.1, rtol=0.85, percent=0.8)
1208 +
1209 +
1210 +# -----------------------------------------------------------------------------
1211 +# ROCm Oracle-based kernel execution tests
1212 +# -----------------------------------------------------------------------------
1213 +# TODO: Further tighten the accuracy threshold.
1214 +# - More accurate ref moe to include activation quantization
1215 +# - Check aiter kernel accuracy. E.g., quant / dequant details.
1216 +ROCM_BACKEND_CONFIGS = {
1217 + "TRITON": {
1218 + "activation": "SWIGLUOAI",
1219 + "rtol": 0.3,
1220 + "percent": 0.95,
1221 + "requires_aiter": False,
1222 + "requires_gfx950": False,
1223 + },
1224 + "TRITON_UNFUSED": {
1225 + "activation": "SWIGLUOAI",
1226 + "rtol": 0.3,
1227 + "percent": 0.95,
1228 + "requires_aiter": False,
1229 + "requires_gfx950": False,
1230 + },
1231 + "AITER_MXFP4_BF16": {
1232 + "activation": "SILU",
1233 + "rtol": 1.0,
1234 + "percent": 0.7,
1235 + "requires_aiter": True,
1236 + "requires_gfx950": True,
1237 + },
1238 + "AITER_MXFP4_FP8": {
1239 + "activation": "SWIGLUOAI",
1240 + "rtol": 0.5,
1241 + "percent": 0.9,
1242 + "requires_aiter": True,
1243 + "requires_gfx950": True,
1244 + },
1245 +}
1246 +
1247 +
1248 +@pytest.mark.parametrize("backend_name", list(ROCM_BACKEND_CONFIGS.keys()))
1249 +@pytest.mark.parametrize("topk", [4])
1250 +@pytest.mark.parametrize("num_experts", [8])
1251 +@pytest.mark.parametrize("num_tokens,hidden_size,intermediate_size", [(16, 256, 256)])
1252 +@pytest.mark.skipif(
1253 + not ROCM_AVAILABLE,
1254 + reason="ROCm is required for this test",
1255 +)
1256 +@torch.inference_mode()
1257 +def test_rocm_mxfp4_moe_oracle(
1258 + backend_name: str,
1259 + topk: int,
1260 + num_experts: int,
1261 + num_tokens: int,
1262 + hidden_size: int,
1263 + intermediate_size: int,
1264 +):
1265 + """
1266 + Test ROCm MXFP4 MoE using oracle functions.
1267 +
1268 + This test validates that the oracle functions work end-to-end:
1269 + - select_mxfp4_moe_backend() selects a valid backend
1270 + - convert_to_mxfp4_moe_kernel_format() converts weights without error
1271 + - make_mxfp4_moe_quant_config() builds a valid quant config
1272 + - make_mxfp4_moe_kernel() creates a kernel that runs without error
1273 + - The kernel output is within accuracy tolerance of reference
1274 + """
1275 + config = ROCM_BACKEND_CONFIGS[backend_name]
1276 +
1277 + # Check platform requirements
1278 + if not ROCM_TRITON_KERNELS_AVAILABLE:
1279 + pytest.skip("triton_kernels required for quantization")
1280 + if config["requires_aiter"] and not ROCM_AITER_AVAILABLE:
1281 + pytest.skip(f"Backend {backend_name} requires AITER")
1282 + if config["requires_gfx950"] and not ROCM_GFX950:
1283 + pytest.skip(f"Backend {backend_name} requires GFX950")
1284 +
1285 + from vllm.config import VllmConfig, set_current_vllm_config
1286 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
1287 + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1288 + Mxfp4MoeBackend,
1289 + backend_to_kernel_cls,
1290 + convert_to_mxfp4_moe_kernel_format,
1291 + make_mxfp4_moe_kernel,
1292 + make_mxfp4_moe_quant_config,
1293 + )
1294 + from vllm.v1.worker.workspace import init_workspace_manager
1295 +
1296 + # Initialize workspace manager (needed for modular kernels)
1297 + init_workspace_manager(torch.accelerator.current_device_index())
1298 +
1299 + # Map string to enum
1300 + backend = Mxfp4MoeBackend[backend_name]
1301 +
1302 + # Get experts class from oracle
1303 + experts_cls_list = backend_to_kernel_cls(backend)
1304 + if experts_cls_list is None or len(experts_cls_list) == 0:
1305 + pytest.skip(f"Backend {backend_name} not available")
1306 +
1307 + # Use first experts class
1308 + experts_cls = experts_cls_list[0]
1309 +
1310 + torch.manual_seed(42)
1311 + dtype = torch.bfloat16
1312 + device = "cuda:0"
1313 +
1314 + # Create MoE config with Renormalize routing (required by monolithic kernels)
1315 + from vllm.model_executor.layers.fused_moe import FusedMoEConfig
1316 + from vllm.model_executor.layers.fused_moe.config import (
1317 + FusedMoEParallelConfig,
1318 + RoutingMethodType,
1319 + )
1320 +
1321 + moe_config = FusedMoEConfig(
1322 + num_experts=num_experts,
1323 + experts_per_token=topk,
1324 + hidden_dim=hidden_size,
1325 + intermediate_size_per_partition=intermediate_size,
1326 + num_local_experts=num_experts,
1327 + num_logical_experts=num_experts,
1328 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1329 + activation=MoEActivation[config["activation"]],
1330 + in_dtype=dtype,
1331 + device="cuda",
1332 + routing_method=RoutingMethodType.Renormalize,
1333 + )
1334 +
1335 + # Create float weights in checkpoint format:
1336 + # w13: [num_experts, 2*intermediate_size, hidden_size]
1337 + # w2: [num_experts, hidden_size, intermediate_size]
1338 + w13_float = torch.randn(
1339 + num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device
1340 + )
1341 + w2_float = torch.randn(
1342 + num_experts, hidden_size, intermediate_size, dtype=dtype, device=device
1343 + )
1344 +
1345 + # dynamic_mxfp4_quant expects 2D input, so reshape 3D weights
1346 + # w13: [E, 2*I, H] -> [E*2*I, H] -> quantize -> [E, 2*I, H//2]
1347 + # w2: [E, H, I] -> [E*H, I] -> quantize -> [E, H, I//2]
1348 + w13_2d = w13_float.reshape(-1, hidden_size)
1349 + w13_quant_2d, w13_scale_2d = dynamic_mxfp4_quant(w13_2d)
1350 + w13_quant = w13_quant_2d.reshape(num_experts, 2 * intermediate_size, -1)
1351 + w13_scale = w13_scale_2d.reshape(num_experts, 2 * intermediate_size, -1)
1352 +
1353 + w2_2d = w2_float.reshape(-1, intermediate_size)
1354 + w2_quant_2d, w2_scale_2d = dynamic_mxfp4_quant(w2_2d)
1355 + w2_quant = w2_quant_2d.reshape(num_experts, hidden_size, -1)
1356 + w2_scale = w2_scale_2d.reshape(num_experts, hidden_size, -1)
1357 +
1358 + w13_bias = torch.randn(
1359 + num_experts, 2 * intermediate_size, dtype=dtype, device=device
1360 + )
1361 + w2_bias = torch.randn(num_experts, hidden_size, dtype=dtype, device=device)
1362 +
1363 + # Create static input scales for W4A8 backend (AITER_MXFP4_FP8)
1364 + w13_input_scale: torch.Tensor | None = None
1365 + w2_input_scale: torch.Tensor | None = None
1366 + if backend_name == "AITER_MXFP4_FP8":
1367 + # Static FP8 scales: one scale per expert
1368 + w13_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device)
1369 + w2_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device)
1370 +
1371 + # Create mock layer for oracle functions
1372 + class MockLayer:
1373 + w13_weight: torch.Tensor
1374 + w2_weight: torch.Tensor
1375 + w13_weight_scale: torch.Tensor
1376 + w2_weight_scale: torch.Tensor
1377 + w13_input_scale: torch.Tensor | None
1378 + w2_input_scale: torch.Tensor | None
1379 +
1380 + layer = MockLayer()
1381 + layer.w13_weight = w13_quant
1382 + layer.w2_weight = w2_quant
1383 + layer.w13_weight_scale = w13_scale
1384 + layer.w2_weight_scale = w2_scale
1385 + layer.w13_input_scale = w13_input_scale
1386 + layer.w2_input_scale = w2_input_scale
1387 +
1388 + # Convert weights using oracle
1389 + w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = (
1390 + convert_to_mxfp4_moe_kernel_format(
1391 + mxfp4_backend=backend,
1392 + layer=layer, # type: ignore[arg-type]
1393 + w13_weight=w13_quant,
1394 + w2_weight=w2_quant,
1395 + w13_weight_scale=w13_scale,
1396 + w2_weight_scale=w2_scale,
1397 + w13_bias=w13_bias,
1398 + w2_bias=w2_bias,
1399 + )
1400 + )
1401 +
1402 + # Build quant config using oracle
1403 + quant_config = make_mxfp4_moe_quant_config(
1404 + mxfp4_backend=backend,
1405 + w1_scale=w13_scale_conv,
1406 + w2_scale=w2_scale_conv,
1407 + w1_bias=w13_bias_conv,
1408 + w2_bias=w2_bias_conv,
1409 + a1_scale=w13_input_scale,
1410 + a2_scale=w2_input_scale,
1411 + )
1412 +
1413 + # Select activation based on backend
1414 + activation_name = str(config["activation"])
1415 + activation = MoEActivation[activation_name]
1416 +
1417 + # Build kernel using oracle
1418 + assert quant_config is not None, "Failed to create quant config"
1419 + with set_current_vllm_config(VllmConfig()):
1420 + kernel = make_mxfp4_moe_kernel(
1421 + moe_quant_config=quant_config,
1422 + moe_config=moe_config,
1423 + mxfp4_backend=backend,
1424 + experts_cls=experts_cls,
1425 + routing_tables=None,
1426 + shared_experts=None,
1427 + )
1428 +
1429 + # Create inputs
1430 + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
1431 + router_logits = torch.randn(
1432 + num_tokens, num_experts, dtype=torch.float32, device=device
1433 + )
1434 + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)
1435 + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
1436 +
1437 + # Run kernel - use appropriate method based on impl type
1438 + if kernel.is_monolithic:
1439 + # Monolithic impl uses router_logits
1440 + out = kernel.apply_monolithic(
1441 + hidden_states=x,
1442 + w1=w13_conv,
1443 + w2=w2_conv,
1444 + router_logits=router_logits,
1445 + activation=activation,
1446 + global_num_experts=num_experts,
1447 + expert_map=None,
1448 + apply_router_weight_on_input=False,
1449 + )
1450 + else:
1451 + # Modular impl uses topk_weights and topk_ids
1452 + out = kernel.apply(
1453 + hidden_states=x,
1454 + w1=w13_conv,
1455 + w2=w2_conv,
1456 + topk_weights=topk_weights,
1457 + topk_ids=topk_ids,
1458 + activation=activation,
1459 + global_num_experts=num_experts,
1460 + expert_map=None,
1461 + apply_router_weight_on_input=False,
1462 + )
1463 +
1464 + # Verify output is valid (no NaN/Inf) and has expected shape
1465 + assert out.shape == (num_tokens, hidden_size), f"Unexpected shape: {out.shape}"
1466 + assert not torch.any(torch.isnan(out)), "Output contains NaN"
1467 + assert not torch.any(torch.isinf(out)), "Output contains Inf"
1468 +
1469 + # Verify output has reasonable magnitude (not all zeros)
1470 + assert out.abs().max() > 0.01, "Output is effectively zero"
1471 +
1472 + # Dequantize weights for reference computation
1473 + w13_dq = upcast_from_mxfp(
1474 + w13_quant.view(torch.uint8), w13_scale, torch.bfloat16, axis=-1
1475 + )
1476 + w2_dq = upcast_from_mxfp(
1477 + w2_quant.view(torch.uint8), w2_scale, torch.bfloat16, axis=-1
1478 + )
1479 +
1480 + # Determine activation type and layout
1481 + # SWIGLUOAI uses interleaved layout (gate/up alternating)
1482 + # SILU uses chunked layout (first half gate, second half up)
1483 + use_interleaved = activation == MoEActivation.SWIGLUOAI
1484 + if activation in [MoEActivation.SWIGLUOAI, MoEActivation.SILU]:
1485 + act_name = "swiglu"
1486 + else:
1487 + act_name = "relu2"
1488 +
1489 + ref = reference_moe(
1490 + router_logits,
1491 + topk,
1492 + num_experts,
1493 + x.to(torch.float32),
1494 + w13_dq.to(torch.float32),
1495 + w13_bias.to(torch.float32),
1496 + w2_dq.to(torch.float32),
1497 + w2_bias.to(torch.float32),
1498 + alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0,
1499 + beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0,
1500 + limit=7.0 if activation == MoEActivation.SWIGLUOAI else None,
1501 + act_type="bf16",
1502 + activation=act_name,
1503 + use_interleaved_layout=use_interleaved,
1504 + )
1505 +
1506 + # Compute and print accuracy statistics
1507 + diff = (ref.float() - out.float()).abs()
1508 + rel_diff = diff / (ref.float().abs() + 1e-6)
1509 +
1510 + print(f"\n[{backend_name}] Accuracy statistics:")
1511 + print(
1512 + f" Reference: min={ref.min():.4f}, max={ref.max():.4f}, mean={ref.mean():.4f}"
1513 + )
1514 + print(
1515 + f" Output: min={out.min():.4f}, max={out.max():.4f}, mean={out.mean():.4f}"
1516 + )
1517 + print(
1518 + f" Abs diff: min={diff.min():.4f}, max={diff.max():.4f}, "
1519 + f"mean={diff.mean():.4f}"
1520 + )
1521 + print(
1522 + f" Rel diff: min={rel_diff.min():.4f}, max={rel_diff.max():.4f}, "
1523 + f"mean={rel_diff.mean():.4f}"
1524 + )
1525 +
1526 + # Check what percentage of values are within various tolerances
1527 + for rtol in [0.1, 0.5, 1.0, 2.0]:
1528 + within_tol = (diff <= rtol * out.float().abs()).float().mean()
1529 + print(f" Within rtol={rtol}: {within_tol * 100:.1f}%")
1530 +
1531 + # Check accuracy using per-backend thresholds
1532 + check_accuracy(ref, out, atol=0.1, rtol=config["rtol"], percent=config["percent"])
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py new
+292
@@ -0,0 +1,292 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm._aiter_ops import rocm_aiter_ops
8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + FusedMoEConfig,
11 + FusedMoEParallelConfig,
12 + FusedMoEQuantConfig,
13 + RoutingMethodType,
14 +)
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + QuantKey,
17 + kFp8StaticTensorSym,
18 + kMxfp4Static,
19 +)
20 +
21 +__all__ = [
22 + "AiterW4A8ExpertsMonolithic",
23 + "aiter_triton_kernel_w4a8_moe_forward",
24 +]
25 +
26 +
27 +def aiter_triton_kernel_w4a8_moe_forward(
28 + hidden_states: torch.Tensor,
29 + w1, # Tensor or triton_kernels.Tensor
30 + w2, # Tensor or triton_kernels.Tensor
31 + gating_output: torch.Tensor,
32 + topk: int,
33 + renormalize: bool,
34 + activation: MoEActivation = MoEActivation.SWIGLUOAI,
35 + quant_config: FusedMoEQuantConfig | None = None,
36 + apply_router_weight_on_input: bool = False,
37 + global_num_experts: int = -1,
38 + expert_map: torch.Tensor | None = None,
39 + unpadded_N_w1=None,
40 + unpadded_K_w1=None,
41 + unpadded_N_w2=None,
42 + unpadded_K_w2=None,
43 +):
44 + assert (
45 + quant_config is not None
46 + and quant_config.use_mxfp4_w4a8
47 + and rocm_aiter_ops.is_enabled()
48 + )
49 + from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
50 +
51 + routing_data, gather_idx, scatter_idx = aiter_routing(
52 + gating_output, topk, sm_first=not renormalize
53 + )
54 + return triton_kernel_fused_mxfp4_w4a8_experts(
55 + None,
56 + hidden_states,
57 + w1,
58 + w2,
59 + routing_data,
60 + gather_idx,
61 + scatter_idx,
62 + activation=activation.value,
63 + quant_config=quant_config,
64 + apply_router_weight_on_input=apply_router_weight_on_input,
65 + global_num_experts=global_num_experts,
66 + expert_map=expert_map,
67 + unpadded_N_w1=unpadded_N_w1,
68 + unpadded_K_w1=unpadded_K_w1,
69 + unpadded_N_w2=unpadded_N_w2,
70 + unpadded_K_w2=unpadded_K_w2,
71 + )
72 +
73 +
74 +def triton_kernel_fused_mxfp4_w4a8_experts(
75 + output_tensor: torch.Tensor,
76 + hidden_states: torch.Tensor,
77 + w1, # Tensor or triton_kernels.Tensor
78 + w2, # Tensor or triton_kernels.Tensor
79 + routing_data, # RoutingData
80 + gather_indx, # GatherIndx
81 + scatter_indx, # ScatterIndx
82 + activation: str = "silu",
83 + quant_config: FusedMoEQuantConfig | None = None,
84 + swiglu_alpha: float = 1.702,
85 + swiglu_limit: float = 7.0,
86 + apply_router_weight_on_input: bool = False,
87 + global_num_experts: int = -1,
88 + expert_map: torch.Tensor | None = None,
89 + a1q_scale: torch.Tensor | None = None,
90 + unpadded_N_w1=None,
91 + unpadded_K_w1=None,
92 + unpadded_N_w2=None,
93 + unpadded_K_w2=None,
94 +) -> torch.Tensor:
95 + assert quant_config is not None
96 + # type check, uint8 means mxfp4
97 + assert hidden_states.dtype == torch.bfloat16
98 + assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
99 + assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
100 +
101 + # Shape check: weights are padded (e.g. hidden_size padded for
102 + # GFX950 swizzle).
103 + assert hidden_states.shape[-1] == w1.shape[-2]
104 + assert w2.shape[-1] == w1.shape[1]
105 +
106 + E, _, N = w1.shape
107 +
108 + if global_num_experts == -1:
109 + global_num_experts = E
110 +
111 + gammas = routing_data.gate_scal if routing_data else None
112 +
113 + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
114 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8
115 +
116 + assert quant_config.w1_precision is not None, (
117 + "w1_precision in quant config can't be None"
118 + )
119 + assert quant_config.w2_precision is not None, (
120 + "w2_precision in quant config can't be None"
121 + )
122 +
123 + hidden_states = downcast_to_static_fp8(
124 + hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
125 + )
126 +
127 + intermediate_cache1 = moe_gemm_a8w4(
128 + hidden_states,
129 + w1.storage.data,
130 + None,
131 + quant_config.w1_precision.weight_scale.storage.data,
132 + quant_config.w1_precision.flex_ctx.lhs_data.scale,
133 + quant_config.w2_precision.flex_ctx.lhs_data.scale,
134 + quant_config.w1_bias,
135 + routing_data,
136 + gather_indx=gather_indx,
137 + gammas=gammas if apply_router_weight_on_input else None,
138 + swizzle_mx_scale="CDNA4_SCALE",
139 + out_dtype=torch.float8_e4m3fn,
140 + apply_swiglu=True,
141 + alpha=swiglu_alpha,
142 + limit=swiglu_limit,
143 + unpadded_N=unpadded_N_w1,
144 + unpadded_K=unpadded_K_w1,
145 + )
146 +
147 + intermediate_cache3 = moe_gemm_a8w4(
148 + intermediate_cache1,
149 + w2.storage.data,
150 + None,
151 + quant_config.w2_precision.weight_scale.storage.data,
152 + quant_config.w2_precision.flex_ctx.lhs_data.scale,
153 + None,
154 + quant_config.w2_bias,
155 + routing_data,
156 + scatter_indx=scatter_indx,
157 + gammas=None if apply_router_weight_on_input else gammas,
158 + swizzle_mx_scale="CDNA4_SCALE",
159 + unpadded_N=unpadded_N_w2,
160 + unpadded_K=unpadded_K_w2,
161 + )
162 +
163 + return intermediate_cache3
164 +
165 +
166 +class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
167 + """
168 + Monolithic MXFP4 W4A8 expert using AITER triton kernels.
169 +
170 + This backend uses:
171 + - aiter.ops.triton.moe_routing.routing for routing
172 + - aiter.ops.triton.moe_op_gemm_a8w4.moe_gemm_a8w4 for computation
173 +
174 + Weight format: MXFP4 weights with GFX950 swizzle
175 + Activation: Static FP8 quantization
176 + """
177 +
178 + def __init__(
179 + self,
180 + moe_config: FusedMoEConfig,
181 + quant_config: FusedMoEQuantConfig,
182 + ):
183 + super().__init__(moe_config, quant_config)
184 + self.topk = moe_config.experts_per_token
185 + self.renormalize = moe_config.routing_method in (
186 + RoutingMethodType.Renormalize,
187 + RoutingMethodType.RenormalizeNaive,
188 + )
189 +
190 + @staticmethod
191 + def activation_format() -> mk.FusedMoEActivationFormat:
192 + return mk.FusedMoEActivationFormat.Standard
193 +
194 + @staticmethod
195 + def _supports_current_device() -> bool:
196 + # Requires AITER and GFX950
197 + if not rocm_aiter_ops.is_enabled():
198 + return False
199 + from vllm.platforms.rocm import on_gfx950
200 +
201 + return on_gfx950()
202 +
203 + @staticmethod
204 + def _supports_no_act_and_mul() -> bool:
205 + return False
206 +
207 + @staticmethod
208 + def _supports_quant_scheme(
209 + weight_key: QuantKey | None,
210 + activation_key: QuantKey | None,
211 + ) -> bool:
212 + # W4A8: MXFP4 weights with static FP8 activations
213 + SUPPORTED_W_A = [
214 + (kMxfp4Static, kFp8StaticTensorSym),
215 + ]
216 + return (weight_key, activation_key) in SUPPORTED_W_A
217 +
218 + @staticmethod
219 + def _supports_activation(activation: MoEActivation) -> bool:
220 + # Only SILU activation (swiglu) is supported
221 + return activation == MoEActivation.SWIGLUOAI
222 +
223 + @staticmethod
224 + def _supports_parallel_config(
225 + moe_parallel_config: FusedMoEParallelConfig,
226 + ) -> bool:
227 + return (
228 + not moe_parallel_config.use_all2all_kernels
229 + and not moe_parallel_config.enable_eplb
230 + and moe_parallel_config.dp_size <= 1
231 + )
232 +
233 + @staticmethod
234 + def _supports_routing_method(
235 + routing_method: RoutingMethodType,
236 + weight_key: QuantKey | None,
237 + activation_key: QuantKey | None,
238 + ) -> bool:
239 + return routing_method in [
240 + RoutingMethodType.Renormalize,
241 + RoutingMethodType.RenormalizeNaive,
242 + ]
243 +
244 + @staticmethod
245 + def _supports_router_logits_dtype(
246 + router_logits_dtype: torch.dtype | None,
247 + routing_method: RoutingMethodType,
248 + ) -> bool:
249 + return True
250 +
251 + def supports_expert_map(self) -> bool:
252 + return False # Expert parallelism not yet supported
253 +
254 + @property
255 + def expects_unquantized_inputs(self) -> bool:
256 + return True
257 +
258 + def apply(
259 + self,
260 + hidden_states: torch.Tensor,
261 + w1: torch.Tensor,
262 + w2: torch.Tensor,
263 + router_logits: torch.Tensor,
264 + activation: MoEActivation,
265 + global_num_experts: int,
266 + expert_map: torch.Tensor | None,
267 + a1q_scale: torch.Tensor | None,
268 + apply_router_weight_on_input: bool,
269 + # grouped topk + fused topk bias parameters
270 + num_expert_group: int | None = None,
271 + e_score_correction_bias: torch.Tensor | None = None,
272 + routed_scaling_factor: float | None = None,
273 + topk_group: int | None = None,
274 + ) -> torch.Tensor:
275 + assert self.moe_config.intermediate_size_per_partition_unpadded is not None
276 + assert self.moe_config.hidden_dim_unpadded is not None
277 + return aiter_triton_kernel_w4a8_moe_forward(
278 + hidden_states=hidden_states,
279 + w1=w1,
280 + w2=w2,
281 + gating_output=router_logits,
282 + topk=self.topk,
283 + renormalize=self.renormalize,
284 + global_num_experts=global_num_experts,
285 + expert_map=expert_map,
286 + quant_config=self.quant_config,
287 + apply_router_weight_on_input=apply_router_weight_on_input,
288 + unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2,
289 + unpadded_K_w1=self.moe_config.hidden_dim_unpadded,
290 + unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
291 + unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
292 + )
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
-123
@@ -5,7 +5,6 @@ import torch
5
6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 from vllm import _custom_ops as ops
8 -from vllm._aiter_ops import rocm_aiter_ops
8 from vllm.logger import init_logger
9 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10 from vllm.model_executor.layers.fused_moe.config import (
@@ -286,35 +285,6 @@ def triton_kernel_moe_forward(
285 unpadded_N_w2=None,
286 unpadded_K_w2=None,
287 ) -> torch.Tensor:
289 - if (
290 - quant_config is not None
291 - and quant_config.use_mxfp4_w4a8
292 - and rocm_aiter_ops.is_enabled()
293 - ):
294 - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
295 -
296 - routing_data, gather_idx, scatter_idx = aiter_routing(
297 - gating_output, topk, sm_first=not renormalize
298 - )
299 - return triton_kernel_fused_mxfp4_w4a8_experts(
300 - None,
301 - hidden_states,
302 - w1,
303 - w2,
304 - routing_data,
305 - gather_idx,
306 - scatter_idx,
307 - activation=activation.value,
308 - quant_config=quant_config,
309 - apply_router_weight_on_input=apply_router_weight_on_input,
310 - global_num_experts=global_num_experts,
311 - expert_map=expert_map,
312 - unpadded_N_w1=unpadded_N_w1,
313 - unpadded_K_w1=unpadded_K_w1,
314 - unpadded_N_w2=unpadded_N_w2,
315 - unpadded_K_w2=unpadded_K_w2,
316 - )
317 -
288 from triton_kernels.topk import topk as topk_fn
289
290 sm_first = not renormalize
@@ -471,99 +441,6 @@ def triton_kernel_fused_experts(
441 return output_tensor
442
443
474 -# This is a triton implementation of the fused_experts function
475 -def triton_kernel_fused_mxfp4_w4a8_experts(
476 - output_tensor: torch.Tensor,
477 - hidden_states: torch.Tensor,
478 - w1, # Tensor or triton_kernels.Tensor
479 - w2, # Tensor or triton_kernels.Tensor
480 - routing_data, # RoutingData
481 - gather_indx, # GatherIndx
482 - scatter_indx, # ScatterIndx
483 - activation: str = "silu",
484 - quant_config: FusedMoEQuantConfig | None = None,
485 - swiglu_alpha: float = 1.702,
486 - swiglu_limit: float = 7.0,
487 - apply_router_weight_on_input: bool = False,
488 - global_num_experts: int = -1,
489 - expert_map: torch.Tensor | None = None,
490 - a1q_scale: torch.Tensor | None = None,
491 - unpadded_N_w1=None,
492 - unpadded_K_w1=None,
493 - unpadded_N_w2=None,
494 - unpadded_K_w2=None,
495 -) -> torch.Tensor:
496 - assert quant_config is not None
497 - # type check, uint8 means mxfp4
498 - assert hidden_states.dtype == torch.bfloat16
499 - assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
500 - assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
501 -
502 - # Shape check: weights are padded (e.g. hidden_size padded for
503 - # GFX950 swizzle).
504 - assert hidden_states.shape[-1] == w1.shape[-2]
505 - assert w2.shape[-1] == w1.shape[1]
506 -
507 - E, _, N = w1.shape
508 -
509 - if global_num_experts == -1:
510 - global_num_experts = E
511 -
512 - gammas = routing_data.gate_scal if routing_data else None
513 -
514 - from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
515 - from aiter.ops.triton.quant_moe import downcast_to_static_fp8
516 -
517 - assert quant_config.w1_precision is not None, (
518 - "w1_precision in quant config can't be None"
519 - )
520 - assert quant_config.w2_precision is not None, (
521 - "w2_precision in quant config can't be None"
522 - )
523 -
524 - hidden_states = downcast_to_static_fp8(
525 - hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
526 - )
527 -
528 - intermediate_cache1 = moe_gemm_a8w4(
529 - hidden_states,
530 - w1.storage.data,
531 - None,
532 - quant_config.w1_precision.weight_scale.storage.data,
533 - quant_config.w1_precision.flex_ctx.lhs_data.scale,
534 - quant_config.w2_precision.flex_ctx.lhs_data.scale,
535 - quant_config.w1_bias,
536 - routing_data,
537 - gather_indx=gather_indx,
538 - gammas=gammas if apply_router_weight_on_input else None,
539 - swizzle_mx_scale="CDNA4_SCALE",
540 - out_dtype=torch.float8_e4m3fn,
541 - apply_swiglu=True,
542 - alpha=swiglu_alpha,
543 - limit=swiglu_limit,
544 - unpadded_N=unpadded_N_w1,
545 - unpadded_K=unpadded_K_w1,
546 - )
547 -
548 - intermediate_cache3 = moe_gemm_a8w4(
549 - intermediate_cache1,
550 - w2.storage.data,
551 - None,
552 - quant_config.w2_precision.weight_scale.storage.data,
553 - quant_config.w2_precision.flex_ctx.lhs_data.scale,
554 - None,
555 - quant_config.w2_bias,
556 - routing_data,
557 - scatter_indx=scatter_indx,
558 - gammas=None if apply_router_weight_on_input else gammas,
559 - swizzle_mx_scale="CDNA4_SCALE",
560 - unpadded_N=unpadded_N_w2,
561 - unpadded_K=unpadded_K_w2,
562 - )
563 -
564 - return intermediate_cache3
565 -
566 -
444 def make_routing_data(
445 topk_ids: torch.Tensor,
446 topk_weights: torch.Tensor,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+126 -13
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (
19 FusedMoEQuantConfig,
20 FusedMoEQuantDesc,
21 mxfp4_mxfp8_moe_quant_config,
22 + mxfp4_w4a8_moe_quant_config,
23 mxfp4_w4a16_moe_quant_config,
24 ocp_mx_moe_quant_config,
25 )
@@ -26,9 +27,11 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_m
27 from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 QuantKey,
29 kFp8Dynamic128Sym,
30 + kFp8StaticTensorSym,
31 kMxfp4Static,
32 kMxfp8Dynamic,
33 )
34 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d
35 from vllm.platforms import current_platform
36 from vllm.utils.import_utils import has_triton_kernels
37 from vllm.utils.math_utils import round_up
@@ -59,8 +62,9 @@ class Mxfp4MoeBackend(Enum):
62 # Marlin
63 BATCHED_MARLIN = "BATCHED_MARLIN"
64 MARLIN = "MARLIN"
62 - # ROCm AITER
63 - AITER = "AITER"
65 + # ROCm AITER backends
66 + AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel
67 + AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel
68 # Triton
69 TRITON = "TRITON"
70 TRITON_UNFUSED = "TRITON_UNFUSED"
@@ -72,6 +76,13 @@ class Mxfp4MoeBackend(Enum):
76 HUMMING = "HUMMING"
77
78
79 +# AITER backends group
80 +AITER_BACKENDS = (
81 + Mxfp4MoeBackend.AITER_MXFP4_BF16,
82 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
83 +)
84 +
85 +
86 # Backends that share the same TRTLLM weight format
87 TRTLLM_BACKENDS = (
88 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
@@ -159,13 +170,20 @@ def backend_to_kernel_cls(
170
171 return [BatchedMarlinExperts]
172
162 - elif backend == Mxfp4MoeBackend.AITER:
173 + elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
174 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
175 AiterExperts,
176 )
177
178 return [AiterExperts]
179
180 + elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
181 + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
182 + AiterW4A8ExpertsMonolithic,
183 + )
184 +
185 + return [AiterW4A8ExpertsMonolithic]
186 +
187 elif backend == Mxfp4MoeBackend.XPU:
188 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMXFp4
189
@@ -194,7 +212,8 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend:
212 "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED,
213 "humming": Mxfp4MoeBackend.HUMMING,
214 "marlin": Mxfp4MoeBackend.MARLIN,
197 - "aiter": Mxfp4MoeBackend.AITER,
215 + "aiter": Mxfp4MoeBackend.AITER_MXFP4_BF16,
216 + "aiter_mxfp4_fp8": Mxfp4MoeBackend.AITER_MXFP4_FP8,
217 "xpu": Mxfp4MoeBackend.XPU,
218 "emulation": Mxfp4MoeBackend.EMULATION,
219 }
@@ -213,7 +232,8 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
232 """
233 _AVAILABLE_BACKENDS = [
234 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
216 - Mxfp4MoeBackend.AITER,
235 + Mxfp4MoeBackend.AITER_MXFP4_BF16,
236 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
237 Mxfp4MoeBackend.TRITON,
238 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
239 # TRITON_UNFUSED has bug with MTP support
@@ -254,16 +274,28 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
274 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
275 ):
276 return kMxfp8Dynamic
257 - return None
277 + if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
278 + return kFp8StaticTensorSym
279 + return None # BF16 activation
280
281
260 -def select_gpt_oss_mxfp4_moe_backend(
282 +def select_mxfp4_moe_backend(
283 config: FusedMoEConfig,
284 + activation_key: QuantKey | None = None,
285 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
286 """
287 Select the primary MXFP4 MoE backend.
288 +
289 + Args:
290 + config: MoE configuration
291 + activation_key: Optional activation quantization key. If provided,
292 + overrides the default activation key for backend selection.
293 + Use kFp8StaticTensorSym for W4A8 scheme.
294 +
295 Note: Shape-specific fallbacks may still occur at runtime.
296 """
297 + # If activation_key is explicitly provided (e.g., W4A8), use it
298 + requested_activation_key = activation_key
299 device_capability = current_platform.get_device_capability()
300 triton_kernels_supported = (
301 has_triton_kernels()
@@ -332,11 +364,17 @@ def select_gpt_oss_mxfp4_moe_backend(
364 and requested_backend == Mxfp4MoeBackend.MARLIN
365 ):
366 requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN
367 + # Use requested_activation_key if provided, otherwise use backend default
368 + act_key = (
369 + requested_activation_key
370 + if requested_activation_key is not None
371 + else _backend_activation_key(requested_backend)
372 + )
373 return _return_or_raise(
374 requested_backend,
375 config,
376 kMxfp4Static,
339 - _backend_activation_key(requested_backend),
377 + act_key,
378 activation_format,
379 )
380
@@ -408,10 +446,15 @@ def select_gpt_oss_mxfp4_moe_backend(
446 )
447
448 for backend in AVAILABLE_BACKENDS:
411 - activation_key = _backend_activation_key(backend)
449 + # Use requested_activation_key if provided, otherwise use backend default
450 + act_key = (
451 + requested_activation_key
452 + if requested_activation_key is not None
453 + else _backend_activation_key(backend)
454 + )
455 for k_cls in backend_to_kernel_cls(backend):
456 supported, reason = k_cls.is_supported_config(
414 - k_cls, config, kMxfp4Static, activation_key, activation_format
457 + k_cls, config, kMxfp4Static, act_key, activation_format
458 )
459 if supported:
460 logger.info_once(_make_log_backend(backend))
@@ -438,7 +481,7 @@ def select_gpt_oss_mxfp4_moe_backend(
481 return Mxfp4MoeBackend.NONE, None
482
483
441 -def select_mxfp4_moe_backend(
484 +def select_deepseek_v4_mxfp4_moe_backend(
485 config: FusedMoEConfig,
486 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
487 """
@@ -836,7 +879,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
879 w2_bias,
880 )
881
839 - elif mxfp4_backend == Mxfp4MoeBackend.AITER:
882 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
883 from vllm._aiter_ops import rocm_aiter_ops
884
885 if w13_bias is not None:
@@ -898,6 +941,63 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
941 w2_bias,
942 )
943
944 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
945 + # W4A8: MXFP4 weights + static FP8 activations (triton kernel)
946 + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
947 + from triton_kernels.numerics import InFlexData
948 +
949 + if w13_bias is not None:
950 + w13_bias = w13_bias.to(torch.float32)
951 + if w2_bias is not None:
952 + w2_bias = w2_bias.to(torch.float32)
953 +
954 + # Process static FP8 input scales (reduce to scalar, warn if not uniform)
955 + w13_input_scale = layer.w13_input_scale
956 + w2_input_scale = layer.w2_input_scale
957 + if w13_input_scale is None or w2_input_scale is None:
958 + raise ValueError(
959 + "W4A8 (AITER_MXFP4_FP8) requires static input scales, but found "
960 + "w13_input_scale or w2_input_scale is None."
961 + )
962 + if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
963 + logger.warning_once(
964 + "Found input_scales that are not equal for "
965 + "fp8 MoE layer. Using the maximum across experts "
966 + "for each layer."
967 + )
968 + w13_input_scale = w13_input_scale.max().to(torch.float32)
969 + w2_input_scale = w2_input_scale.max().to(torch.float32)
970 +
971 + # Swizzle weights for GFX950
972 + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(w13_weight, w13_weight_scale)
973 + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(w2_weight, w2_weight_scale)
974 +
975 + # Create InFlexData for activation scales
976 + lhs_data13 = InFlexData(scale=w13_input_scale)
977 + lhs_data2 = InFlexData(scale=w2_input_scale)
978 +
979 + # Create PrecisionConfig with both weight and activation info
980 + w13_precision_config = PrecisionConfig(
981 + weight_scale=w13_scale,
982 + flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
983 + )
984 + w2_precision_config = PrecisionConfig(
985 + weight_scale=w2_scale,
986 + flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
987 + )
988 +
989 + del layer.w13_weight
990 + del layer.w2_weight
991 +
992 + return (
993 + w13_weight,
994 + w2_weight,
995 + w13_precision_config,
996 + w2_precision_config,
997 + w13_bias,
998 + w2_bias,
999 + )
1000 +
1001 elif mxfp4_backend in TRITON_BACKENDS:
1002 from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1003
@@ -1220,6 +1320,8 @@ def make_mxfp4_moe_quant_config(
1320 swiglu_limit: float | None = None,
1321 w1_bias: torch.Tensor | None = None,
1322 w2_bias: torch.Tensor | None = None,
1323 + a1_scale: torch.Tensor | None = None,
1324 + a2_scale: torch.Tensor | None = None,
1325 layer: torch.nn.Module | None = None,
1326 ) -> FusedMoEQuantConfig | None:
1327 """Create a FusedMoEQuantConfig for the given MXFP4 backend."""
@@ -1262,6 +1364,17 @@ def make_mxfp4_moe_quant_config(
1364 gemm1_beta=gemm1_beta,
1365 gemm1_clamp_limit=swiglu_limit,
1366 )
1367 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
1368 + # W4A8: MXFP4 weights + static FP8 activations
1369 + return mxfp4_w4a8_moe_quant_config(
1370 + w1_scale=w1_scale,
1371 + w2_scale=w2_scale,
1372 + a1_scale=a1_scale,
1373 + a2_scale=a2_scale,
1374 + w1_bias=w1_bias,
1375 + w2_bias=w2_bias,
1376 + block_shape=None,
1377 + )
1378 elif mxfp4_backend in (
1379 Mxfp4MoeBackend.MARLIN,
1380 Mxfp4MoeBackend.BATCHED_MARLIN,
@@ -1269,7 +1382,7 @@ def make_mxfp4_moe_quant_config(
1382 Mxfp4MoeBackend.TRITON_UNFUSED,
1383 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
1384 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
1272 - Mxfp4MoeBackend.AITER,
1385 + Mxfp4MoeBackend.AITER_MXFP4_BF16,
1386 ):
1387 return mxfp4_w4a16_moe_quant_config(
1388 w1_bias=w1_bias,
vllm/model_executor/layers/quantization/mxfp4.py
+3 -3
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
24 make_mxfp4_moe_kernel,
25 make_mxfp4_moe_quant_config,
26 mxfp4_round_up_hidden_size_and_intermediate_size,
27 - select_gpt_oss_mxfp4_moe_backend,
27 + select_deepseek_v4_mxfp4_moe_backend,
28 select_mxfp4_moe_backend,
29 )
30 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
@@ -140,7 +140,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
140 def __init__(self, moe: FusedMoEConfig):
141 super().__init__(moe)
142 self.weight_dtype = "gpt_oss_mxfp4"
143 - self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)
143 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
144
145 self.max_capture_size = (
146 get_current_vllm_config().compilation_config.max_cudagraph_capture_size
@@ -468,7 +468,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
468 def __init__(self, moe: FusedMoEConfig):
469 super().__init__(moe)
470 self.weight_dtype = "mxfp4"
471 - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
471 + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)
472
473 self.max_capture_size = (
474 get_current_vllm_config().compilation_config.max_cudagraph_capture_size
vllm/model_executor/layers/quantization/quark/quark_moe.py
+66 -200
@@ -35,19 +35,19 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
35 make_mxfp4_moe_kernel,
36 make_mxfp4_moe_quant_config,
37 mxfp4_round_up_hidden_size_and_intermediate_size,
38 - select_gpt_oss_mxfp4_moe_backend,
38 + select_mxfp4_moe_backend,
39 )
40 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
41 prepare_fp8_moe_layer_for_marlin,
42 )
43 -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
44 - _swizzle_mxfp4,
45 -)
43 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
44 OCP_MX_BLOCK_SIZE,
45 OCP_MX_Scheme,
46 )
50 -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
47 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
48 + GroupShape,
49 + kFp8StaticTensorSym,
50 +)
51 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
52 all_close_1d,
53 normalize_e4m3fn_to_e4m3fnuz,
@@ -62,7 +62,6 @@ logger = init_logger(__name__)
62 __all__ = [
63 "QuarkMoEMethod",
64 "QuarkOCP_MX_MoEMethod",
65 - "QuarkOCP_MX_MoEMethod_OSS",
65 ]
66
67
@@ -94,22 +93,9 @@ class QuarkMoEMethod(FusedMoEMethodBase):
93 elif quant_config._is_fp8_w8a8(weight_config, input_config):
94 return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
95 elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config):
97 - emulate = not current_platform.supports_mx() or not (
98 - rocm_aiter_ops.is_fused_moe_enabled()
99 - )
100 - if (
101 - input_config is not None
102 - and input_config.get("dtype") == "fp8_e4m3"
103 - and not input_config.get("is_dynamic")
104 - and not emulate
105 - ):
106 - return QuarkOCP_MX_MoEMethod_OSS(
107 - weight_config, input_config, module.moe_config
108 - )
109 - else:
110 - return QuarkOCP_MX_MoEMethod(
111 - weight_config, input_config, module.moe_config
112 - )
96 + # All OCP MX schemes (W4A16, W4A8, etc.) handled by QuarkOCP_MX_MoEMethod
97 + # Backend selection happens inside via oracle
98 + return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
99 elif quant_config._is_static_tensor_w8a8(
100 weight_config, input_config
101 ) or quant_config._is_dynamic_per_token_w8a8(weight_config, input_config):
@@ -993,7 +979,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
979 self.experts_cls: type[mk.FusedMoEExperts] | None = None
980 self.moe_kernel: mk.FusedMoEKernel | None = None
981
996 - # Used for triton kernel precision configs
982 + # Used for triton kernel precision configs (W4A8, TRITON backends)
983 self.w13_precision_config = None
984 self.w2_precision_config = None
985
@@ -1002,6 +988,17 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
988 else:
989 self.static_input_scales = False
990
991 + # Select backend based on OCP MX scheme
992 + if self.ocp_mx_scheme == "w_mxfp4":
993 + # W4A16: weight-only MXFP4
994 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
995 + elif self.ocp_mx_scheme == "w_mxfp4_a_fp8" and self.static_input_scales:
996 + # W4A8: MXFP4 weights + static FP8 activations
997 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(
998 + moe, activation_key=kFp8StaticTensorSym
999 + )
1000 +
1001 + # Validation for unsupported schemes
1002 if any(
1003 self.ocp_mx_scheme.endswith(a_scheme)
1004 for a_scheme in ["a_mxfp4", "a_mxfp6_e3m2", "a_mxfp6_e2m3"]
@@ -1026,7 +1023,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1023 )
1024
1025 # TODO: Remove once all OCP MX schemes use the kernel abstraction
1029 - _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4")
1026 + _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4", "w_mxfp4_a_fp8")
1027 self.emulate = (
1028 not current_platform.supports_mx()
1029 or self.ocp_mx_scheme not in _AITER_NATIVE_OCP_MX_SCHEMES
@@ -1034,9 +1031,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1031 self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe
1032 )
1033
1037 - if self.ocp_mx_scheme == "w_mxfp4":
1038 - self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe)
1039 -
1034 if self.emulate:
1035 # We use the same code path between MXFP4/MXFP6 emulation.
1036 self.mxfp4_backend = Mxfp4MoeBackend.EMULATION
@@ -1046,7 +1040,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1040 if self.mxfp4_backend != Mxfp4MoeBackend.NONE:
1041 self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
1042
1049 - if self.emulate:
1043 + # Log backend selection
1044 + if self.mxfp4_backend != Mxfp4MoeBackend.NONE:
1045 + logger.info_once(
1046 + f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}"
1047 + )
1048 + elif self.emulate:
1049 logger.warning_once(
1050 f"The current mode (supports_mx={current_platform.supports_mx()}, "
1051 f"use_rocm_aiter_moe={self.use_rocm_aiter_moe}, "
@@ -1056,10 +1055,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1055 "QDQ (quantize and dequantize) will be used, with the linear "
1056 "layers computed in high precision."
1057 )
1059 - else:
1060 - logger.warning_once(
1061 - "The current mode supports native MoE MXFP4 computation"
1062 - )
1058
1059 def maybe_roundup_sizes(
1060 self,
@@ -1204,6 +1199,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1199 layer.w2_input_scale = None
1200
1201 def process_weights_after_loading(self, layer):
1202 + # For MXFP4 schemes with native backend, use oracle
1203 + if self.mxfp4_backend != Mxfp4MoeBackend.NONE:
1204 + self._setup_kernel(layer)
1205 + return
1206 +
1207 if self.static_input_scales and self.input_dtype == "fp8":
1208 # firstly, process activations if fp8 static input
1209 if layer.w13_input_scale is None or layer.w2_input_scale is None:
@@ -1252,14 +1252,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1252 w2_input_scale, requires_grad=False
1253 )
1254
1255 - # For w_mxfp4, use oracle functions
1256 - if self.emulate or (
1257 - self.ocp_mx_scheme == "w_mxfp4"
1258 - and self.mxfp4_backend != Mxfp4MoeBackend.NONE
1259 - ):
1260 - self._setup_kernel_via_oracle(layer)
1261 - return
1262 -
1255 # TODO(bowenbao): gradually migrate to oracles.
1256 # Existing AITER path for w_mxfp4_a_mxfp4 and other schemes
1257 from aiter.utility.fp4_utils import e8m0_shuffle
@@ -1298,46 +1290,48 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1290 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1291 torch.accelerator.empty_cache()
1292
1301 - def _setup_kernel_via_oracle(self, layer: FusedMoE):
1302 - """Setup kernel using oracle functions for w_mxfp4 scheme."""
1303 - w13 = layer.w13_weight
1304 - w2 = layer.w2_weight
1305 - w13_scale = layer.w13_weight_scale
1306 - w2_scale = layer.w2_weight_scale
1293 + def _setup_kernel(self, layer: FusedMoE):
1294 + """Setup kernel using oracle functions for MXFP4 schemes (W4A16, W4A8)."""
1295 w13_bias = getattr(layer, "w13_bias", None)
1296 w2_bias = getattr(layer, "w2_bias", None)
1297
1310 - # Convert weights to kernel format
1298 + # Convert weights to kernel format (handles all backend-specific logic)
1299 w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
1300 convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1301 mxfp4_backend=self.mxfp4_backend,
1302 layer=layer,
1315 - w13_weight=w13,
1316 - w2_weight=w2,
1317 - w13_weight_scale=w13_scale,
1318 - w2_weight_scale=w2_scale,
1303 + w13_weight=layer.w13_weight,
1304 + w2_weight=layer.w2_weight,
1305 + w13_weight_scale=layer.w13_weight_scale,
1306 + w2_weight_scale=layer.w2_weight_scale,
1307 w13_bias=w13_bias,
1308 w2_bias=w2_bias,
1309 )
1310 )
1311
1324 - # For TRITON backends, weights are wrapped tensors from triton_kernels
1325 - # that don't support .detach(). Manually assign parameters.
1326 - if self.mxfp4_backend not in TRITON_BACKENDS:
1327 - replace_parameter(layer, "w13_weight", w13)
1328 - replace_parameter(layer, "w2_weight", w2)
1329 - replace_parameter(layer, "w13_weight_scale", w13_scale)
1330 - replace_parameter(layer, "w2_weight_scale", w2_scale)
1331 - else:
1312 + # Handle weight/scale assignment based on backend type
1313 + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in (
1314 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
1315 + ):
1316 + # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor
1317 + # Store on layer for apply(), scales are PrecisionConfig
1318 layer.w13_weight = w13
1319 layer.w2_weight = w2
1320 self.w13_precision_config = w13_scale
1321 self.w2_precision_config = w2_scale
1322 + else:
1323 + # Standard backends: replace parameters
1324 + replace_parameter(layer, "w13_weight", w13)
1325 + replace_parameter(layer, "w2_weight", w2)
1326 + replace_parameter(layer, "w13_weight_scale", w13_scale)
1327 + replace_parameter(layer, "w2_weight_scale", w2_scale)
1328
1329 if w13_bias is not None and w2_bias is not None:
1330 replace_parameter(layer, "w13_bias", w13_bias)
1331 replace_parameter(layer, "w2_bias", w2_bias)
1332
1333 + torch.accelerator.empty_cache()
1334 +
1335 # Build quant config and kernel
1336 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1337 if self.moe_quant_config is not None and self.experts_cls is not None:
@@ -1353,22 +1347,26 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1347 def get_fused_moe_quant_config(
1348 self, layer: torch.nn.Module
1349 ) -> FusedMoEQuantConfig | None:
1356 - # For w_mxfp4 with oracle backend, use oracle function
1357 - if self.ocp_mx_scheme == "w_mxfp4" and self.mxfp4_backend not in (
1358 - Mxfp4MoeBackend.NONE,
1359 - Mxfp4MoeBackend.EMULATION,
1360 - ):
1361 - w1_scale = layer.w13_weight_scale
1362 - w2_scale = layer.w2_weight_scale
1363 - if self.mxfp4_backend in TRITON_BACKENDS:
1350 + # For oracle-based backends (W4A16, W4A8), use make_mxfp4_moe_quant_config
1351 + if self.mxfp4_backend not in (Mxfp4MoeBackend.NONE, Mxfp4MoeBackend.EMULATION):
1352 + # Determine scale source based on backend type
1353 + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in (
1354 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
1355 + ):
1356 w1_scale = self.w13_precision_config
1357 w2_scale = self.w2_precision_config
1358 + else:
1359 + w1_scale = layer.w13_weight_scale
1360 + w2_scale = layer.w2_weight_scale
1361 +
1362 return make_mxfp4_moe_quant_config(
1363 mxfp4_backend=self.mxfp4_backend,
1364 w1_scale=w1_scale,
1365 w2_scale=w2_scale,
1366 w1_bias=getattr(layer, "w13_bias", None),
1367 w2_bias=getattr(layer, "w2_bias", None),
1368 + a1_scale=getattr(layer, "w13_input_scale", None),
1369 + a2_scale=getattr(layer, "w2_input_scale", None),
1370 )
1371
1372 # Emulation and other schemes
@@ -1421,7 +1419,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1419 topk_ids: torch.Tensor,
1420 shared_experts_input: torch.Tensor | None,
1421 ) -> torch.Tensor:
1424 - # For oracle kernel or emulation kernel
1422 + # For oracle-based kernels (W4A16, W4A8) or emulation kernel
1423 if self.moe_kernel is not None:
1424 return self.moe_kernel.apply(
1425 hidden_states=x,
@@ -1473,135 +1471,3 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1471 expert_map=layer.expert_map,
1472 apply_router_weight_on_input=layer.apply_router_weight_on_input,
1473 )
1476 -
1477 -
1478 -class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
1479 - def __init__(
1480 - self,
1481 - weight_config: dict[str, Any],
1482 - input_config: dict[str, Any],
1483 - moe: FusedMoEConfig,
1484 - ):
1485 - super().__init__(weight_config, input_config, moe)
1486 -
1487 - def process_weights_after_loading(self, layer):
1488 - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1489 -
1490 - w13_bias = layer.w13_bias.to(torch.float32)
1491 - w2_bias = layer.w2_bias.to(torch.float32)
1492 -
1493 - layer.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
1494 - layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
1495 -
1496 - # FIXME warp need to be adjusted based on batch size
1497 - # only apply to batched mode
1498 - if self.moe.use_ep:
1499 - num_warps = 4 if self.moe.max_num_tokens <= 512 else 8
1500 - else:
1501 - num_warps = 8
1502 -
1503 - w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
1504 - layer.w13_weight, layer.w13_weight_scale, num_warps
1505 - )
1506 - w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
1507 - layer.w2_weight, layer.w2_weight_scale, num_warps
1508 - )
1509 -
1510 - self.w13_weight_triton_tensor = w13_weight
1511 - self.w2_weight_triton_tensor = w2_weight
1512 -
1513 - # need to delete the original weights to save memory on single GPU
1514 - del layer.w13_weight
1515 - del layer.w2_weight
1516 - layer.w13_weight = None
1517 - layer.w2_weight = None
1518 - torch.accelerator.empty_cache()
1519 -
1520 - if self.static_input_scales:
1521 - if layer.w13_input_scale is None or layer.w2_input_scale is None:
1522 - raise ValueError(
1523 - "QuantConfig has static quantization, but found "
1524 - "activation scales are None."
1525 - )
1526 - if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
1527 - layer.w2_input_scale
1528 - ):
1529 - logger.warning_once(
1530 - "Found input_scales that are not equal for "
1531 - "fp8 MoE layer. Using the maximum across experts "
1532 - "for each layer."
1533 - )
1534 -
1535 - layer.w13_input_scale = torch.nn.Parameter(
1536 - layer.w13_input_scale.max().to(torch.float32), requires_grad=False
1537 - )
1538 - layer.w2_input_scale = torch.nn.Parameter(
1539 - layer.w2_input_scale.max().to(torch.float32), requires_grad=False
1540 - )
1541 -
1542 - from triton_kernels.numerics import InFlexData
1543 -
1544 - lhs_data13 = InFlexData(scale=layer.w13_input_scale)
1545 - lhs_data2 = InFlexData(scale=layer.w2_input_scale)
1546 -
1547 - self.w13_precision_config = PrecisionConfig(
1548 - weight_scale=w13_scale,
1549 - flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
1550 - )
1551 -
1552 - self.w2_precision_config = PrecisionConfig(
1553 - weight_scale=w2_scale,
1554 - flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
1555 - )
1556 -
1557 - def get_fused_moe_quant_config(
1558 - self, layer: torch.nn.Module
1559 - ) -> FusedMoEQuantConfig | None:
1560 - return mxfp4_w4a8_moe_quant_config(
1561 - w1_scale=self.w13_precision_config,
1562 - w2_scale=self.w2_precision_config,
1563 - a1_scale=layer.w13_input_scale,
1564 - a2_scale=layer.w2_input_scale,
1565 - w1_bias=layer.w13_bias,
1566 - w2_bias=layer.w2_bias,
1567 - block_shape=None,
1568 - )
1569 -
1570 - @property
1571 - def is_monolithic(self) -> bool:
1572 - return True
1573 -
1574 - def apply_monolithic(
1575 - self,
1576 - layer: FusedMoE,
1577 - x: torch.Tensor,
1578 - router_logits: torch.Tensor,
1579 - input_ids: torch.Tensor | None = None,
1580 - ) -> torch.Tensor:
1581 - if layer.enable_eplb:
1582 - raise NotImplementedError(
1583 - f"EPLB not supported for {self.__class__.__name__} yet."
1584 - )
1585 -
1586 - from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
1587 - triton_kernel_moe_forward,
1588 - )
1589 -
1590 - assert self.moe.hidden_dim_unpadded is not None
1591 - assert self.moe.intermediate_size_per_partition_unpadded is not None
1592 - return triton_kernel_moe_forward(
1593 - hidden_states=x,
1594 - w1=self.w13_weight_triton_tensor,
1595 - w2=self.w2_weight_triton_tensor,
1596 - gating_output=router_logits,
1597 - topk=layer.top_k,
1598 - renormalize=layer.renormalize,
1599 - global_num_experts=layer.global_num_experts,
1600 - expert_map=layer.expert_map,
1601 - quant_config=self.moe_quant_config,
1602 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
1603 - unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2,
1604 - unpadded_K_w1=self.moe.hidden_dim_unpadded,
1605 - unpadded_N_w2=self.moe.hidden_dim_unpadded,
1606 - unpadded_K_w2=self.moe.intermediate_size_per_partition_unpadded,
1607 - )