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 7, 14:00 UTC
[ROCm][Quantization] Add Composable Kernel (CK) backend support for M… (#34301)
Signed-off-by: Doug Lehr <douglehr@amd.com> Signed-off-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Signed-off-by: Douglas Lehr <Doug.Lehr@amd.com> Co-authored-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Douglas Lehr committed
Feb 27, 2026 at 21:37 UTC
d5b6f3ba362552d582f93c25958a6a612d28a700
2 files changed
+260
-6
vllm/_aiter_ops.py
+160
@@ -87,6 +87,10 @@ def _rocm_aiter_fused_moe_impl(
87
a2_scale: torch.Tensor | None = None,
88
num_local_tokens: torch.Tensor | None = None,
89
output_dtype: torch.dtype | None = None,
90
+ hidden_pad: int = 0,
91
+ intermediate_pad: int = 0,
92
+ bias1: torch.Tensor | None = None,
93
+ bias2: torch.Tensor | None = None,
94
) -> torch.Tensor:
95
from aiter import ActivationType, QuantType
96
from aiter.fused_moe import fused_moe
@@ -110,6 +114,10 @@ def _rocm_aiter_fused_moe_impl(
114
a2_scale,
115
num_local_tokens=num_local_tokens,
116
dtype=output_dtype,
117
+ hidden_pad=hidden_pad,
118
+ intermediate_pad=intermediate_pad,
119
+ bias1=bias1,
120
+ bias2=bias2,
121
)
122
123
@@ -307,6 +315,28 @@ def _rocm_aiter_grouped_topk_fake(
315
pass
316
317
318
+def _rocm_aiter_fused_topk_impl(
319
+ x: torch.Tensor,
320
+ router_logits: torch.Tensor,
321
+ top_k: int,
322
+ gate_up: bool,
323
+) -> tuple[torch.Tensor, torch.Tensor]:
324
+ from aiter.fused_moe import fused_topk
325
+
326
+ # fused_topk returns (topk_weights, topk_indices)
327
+ return fused_topk(x, router_logits, top_k, gate_up)
328
+
329
+
330
+def _rocm_aiter_fused_topk_fake(
331
+ x: torch.Tensor,
332
+ router_logits: torch.Tensor,
333
+ top_k: int,
334
+ gate_up: bool,
335
+) -> None:
336
+ # tuple[torch.Tensor, torch.Tensor]:
337
+ pass
338
+
339
+
340
# Cache whether aiter supports FP8 MLA parameters
341
_AITER_MLA_SUPPORTS_FP8: bool | None = None
342
@@ -994,6 +1024,70 @@ class rocm_aiter_ops:
1024
cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
1025
cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
1026
1027
+ @staticmethod
1028
+ def get_aiter_activation_type(activation_str: str):
1029
+ """
1030
+ Given an activation type as a string, returns the corresponding aiter ActivationType enum.
1031
+ Supported activation types: "no", "none", "silu", "gelu", "swiglu".
1032
+ Returns None if the mapping fails.
1033
+
1034
+ Args:
1035
+ activation_str (str): Activation type as string.
1036
+
1037
+ Returns:
1038
+ Aiter ActivationType enum value, or None if not found.
1039
+ """
1040
+ # Import only locally, since aiter may not always be available.
1041
+ try:
1042
+ from aiter import ActivationType
1043
+ except ImportError:
1044
+ return None
1045
+
1046
+ if not isinstance(activation_str, str):
1047
+ return None
1048
+
1049
+ name = activation_str.strip().lower()
1050
+ mapping = {
1051
+ "none": ActivationType.No,
1052
+ "no": ActivationType.No,
1053
+ "silu": ActivationType.Silu,
1054
+ "gelu": ActivationType.Gelu,
1055
+ "swiglu": ActivationType.Swiglu,
1056
+ }
1057
+ return mapping.get(name)
1058
+
1059
+ @staticmethod
1060
+ def get_aiter_quant_type(quant_type_str: str):
1061
+ """
1062
+ Given a quantization type as a string, returns the corresponding aiter QuantType enum.
1063
+ Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
1064
+ Returns None if the mapping fails.
1065
+
1066
+ Args:
1067
+ quant_type_str (str): Quantization type as string.
1068
+
1069
+ Returns:
1070
+ Aiter QuantType enum value, or None if not found.
1071
+ """
1072
+ try:
1073
+ from aiter import QuantType
1074
+ except ImportError:
1075
+ return None
1076
+
1077
+ if not isinstance(quant_type_str, str):
1078
+ return None
1079
+
1080
+ name = quant_type_str.strip().lower()
1081
+ mapping = {
1082
+ "no": QuantType.No,
1083
+ "per_tensor": QuantType.per_Tensor,
1084
+ "per_token": QuantType.per_Token,
1085
+ "per_1x32": QuantType.per_1x32,
1086
+ "per_1x128": QuantType.per_1x128,
1087
+ "per_128x128": QuantType.per_128x128,
1088
+ }
1089
+ return mapping.get(name)
1090
+
1091
@classmethod
1092
@if_aiter_supported
1093
def is_enabled(cls) -> bool:
@@ -1127,6 +1221,14 @@ class rocm_aiter_ops:
1221
dispatch_key=current_platform.dispatch_key,
1222
)
1223
1224
+ direct_register_custom_op(
1225
+ op_name="rocm_aiter_fused_topk",
1226
+ op_func=_rocm_aiter_fused_topk_impl,
1227
+ mutates_args=[],
1228
+ fake_impl=_rocm_aiter_fused_topk_fake,
1229
+ dispatch_key=current_platform.dispatch_key,
1230
+ )
1231
+
1232
direct_register_custom_op(
1233
op_name="rocm_aiter_mla_decode_fwd",
1234
op_func=_rocm_aiter_mla_decode_fwd_impl,
@@ -1360,6 +1462,10 @@ class rocm_aiter_ops:
1462
a2_scale: torch.Tensor | None = None,
1463
num_local_tokens: torch.Tensor | None = None,
1464
output_dtype: torch.dtype | None = None,
1465
+ hidden_pad: int = 0,
1466
+ intermediate_pad: int = 0,
1467
+ bias1: torch.Tensor | None = None,
1468
+ bias2: torch.Tensor | None = None,
1469
) -> torch.Tensor:
1470
return torch.ops.vllm.rocm_aiter_fused_moe(
1471
hidden_states,
@@ -1377,6 +1483,10 @@ class rocm_aiter_ops:
1483
a2_scale,
1484
num_local_tokens,
1485
output_dtype,
1486
+ hidden_pad,
1487
+ intermediate_pad,
1488
+ bias1,
1489
+ bias2,
1490
)
1491
1492
@staticmethod
@@ -1481,6 +1591,15 @@ class rocm_aiter_ops:
1591
routed_scaling_factor,
1592
)
1593
1594
+ @staticmethod
1595
+ def fused_topk(
1596
+ x: torch.Tensor,
1597
+ router_logits: torch.Tensor,
1598
+ top_k: int,
1599
+ gate_up: bool,
1600
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1601
+ return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
1602
+
1603
@staticmethod
1604
def mla_decode_fwd(
1605
q: torch.Tensor,
@@ -1701,6 +1820,47 @@ class rocm_aiter_ops:
1820
1821
return shuffle_weight(tensor, layout=layout)
1822
1823
+ @staticmethod
1824
+ def shuffle_weight_a16w4(
1825
+ tensor: "torch.Tensor",
1826
+ nLane: int,
1827
+ gate_up: bool,
1828
+ ) -> "torch.Tensor":
1829
+ """
1830
+ Shuffles the weight tensor into (A16W4) layout for AITER kernels.
1831
+
1832
+ Args:
1833
+ tensor: The input weight tensor to be shuffled.
1834
+ layout: The block layout to use, defaults to (16, 4).
1835
+
1836
+ Returns:
1837
+ torch.Tensor: The shuffled tensor.
1838
+ """
1839
+ from aiter.ops.shuffle import shuffle_weight_a16w4
1840
+
1841
+ return shuffle_weight_a16w4(tensor, nLane, gate_up)
1842
+
1843
+ @staticmethod
1844
+ def shuffle_scale_a16w4(
1845
+ tensor: "torch.Tensor",
1846
+ num_experts: int,
1847
+ gate_up: bool,
1848
+ ) -> "torch.Tensor":
1849
+ """
1850
+ Shuffles the scale tensor into (A16W4) layout for AITER kernels.
1851
+
1852
+ Args:
1853
+ tensor: The input scale tensor to be shuffled.
1854
+ num_experts: Number of experts, needed for reshaping logic.
1855
+ gate_up: Whether the scale is for w13 (True) or w2 (False).
1856
+
1857
+ Returns:
1858
+ torch.Tensor: The shuffled scale tensor.
1859
+ """
1860
+ from aiter.ops.shuffle import shuffle_scale_a16w4
1861
+
1862
+ return shuffle_scale_a16w4(tensor, num_experts, gate_up)
1863
+
1864
@staticmethod
1865
def shuffle_weights(
1866
*tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
vllm/model_executor/layers/quantization/mxfp4.py
+100
-6
@@ -6,6 +6,7 @@ import torch
6
from torch.nn.parameter import Parameter
7
8
from vllm import envs
9
+from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.config import get_current_vllm_config
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.attention import Attention
@@ -77,6 +78,8 @@ class Mxfp4Backend(Enum):
78
# Triton Backend
79
TRITON = 6
80
81
+ CK = 7
82
+
83
84
def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
85
"""
@@ -167,9 +170,15 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
170
elif current_platform.is_xpu():
171
logger.info_once("Using xpu backend on XPU")
172
return Mxfp4Backend.MARLIN
170
- elif current_platform.is_rocm() and has_triton_kernels():
171
- logger.info_once("Using Triton backend")
172
- return Mxfp4Backend.TRITON
173
+ elif current_platform.is_rocm():
174
+ from vllm.platforms.rocm import on_gfx950
175
+
176
+ if rocm_aiter_ops.is_enabled() and on_gfx950():
177
+ logger.info_once("Using CK MXFP4 MoE backend (Aiter ROCm)")
178
+ return Mxfp4Backend.CK
179
+ elif has_triton_kernels():
180
+ logger.info_once("Using Triton backend")
181
+ return Mxfp4Backend.TRITON
182
183
return Mxfp4Backend.NONE
184
@@ -338,6 +347,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
347
348
self.intermediate_size = intermediate_size_per_partition_after_pad
349
self.hidden_size = hidden_size
350
+ self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)
351
+ self.intermediate_pad = (
352
+ intermediate_size_per_partition_after_pad - intermediate_size_per_partition
353
+ )
354
# Fused gate_up_proj (column parallel)
355
w13_weight = torch.nn.Parameter(
356
torch.zeros(
@@ -784,6 +797,66 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
797
),
798
shared_experts=None,
799
)
800
+ elif self.mxfp4_backend == Mxfp4Backend.CK:
801
+ if layer.w13_bias is not None:
802
+ layer.w13_bias.data = layer.w13_bias.data.to(torch.float32)
803
+ if layer.w2_bias.data is not None:
804
+ layer.w2_bias.data = layer.w2_bias.data.to(torch.float32)
805
+
806
+ e, n, k = layer.w13_weight.shape
807
+ layer.w13_weight.view(torch.uint8).copy_(
808
+ layer.w13_weight.data.view(torch.uint8)
809
+ .view(e, n // 2, 2, k)
810
+ .permute(0, 2, 1, 3)
811
+ .contiguous()
812
+ .view(e, n, k)
813
+ )
814
+ layer.w13_weight_scale.data = (
815
+ layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
816
+ .permute(0, 2, 1, 3)
817
+ .contiguous()
818
+ .view(e, n, -1)
819
+ )
820
+ layer.w13_weight.data = layer.w13_weight.data.view(torch.float4_e2m1fn_x2)
821
+ layer.w2_weight.data = layer.w2_weight.data.view(torch.float4_e2m1fn_x2)
822
+
823
+ layer.w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(
824
+ layer.w13_weight, 16, True
825
+ )
826
+ shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
827
+ layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
828
+ self.num_experts,
829
+ True,
830
+ )
831
+
832
+ layer.w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(
833
+ layer.w2_weight, 16, False
834
+ )
835
+ shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
836
+ layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
837
+ self.num_experts,
838
+ False,
839
+ )
840
+
841
+ layer.w13_bias.data = (
842
+ layer.w13_bias.data.view(-1, n // 2, 2)
843
+ .permute(0, 2, 1)
844
+ .contiguous()
845
+ .view(-1, n)
846
+ )
847
+
848
+ layer.w13_weight_scale = torch.nn.Parameter(
849
+ shuffled_w13_scale, requires_grad=False
850
+ )
851
+ layer.w2_weight_scale = torch.nn.Parameter(
852
+ shuffled_w2_scale, requires_grad=False
853
+ )
854
+ # replace_parameter(layer, "w13_bias", w13_bias)
855
+ # replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
856
+ # replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
857
+ # replace_parameter(layer, "w13_weight", w13_weight)
858
+ # replace_parameter(layer, "w2_weight", w2_weight)
859
+
860
elif self.mxfp4_backend == Mxfp4Backend.TRITON:
861
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
862
@@ -792,7 +865,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
865
866
layer.w13_bias = Parameter(w13_bias, requires_grad=False)
867
layer.w2_bias = Parameter(w2_bias, requires_grad=False)
795
-
868
# Ideally we'd use FusedMoEModularKernel.prepare_finalize object
869
# (stored in self.fused_experts) to determine if the MoE has a
870
# batched activation format. As self.fused_experts is not
@@ -803,7 +875,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
875
num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
876
else:
877
num_warps = 8
806
-
878
w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
879
layer.w13_weight, layer.w13_weight_scale, num_warps
880
)
@@ -817,13 +888,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
888
self.w2_precision_config = PrecisionConfig(
889
weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
890
)
820
-
891
self.w13_weight = w13_weight
892
self.w2_weight = w2_weight
893
del layer.w13_weight
894
del layer.w2_weight
895
layer.w13_weight = w13_weight
896
layer.w2_weight = w2_weight
897
+
898
else:
899
raise ValueError(
900
f"Unsupported mxfp4_backend: {self.mxfp4_backend}: "
@@ -862,6 +933,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
933
elif self.mxfp4_backend in [
934
Mxfp4Backend.SM100_FI_MXFP4_BF16,
935
Mxfp4Backend.SM90_FI_MXFP4_BF16,
936
+ Mxfp4Backend.CK,
937
]:
938
return mxfp4_w4a16_moe_quant_config(
939
w1_bias=layer.w13_bias,
@@ -933,6 +1005,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1005
self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
1006
or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
1007
or self.mxfp4_backend == Mxfp4Backend.TRITON
1008
+ or self.mxfp4_backend == Mxfp4Backend.CK
1009
)
1010
1011
def apply(
@@ -1054,6 +1127,27 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1127
tune_max_num_tokens=max(self.max_capture_size, 1),
1128
)[0]
1129
return trtllm_gen_output
1130
+ elif self.mxfp4_backend == Mxfp4Backend.CK:
1131
+ topk_weights, topk_ids = rocm_aiter_ops.fused_topk(
1132
+ x, router_logits, layer.top_k, True
1133
+ )
1134
+ output = rocm_aiter_ops.fused_moe(
1135
+ x,
1136
+ layer.w13_weight,
1137
+ layer.w2_weight,
1138
+ topk_weights,
1139
+ topk_ids,
1140
+ activation_method=rocm_aiter_ops.get_aiter_activation_type("swiglu"),
1141
+ quant_method=rocm_aiter_ops.get_aiter_quant_type("per_1x32"),
1142
+ w1_scale=layer.w13_weight_scale,
1143
+ w2_scale=layer.w2_weight_scale,
1144
+ doweight_stage1=False,
1145
+ hidden_pad=self.hidden_pad // 128 * 128,
1146
+ intermediate_pad=self.intermediate_pad // 64 * 64 * 2,
1147
+ bias1=layer.w13_bias,
1148
+ bias2=layer.w2_bias,
1149
+ )
1150
+ return output
1151
elif self.mxfp4_backend == Mxfp4Backend.TRITON:
1152
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
1153
triton_kernel_moe_forward,