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 9, 07:30 UTC
[Hardware][SM100] Add TRTLLM Kernel for INT4 W4A16 Kernel. (#32437)
Signed-off-by: Pavani Majety <pmajety@nvidia.com>
Pavani Majety committed
Jan 30, 2026 at 10:30 UTC
c3a9752b0c11f87677e2ab918e524af7a368c664
6 files changed
+727
-23
tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
new
+272
@@ -0,0 +1,272 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Test comparing Marlin INT4 MoE vs FlashInfer TRT-LLM MXINT4 MoE."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
9
+ fused_marlin_moe,
10
+)
11
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
12
+ grouped_topk,
13
+)
14
+from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
15
+ prepare_static_weights_for_trtllm_mxint4_moe,
16
+)
17
+from vllm.platforms import current_platform
18
+from vllm.scalar_type import scalar_types
19
+
20
+
21
+def mxint4_quantize(
22
+ x: torch.Tensor, sf_vec_size: int = 32
23
+) -> tuple[torch.Tensor, torch.Tensor]:
24
+ """Quantize BF16 tensor to MXINT4 with block scaling (group_size=sf_vec_size).
25
+
26
+ Returns:
27
+ - uint8 packed (2 INT4/byte): [..., k//2] - stores SIGNED INT4 [-8, 7]
28
+ - scales in BF16: [..., k//sf_vec_size]
29
+ """
30
+ x_reshaped = x.reshape(-1, sf_vec_size)
31
+ x_max = x_reshaped.max(dim=-1, keepdim=True)[0].to(torch.float32)
32
+ x_min = x_reshaped.min(dim=-1, keepdim=True)[0].to(torch.float32)
33
+ x_max = x_max * 8.0 / 7.0
34
+ amax = torch.where(x_max > -x_min, x_max, -x_min)
35
+ scales = amax / 8.0
36
+ x_scaled = x_reshaped * scales.reciprocal()
37
+ x_int8 = (
38
+ x_scaled.round().clamp(-8, 7).to(torch.int8).reshape(-1, sf_vec_size // 2, 2)
39
+ )
40
+ x_int4 = (x_int8[..., 0] & 0x0F) | ((x_int8[..., 1] & 0x0F) << 4)
41
+ return (
42
+ x_int4.to(torch.uint8).reshape(*x.shape[:-1], x.shape[-1] // 2),
43
+ scales.to(x.dtype).reshape(*x.shape[:-1], x.shape[-1] // sf_vec_size),
44
+ )
45
+
46
+
47
+def mxint4_quantize_moe_weights(
48
+ weights_bf16: torch.Tensor, group_size: int = 32
49
+) -> tuple[torch.Tensor, torch.Tensor]:
50
+ """Quantize MoE weights [e, n, k] to MxInt4 format.
51
+
52
+ Args:
53
+ weights_bf16: BF16 weights of shape [num_experts, out_features, in_features]
54
+ group_size: Quantization group size (default: 32)
55
+
56
+ Returns:
57
+ - weights_mxint4: Quantized weights [e, n, k//2] uint8
58
+ - scales_mxint4: Quantization scales [e, n, k//group_size] bf16
59
+ """
60
+ e = weights_bf16.shape[0]
61
+ weight_list = []
62
+ scale_list = []
63
+
64
+ for i in range(e):
65
+ w_q, w_s = mxint4_quantize(weights_bf16[i], sf_vec_size=group_size)
66
+ weight_list.append(w_q)
67
+ scale_list.append(w_s)
68
+
69
+ return torch.stack(weight_list), torch.stack(scale_list)
70
+
71
+
72
+__all__ = [
73
+ "mxint4_quantize",
74
+ "mxint4_quantize_moe_weights",
75
+ "marlin_quantize_moe_weights",
76
+]
77
+
78
+
79
+def marlin_quantize_moe_weights(
80
+ weights_bf16: torch.Tensor, group_size: int = 32
81
+) -> tuple[torch.Tensor, torch.Tensor]:
82
+ """Quantize MoE weights [e, n, k] to Marlin INT4 format.
83
+
84
+ Args:
85
+ weights_bf16: BF16 weights of shape [num_experts, out_features, in_features]
86
+ group_size: Quantization group size (default: 32)
87
+
88
+ Returns:
89
+ - weights_marlin: Marlin quantized weights [e, k//8, n] int32
90
+ - scales_marlin: Marlin quantization scales [e, k//group_size, n] bf16
91
+ """
92
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
93
+ marlin_quantize,
94
+ )
95
+
96
+ e, n, k = weights_bf16.shape
97
+ weight_list = []
98
+ scale_list = []
99
+
100
+ for i in range(e):
101
+ # Transpose for Marlin: [n, k] → [k, n]
102
+ w_t = weights_bf16[i].T.contiguous()
103
+ _, w_q, w_s, _, _, _ = marlin_quantize(
104
+ w_t, scalar_types.uint4b8, group_size, act_order=False
105
+ )
106
+ weight_list.append(w_q)
107
+ scale_list.append(w_s)
108
+
109
+ # Stack to get [e, ...] shape
110
+ weights_marlin = torch.stack(weight_list) # [e, k // 8, n]
111
+ scales_marlin = torch.stack(scale_list) # [e, k // group_size, n]
112
+
113
+ return weights_marlin, scales_marlin
114
+
115
+
116
+TRTLLM_GEN_AVAILABLE = (
117
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
118
+)
119
+
120
+
121
+@pytest.mark.skipif(not TRTLLM_GEN_AVAILABLE, reason="Skip for non SM100")
122
+@pytest.mark.parametrize("m", [1, 33])
123
+@pytest.mark.parametrize("n", [7168])
124
+@pytest.mark.parametrize("k", [512])
125
+@pytest.mark.parametrize("e", [384])
126
+@pytest.mark.parametrize("topk", [8])
127
+@pytest.mark.parametrize("group_size", [32])
128
+def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group_size):
129
+ """Compare Marlin INT4 MoE vs FlashInfer TRT-LLM MXINT4 MoE.
130
+
131
+ Uses mxint4_quantize() to generate common INT4 weights + BF16 scales,
132
+ then runs both Marlin and TRT-LLM kernels and compares outputs.
133
+ """
134
+ pytest.importorskip("flashinfer")
135
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_INT4", "1")
136
+
137
+ torch.cuda.manual_seed(0)
138
+
139
+ dtype = torch.bfloat16
140
+
141
+ # DeepSeekV3 routing config (from Kimi-K2-Thinking config.json)
142
+ n_group = 1 # n_group from model config
143
+ topk_group = 1 # topk_group from model config
144
+ routed_scaling = 2.827 # routed_scaling_factor from model config
145
+
146
+ # Input - realistic activation range for LLM (after LayerNorm: mean~0, std~1)
147
+ a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.5
148
+
149
+ # Generate routing logits and bias (DeepSeekV3 expects float logits)
150
+ # Realistic ranges: logits typically [-3, 3], bias [-2, 2]
151
+ routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32) * 1.5
152
+ routing_bias = torch.randn(e, device="cuda", dtype=torch.float32) * 0.8
153
+
154
+ # 1. Generate BF16 weights (SHARED between both paths)
155
+ # Realistic weight initialization: Xavier/Glorot uniform scaling
156
+ # std = sqrt(2 / (fan_in + fan_out))
157
+ std_w1 = (2.0 / (k + 2 * n)) ** 0.5
158
+ std_w2 = (2.0 / (n + k)) ** 0.5
159
+ w1_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) * std_w1
160
+ w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) * std_w2
161
+
162
+ # === Path 1: TRT-LLM FlashInfer MXINT4 MoE ===
163
+ # Similar to: if self.use_flashinfer_mxint4_moe
164
+ # Quantize using MXINT4 method (signed INT4)
165
+ w1_int4, w1_scales = mxint4_quantize_moe_weights(w1_bf16, group_size)
166
+ w2_int4, w2_scales = mxint4_quantize_moe_weights(w2_bf16, group_size)
167
+
168
+ trtllm_weights = prepare_static_weights_for_trtllm_mxint4_moe(
169
+ gemm1_weights=w1_int4,
170
+ gemm1_scales=w1_scales,
171
+ gemm2_weights=w2_int4,
172
+ gemm2_scales=w2_scales,
173
+ )
174
+
175
+ from flashinfer import RoutingMethodType
176
+ from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe
177
+
178
+ # Routing handled internally by trtllm_mxint4_block_scale_moe
179
+ trtllm_output = trtllm_mxint4_block_scale_moe(
180
+ routing_logits=routing_logits,
181
+ routing_bias=routing_bias.to(torch.bfloat16),
182
+ hidden_states=a,
183
+ gemm1_weights=trtllm_weights["gemm1_weights"],
184
+ gemm1_weights_scale=trtllm_weights["gemm1_scales"],
185
+ gemm1_alpha=None,
186
+ gemm1_beta=None,
187
+ gemm1_clamp_limit=None,
188
+ gemm2_weights=trtllm_weights["gemm2_weights"],
189
+ gemm2_weights_scale=trtllm_weights["gemm2_scales"],
190
+ num_experts=e,
191
+ top_k=topk,
192
+ n_group=n_group,
193
+ topk_group=topk_group,
194
+ intermediate_size=n,
195
+ local_expert_offset=0,
196
+ local_num_experts=e,
197
+ routed_scaling_factor=routed_scaling,
198
+ routing_method_type=RoutingMethodType.DeepSeekV3,
199
+ enable_pdl=None,
200
+ output=None,
201
+ tune_max_num_tokens=8192,
202
+ ).to(dtype)
203
+
204
+ # === Path 2: Marlin INT4 MoE ===
205
+ # Similar to: else (non-flashinfer path)
206
+ # Quantize using Marlin's method (UINT4b8)
207
+ w1_marlin, w1_scales_marlin = marlin_quantize_moe_weights(w1_bf16, group_size)
208
+ w2_marlin, w2_scales_marlin = marlin_quantize_moe_weights(w2_bf16, group_size)
209
+
210
+ # Use production routing kernel (same as router.select_experts internally uses)
211
+ topk_weights, topk_ids = grouped_topk(
212
+ hidden_states=a,
213
+ gating_output=routing_logits,
214
+ topk=topk,
215
+ renormalize=False, # DeepSeekV3 doesn't renormalize
216
+ num_expert_group=n_group,
217
+ topk_group=topk_group,
218
+ scoring_func="sigmoid", # DeepSeekV3 uses sigmoid
219
+ routed_scaling_factor=routed_scaling,
220
+ e_score_correction_bias=routing_bias,
221
+ )
222
+
223
+ marlin_output = fused_marlin_moe(
224
+ a,
225
+ w1_marlin,
226
+ w2_marlin,
227
+ None,
228
+ None,
229
+ w1_scales_marlin,
230
+ w2_scales_marlin,
231
+ None, # gating_output not needed when topk_weights/ids provided
232
+ topk_weights,
233
+ topk_ids,
234
+ global_num_experts=e,
235
+ expert_map=None,
236
+ global_scale1=None,
237
+ global_scale2=None,
238
+ g_idx1=None,
239
+ g_idx2=None,
240
+ input_global_scale1=None,
241
+ input_global_scale2=None,
242
+ sort_indices1=None,
243
+ sort_indices2=None,
244
+ w1_zeros=None,
245
+ w2_zeros=None,
246
+ input_dtype=dtype,
247
+ quant_type_id=scalar_types.uint4b8.id,
248
+ is_k_full=True,
249
+ )
250
+
251
+ # Sanity check: manually compute BF16 reference for comparison
252
+ # Use same routing as Marlin path for consistency
253
+ bf16_output = torch.zeros((m, k), device="cuda", dtype=dtype)
254
+ for token_idx in range(m):
255
+ for expert_rank in range(topk):
256
+ expert_id = topk_ids[token_idx, expert_rank].item()
257
+ weight = topk_weights[token_idx, expert_rank].item()
258
+ # w1: [2*n, k] @ [k] -> [2*n]
259
+ up_gate = a[token_idx] @ w1_bf16[expert_id].T # [2*n]
260
+ gate, up = up_gate.chunk(2, dim=0)
261
+ intermediate = torch.nn.functional.silu(gate) * up # [n]
262
+ # w2: [k, n] @ [n] -> [k]
263
+ expert_out = intermediate @ w2_bf16[expert_id].T # [k]
264
+ bf16_output[token_idx] += weight * expert_out
265
+ # Compare against BF16 reference.
266
+ torch.testing.assert_close(marlin_output, bf16_output, atol=0.3, rtol=1.0)
267
+ torch.testing.assert_close(trtllm_output, bf16_output, atol=0.3, rtol=1.0)
268
+
269
+ # Compare against each other for sanity.
270
+ # Note: Different quantization schemes (UINT4b8 vs signed MXINT4) cause
271
+ # some differences
272
+ torch.testing.assert_close(marlin_output, trtllm_output, atol=0.3, rtol=6.0)
vllm/envs.py
+8
-3
@@ -174,6 +174,7 @@ if TYPE_CHECKING:
174
VLLM_USE_FLASHINFER_MOE_FP16: bool = False
175
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
176
VLLM_USE_FLASHINFER_MOE_FP4: bool = False
177
+ VLLM_USE_FLASHINFER_MOE_INT4: bool = False
178
VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
179
"latency"
180
)
@@ -1240,18 +1241,22 @@ environment_variables: dict[str, Callable[[], Any]] = {
1241
"VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
1242
int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "0"))
1243
),
1243
- # Allow use of FlashInfer MoE kernels for fused moe ops.
1244
+ # Allow use of FlashInfer BF16 MoE kernels for fused moe ops.
1245
"VLLM_USE_FLASHINFER_MOE_FP16": lambda: bool(
1246
int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP16", "0"))
1247
),
1247
- # Allow use of FlashInfer MoE kernels for fused moe ops.
1248
+ # Allow use of FlashInfer FP8 MoE kernels for fused moe ops.
1249
"VLLM_USE_FLASHINFER_MOE_FP8": lambda: bool(
1250
int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP8", "0"))
1251
),
1251
- # Allow use of FlashInfer CUTLASS kernels for fused moe ops.
1252
+ # Allow use of FlashInfer NVFP4 MoE kernels for fused moe ops.
1253
"VLLM_USE_FLASHINFER_MOE_FP4": lambda: bool(
1254
int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP4", "0"))
1255
),
1256
+ # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops.
1257
+ "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool(
1258
+ int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0"))
1259
+ ),
1260
# If set to 1, use the FlashInfer
1261
# MXFP8 (activation) x MXFP4 (weight) MoE backend.
1262
"VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8": lambda: bool(
vllm/model_executor/layers/fused_moe/layer.py
+9
-5
@@ -1138,6 +1138,11 @@ class FusedMoE(CustomOp):
1138
return False if return_success else None
1139
# Hereafter, `expert_id` is local physical id
1140
1141
+ # is_transposed: if the dim to shard the weight
1142
+ # should be flipped. Required by GPTQ, compressed-tensors
1143
+ # should be whatever dimension intermediate_size_per_partition is
1144
+ is_transposed = getattr(param, "is_transposed", False)
1145
+
1146
# compressed-tensors checkpoints with packed weights are stored flipped
1147
# TODO (mgoin): check self.quant_method.quant_config.quant_format
1148
# against known CompressionFormat enum values that have this quality
@@ -1145,7 +1150,10 @@ class FusedMoE(CustomOp):
1150
"CompressedTensorsWNA16MarlinMoEMethod",
1151
"CompressedTensorsWNA16MoEMethod",
1152
):
1148
- loaded_weight = loaded_weight.t().contiguous()
1153
+ if is_transposed:
1154
+ loaded_weight = loaded_weight.t().contiguous()
1155
+ else:
1156
+ loaded_weight = loaded_weight
1157
1158
if shard_id not in ("w1", "w2", "w3"):
1159
raise ValueError(f"shard_id must be ['w1','w2','w3'] but got {shard_id}.")
@@ -1183,10 +1191,6 @@ class FusedMoE(CustomOp):
1191
)
1192
return True if return_success else None
1193
1186
- # is_transposed: if the dim to shard the weight
1187
- # should be flipped. Required by GPTQ, compressed-tensors
1188
- # should be whatever dimension intermediate_size_per_partition is
1189
- is_transposed = getattr(param, "is_transposed", False)
1194
shard_dim = SHARD_ID_TO_SHARDED_DIM[shard_id]
1195
if is_transposed:
1196
shard_dim = int(not shard_dim)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+170
-13
@@ -63,6 +63,11 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
63
flashinfer_trtllm_fp4_moe,
64
flashinfer_trtllm_fp4_routed_moe,
65
)
66
+from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
67
+ flashinfer_trtllm_mxint4_moe,
68
+ is_flashinfer_mxint4_moe_available,
69
+ prepare_static_weights_for_trtllm_mxint4_moe,
70
+)
71
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
72
apply_fi_trtllm_fp8_per_tensor_moe,
73
)
@@ -1247,8 +1252,89 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1252
self.actorder = weight_quant.actorder
1253
1254
self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
1250
- self.use_marlin = True
1255
+
1256
self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
1257
+ self.use_flashinfer_mxint4_moe = (
1258
+ is_flashinfer_mxint4_moe_available()
1259
+ and self.group_size == 32
1260
+ and weight_quant.num_bits == 4
1261
+ )
1262
+ self.kernel_backend = (
1263
+ "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin"
1264
+ )
1265
+ logger.info_once(
1266
+ f"Using {self.kernel_backend} backend for WNA16 MoE "
1267
+ f"(group_size={self.group_size}, num_bits={self.num_bits})",
1268
+ scope="local",
1269
+ )
1270
+
1271
+ def get_weight_shape(
1272
+ self,
1273
+ weight_name: str,
1274
+ num_experts: int,
1275
+ hidden_size: int,
1276
+ intermediate_size_per_partition: int,
1277
+ num_groups_w2: int | None = None,
1278
+ num_groups_w13: int | None = None,
1279
+ ) -> tuple[int, int, int]:
1280
+ """
1281
+ Get the shape of the weight based on the weight name, number of experts
1282
+ hidden size, intermediate size per partition, number of groups for w2,
1283
+ and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
1284
+ for weight scales.
1285
+ """
1286
+ if weight_name == "w13_scale":
1287
+ assert num_groups_w13 is not None, (
1288
+ "num_groups_w13 must be provided for weight scales"
1289
+ )
1290
+ if weight_name == "w2_scale":
1291
+ assert num_groups_w2 is not None, (
1292
+ "num_groups_w2 must be provided for weight scales"
1293
+ )
1294
+ w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1295
+ shape_map = {
1296
+ "w13_weight": {
1297
+ "Flashinfer": (
1298
+ num_experts,
1299
+ w13_num_shards * intermediate_size_per_partition,
1300
+ hidden_size // self.packed_factor,
1301
+ ),
1302
+ "Marlin": (
1303
+ num_experts,
1304
+ hidden_size // self.packed_factor,
1305
+ w13_num_shards * intermediate_size_per_partition,
1306
+ ),
1307
+ },
1308
+ "w13_scale": {
1309
+ "Flashinfer": (
1310
+ num_experts,
1311
+ w13_num_shards * intermediate_size_per_partition,
1312
+ num_groups_w13,
1313
+ ),
1314
+ "Marlin": (
1315
+ num_experts,
1316
+ num_groups_w13,
1317
+ w13_num_shards * intermediate_size_per_partition,
1318
+ ),
1319
+ },
1320
+ "w2_weight": {
1321
+ "Flashinfer": (
1322
+ num_experts,
1323
+ hidden_size,
1324
+ intermediate_size_per_partition // self.packed_factor,
1325
+ ),
1326
+ "Marlin": (
1327
+ num_experts,
1328
+ intermediate_size_per_partition // self.packed_factor,
1329
+ hidden_size,
1330
+ ),
1331
+ },
1332
+ "w2_scale": {
1333
+ "Flashinfer": (num_experts, hidden_size, num_groups_w2),
1334
+ "Marlin": (num_experts, num_groups_w2, hidden_size),
1335
+ },
1336
+ }
1337
+ return shape_map[weight_name][self.kernel_backend]
1338
1339
def create_weights(
1340
self,
@@ -1260,19 +1346,23 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1346
**extra_weight_attrs,
1347
):
1348
intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
1263
- w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1349
1350
# Will transpose the loaded weight along the
1351
# intermediate and hidden dim sizes. Will
1352
# shard for TP along the transposed dims
1353
+ is_transposed = self.kernel_backend != "Flashinfer"
1354
extra_weight_attrs.update(
1269
- {"is_transposed": True, "quant_method": self.strategy}
1355
+ {"is_transposed": is_transposed, "quant_method": self.strategy}
1356
)
1357
+
1358
w13_weight = torch.nn.Parameter(
1359
torch.empty(
1273
- num_experts,
1274
- hidden_size // self.packed_factor,
1275
- w13_num_shards * intermediate_size_per_partition,
1360
+ *self.get_weight_shape(
1361
+ "w13_weight",
1362
+ num_experts,
1363
+ hidden_size,
1364
+ intermediate_size_per_partition,
1365
+ ),
1366
dtype=torch.int32,
1367
),
1368
requires_grad=False,
@@ -1282,9 +1372,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1372
1373
w2_weight = torch.nn.Parameter(
1374
torch.empty(
1285
- num_experts,
1286
- intermediate_size_per_partition // self.packed_factor,
1287
- hidden_size,
1375
+ *self.get_weight_shape(
1376
+ "w2_weight",
1377
+ num_experts,
1378
+ hidden_size,
1379
+ intermediate_size_per_partition,
1380
+ ),
1381
dtype=torch.int32,
1382
),
1383
requires_grad=False,
@@ -1315,9 +1408,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1408
1409
w13_scale = torch.nn.Parameter(
1410
torch.ones(
1318
- num_experts,
1319
- num_groups_w13,
1320
- w13_num_shards * intermediate_size_per_partition,
1411
+ *self.get_weight_shape(
1412
+ "w13_scale",
1413
+ num_experts,
1414
+ hidden_size,
1415
+ intermediate_size_per_partition,
1416
+ num_groups_w13=num_groups_w13,
1417
+ ),
1418
dtype=params_dtype,
1419
),
1420
requires_grad=False,
@@ -1326,7 +1423,16 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1423
set_weight_attrs(w13_scale, extra_weight_attrs)
1424
1425
w2_scale = torch.nn.Parameter(
1329
- torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
1426
+ torch.ones(
1427
+ *self.get_weight_shape(
1428
+ "w2_scale",
1429
+ num_experts,
1430
+ hidden_size,
1431
+ intermediate_size_per_partition,
1432
+ num_groups_w2=num_groups_w2,
1433
+ ),
1434
+ dtype=params_dtype,
1435
+ ),
1436
requires_grad=False,
1437
)
1438
layer.register_parameter("w2_weight_scale", w2_scale)
@@ -1396,6 +1502,27 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1502
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1503
num_experts = layer.w13_weight_g_idx.shape[0]
1504
device = layer.w13_weight_g_idx.device
1505
+ if self.kernel_backend == "Flashinfer":
1506
+ dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
1507
+ layer.w13_weight_packed,
1508
+ layer.w13_weight_scale,
1509
+ layer.w2_weight_packed,
1510
+ layer.w2_weight_scale,
1511
+ )
1512
+ replace_parameter(
1513
+ layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"]
1514
+ )
1515
+ replace_parameter(
1516
+ layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"]
1517
+ )
1518
+ replace_parameter(
1519
+ layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"]
1520
+ )
1521
+ replace_parameter(
1522
+ layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"]
1523
+ )
1524
+ return None
1525
+
1526
is_a_8bit = (
1527
self.marlin_input_dtype is not None
1528
and self.marlin_input_dtype.itemsize == 1
@@ -1560,6 +1687,35 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1687
is_k_full=self.is_k_full,
1688
)
1689
1690
+ @property
1691
+ def is_monolithic(self) -> bool:
1692
+ return self.kernel_backend == "Flashinfer"
1693
+
1694
+ def apply_monolithic(
1695
+ self,
1696
+ layer: FusedMoE,
1697
+ x: torch.Tensor,
1698
+ router_logits: torch.Tensor,
1699
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1700
+ assert self.kernel_backend == "Flashinfer"
1701
+ return flashinfer_trtllm_mxint4_moe(
1702
+ x=x,
1703
+ router_logits=router_logits,
1704
+ w13_weight_packed=layer.w13_weight_packed,
1705
+ w13_weight_scale=layer.w13_weight_scale,
1706
+ w2_weight_packed=layer.w2_weight_packed,
1707
+ w2_weight_scale=layer.w2_weight_scale,
1708
+ global_num_experts=layer.global_num_experts,
1709
+ top_k=layer.top_k,
1710
+ intermediate_size_per_partition=layer.intermediate_size_per_partition,
1711
+ local_num_experts=layer.local_num_experts,
1712
+ ep_rank=layer.ep_rank,
1713
+ num_expert_group=layer.num_expert_group,
1714
+ topk_group=layer.topk_group,
1715
+ e_score_correction_bias=layer.e_score_correction_bias,
1716
+ routing_method_type=layer.routing_method_type,
1717
+ )
1718
+
1719
def apply(
1720
self,
1721
layer: FusedMoE,
@@ -1567,6 +1723,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1723
topk_weights: torch.Tensor,
1724
topk_ids: torch.Tensor,
1725
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1726
+ assert self.kernel_backend == "Marlin"
1727
return fused_marlin_moe(
1728
x,
1729
layer.w13_weight_packed,
vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
new
+266
@@ -0,0 +1,266 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Utility helpers for MxInt4 + FlashInfer fused-MoE path"""
4
+
5
+import functools
6
+
7
+import torch
8
+
9
+import vllm.envs as envs
10
+from vllm.logger import init_logger
11
+from vllm.platforms import current_platform
12
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
13
+
14
+__all__ = [
15
+ "prepare_static_weights_for_trtllm_mxint4_moe",
16
+ "flashinfer_trtllm_mxint4_moe",
17
+ "is_flashinfer_mxint4_moe_available",
18
+]
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+@functools.cache
24
+def is_flashinfer_mxint4_moe_available() -> bool:
25
+ """Return `True` when FlashInfer MxInt4 kernels can be used."""
26
+ return (
27
+ envs.VLLM_USE_FLASHINFER_MOE_INT4
28
+ and has_flashinfer_trtllm_fused_moe()
29
+ and current_platform.is_cuda()
30
+ and current_platform.is_device_capability_family(100)
31
+ )
32
+
33
+
34
+def prepare_static_weights_for_trtllm_mxint4_moe(
35
+ gemm1_weights: torch.Tensor,
36
+ gemm1_scales: torch.Tensor,
37
+ gemm2_weights: torch.Tensor,
38
+ gemm2_scales: torch.Tensor,
39
+) -> dict[str, torch.Tensor]:
40
+ """
41
+ Prepare MxInt4 weights for TRT-LLM kernel.
42
+
43
+ Input:
44
+ gemm1_weights: [num_experts, 2*intermediate_size, hidden_size//8] int32
45
+ (checkpoint uint4b8 packed) or uint8 (already packed signed int4)
46
+ gemm1_scales: [num_experts, 2*intermediate_size, hidden_size//32] bf16
47
+ gemm2_weights: [num_experts, hidden_size, intermediate_size//8] int32
48
+ (checkpoint uint4b8 packed) or uint8 (already packed signed int4)
49
+ gemm2_scales: [num_experts, hidden_size, intermediate_size//32] bf16
50
+
51
+ Returns:
52
+ Dict with keys 'gemm1_weights', 'gemm1_scales', 'gemm2_weights',
53
+ 'gemm2_scales' containing shuffled/packed tensors ready for kernel
54
+ """
55
+ from flashinfer import block_scale_interleave
56
+ from flashinfer.fused_moe import (
57
+ convert_to_block_layout,
58
+ )
59
+ from flashinfer.fused_moe.core import (
60
+ _maybe_get_cached_w3_w1_permute_indices,
61
+ get_w2_permute_indices_with_cache,
62
+ )
63
+
64
+ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
65
+ reorder_w1w3_to_w3w1,
66
+ )
67
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
68
+ convert_packed_uint4b8_to_signed_int4_inplace,
69
+ )
70
+
71
+ device = gemm1_weights.device
72
+ assert gemm1_weights.ndim == 3, (
73
+ f"Expected a 3D gemm1_weights tensor, got {gemm1_weights.shape}"
74
+ )
75
+ assert gemm1_scales.ndim == 3, (
76
+ f"Expected a 3D gemm1_scales tensor, got {gemm1_scales.shape}"
77
+ )
78
+ assert gemm2_weights.ndim == 3, (
79
+ f"Expected a 3D gemm2_weights tensor, got {gemm2_weights.shape}"
80
+ )
81
+ assert gemm2_scales.ndim == 3, (
82
+ f"Expected a 3D gemm2_scales tensor, got {gemm2_scales.shape}"
83
+ )
84
+
85
+ # Convert checkpoint format (uint4b8 in int32) to signed int4
86
+ # Checkpoint stores INT4 as unsigned [0, 15], kernel expects signed [-8, 7]
87
+ if gemm1_weights.dtype == torch.int32 and gemm2_weights.dtype == torch.int32:
88
+ convert_packed_uint4b8_to_signed_int4_inplace(gemm1_weights)
89
+ convert_packed_uint4b8_to_signed_int4_inplace(gemm2_weights)
90
+
91
+ gemm1_weights, gemm1_scales = reorder_w1w3_to_w3w1(
92
+ gemm1_weights, gemm1_scales, dim=-2
93
+ )
94
+
95
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
96
+ num_experts = gemm1_weights.shape[0]
97
+
98
+ # Convert quantized weights to proper formats -
99
+ gemm1_weights_mxint4 = gemm1_weights.view(torch.uint8)
100
+ assert gemm1_scales.dtype == torch.bfloat16
101
+ gemm2_weights_mxint4 = gemm2_weights.view(torch.uint8)
102
+ assert gemm2_scales.dtype == torch.bfloat16
103
+
104
+ epilogue_tile_m = 128
105
+ gemm1_weights_mxint4_shuffled = []
106
+ gemm1_scales_shuffled = []
107
+ gemm2_weights_mxint4_shuffled = []
108
+ gemm2_scales_shuffled = []
109
+
110
+ for i in range(num_experts):
111
+ # Calculate the permute indices for the following:
112
+ # 1. Reorder rows of W1 and scales for fused gated activation
113
+ # 2. Shuffle weights and scaling factors for transposed mma output
114
+ # for both w3_w1 and w2 weights and scale factors
115
+ permute_indices = _maybe_get_cached_w3_w1_permute_indices(
116
+ _cache_permute_indices,
117
+ gemm1_weights_mxint4[i],
118
+ epilogue_tile_m,
119
+ )
120
+ gemm1_weights_shuffled = gemm1_weights_mxint4[i][
121
+ permute_indices.to(gemm1_weights.device)
122
+ ].contiguous()
123
+ permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
124
+ _cache_permute_indices,
125
+ gemm1_scales[i],
126
+ epilogue_tile_m,
127
+ num_elts_per_sf=32,
128
+ ).to(device)
129
+ gemm1_scales_shuffled.append(
130
+ block_scale_interleave(gemm1_scales[i][permute_sf_indices].contiguous())
131
+ )
132
+
133
+ permute_indices = get_w2_permute_indices_with_cache(
134
+ _cache_permute_indices,
135
+ gemm2_weights_mxint4[i],
136
+ epilogue_tile_m,
137
+ )
138
+ gemm2_weights_shuffled = gemm2_weights_mxint4[i][
139
+ permute_indices.to(gemm2_weights.device)
140
+ ].contiguous()
141
+
142
+ permute_sf_indices = get_w2_permute_indices_with_cache(
143
+ _cache_permute_indices,
144
+ gemm2_scales[i],
145
+ epilogue_tile_m,
146
+ num_elts_per_sf=16,
147
+ )
148
+ gemm2_scales_shuffled.append(
149
+ block_scale_interleave(
150
+ gemm2_scales[i][permute_sf_indices.to(gemm2_scales.device)].contiguous()
151
+ )
152
+ )
153
+
154
+ block_k = 128
155
+ gemm1_weights_shuffled = convert_to_block_layout(
156
+ gemm1_weights_shuffled.view(torch.uint8), block_k
157
+ )
158
+ gemm2_weights_shuffled = convert_to_block_layout(
159
+ gemm2_weights_shuffled.view(torch.uint8), block_k
160
+ )
161
+
162
+ gemm1_weights_mxint4_shuffled.append(gemm1_weights_shuffled)
163
+ gemm2_weights_mxint4_shuffled.append(gemm2_weights_shuffled)
164
+
165
+ gemm1_weights_mxint4_shuffled = torch.stack(gemm1_weights_mxint4_shuffled)
166
+ gemm2_weights_mxint4_shuffled = torch.stack(gemm2_weights_mxint4_shuffled)
167
+ gemm1_scales_shuffled = torch.stack(gemm1_scales_shuffled).view(torch.bfloat16)
168
+ gemm2_scales_shuffled = torch.stack(gemm2_scales_shuffled).view(torch.bfloat16)
169
+ return {
170
+ "gemm1_weights": gemm1_weights_mxint4_shuffled,
171
+ "gemm1_scales": gemm1_scales_shuffled,
172
+ "gemm2_weights": gemm2_weights_mxint4_shuffled,
173
+ "gemm2_scales": gemm2_scales_shuffled,
174
+ }
175
+
176
+
177
+def flashinfer_trtllm_mxint4_moe(
178
+ x: torch.Tensor,
179
+ router_logits: torch.Tensor,
180
+ w13_weight_packed: torch.Tensor,
181
+ w13_weight_scale: torch.Tensor,
182
+ w2_weight_packed: torch.Tensor,
183
+ w2_weight_scale: torch.Tensor,
184
+ global_num_experts: int,
185
+ top_k: int,
186
+ intermediate_size_per_partition: int,
187
+ local_num_experts: int,
188
+ ep_rank: int = 0,
189
+ num_expert_group: int | None = None,
190
+ topk_group: int | None = None,
191
+ e_score_correction_bias: torch.Tensor | None = None,
192
+ routing_method_type: int | None = None,
193
+) -> torch.Tensor:
194
+ """
195
+ Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
196
+
197
+ Args:
198
+ x: Input hidden states. dtype: bfloat16
199
+ router_logits: Router logits for expert selection. dtype: bfloat16/float32
200
+ w13_weight_packed: Packed gate+up weights. dtype: uint8
201
+ w13_weight_scale: Scales for gate+up weights. dtype: bfloat16
202
+ w2_weight_packed: Packed down weights. dtype: uint8
203
+ w2_weight_scale: Scales for down weights. dtype: bfloat16
204
+ global_num_experts: Total number of experts across all ranks
205
+ top_k: Number of experts to select per token
206
+ intermediate_size_per_partition: Intermediate size per partition
207
+ local_num_experts: Number of experts on this rank
208
+ ep_rank: Expert parallelism rank (default: 0)
209
+ num_expert_group: Number of expert groups (default: None -> 0)
210
+ topk_group: Top-k within groups (default: None -> 0)
211
+ e_score_correction_bias: Optional routing bias. dtype: bfloat16
212
+ routing_method_type: FlashInfer RoutingMethodType enum value
213
+
214
+ Returns:
215
+ Output tensor from MoE layer. dtype: same as x (bfloat16)
216
+ """
217
+ from flashinfer import RoutingMethodType
218
+ from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe
219
+
220
+ assert x.dtype == torch.bfloat16, f"x dtype must be bfloat16, got {x.dtype}"
221
+ assert w13_weight_packed.dtype == torch.uint8, (
222
+ f"w13_weight_packed dtype must be uint8, got {w13_weight_packed.dtype}"
223
+ )
224
+ assert w13_weight_scale.dtype == torch.bfloat16, (
225
+ f"w13_weight_scale dtype must be bfloat16, got {w13_weight_scale.dtype}"
226
+ )
227
+ assert w2_weight_packed.dtype == torch.uint8, (
228
+ f"w2_weight_packed dtype must be uint8, got {w2_weight_packed.dtype}"
229
+ )
230
+ assert w2_weight_scale.dtype == torch.bfloat16, (
231
+ f"w2_weight_scale dtype must be bfloat16, got {w2_weight_scale.dtype}"
232
+ )
233
+
234
+ routing_bias = None
235
+ if e_score_correction_bias is not None:
236
+ routing_bias = e_score_correction_bias.to(torch.bfloat16)
237
+
238
+ if routing_method_type == RoutingMethodType.DeepSeekV3:
239
+ router_logits = router_logits.to(torch.float32)
240
+
241
+ out = trtllm_mxint4_block_scale_moe(
242
+ routing_logits=router_logits,
243
+ routing_bias=routing_bias,
244
+ hidden_states=x,
245
+ gemm1_weights=w13_weight_packed.data,
246
+ gemm1_weights_scale=w13_weight_scale.data,
247
+ gemm1_alpha=None,
248
+ gemm1_beta=None,
249
+ gemm1_clamp_limit=None,
250
+ gemm2_weights=w2_weight_packed.data,
251
+ gemm2_weights_scale=w2_weight_scale.data,
252
+ num_experts=global_num_experts,
253
+ top_k=top_k,
254
+ n_group=num_expert_group if num_expert_group is not None else 0,
255
+ topk_group=topk_group if topk_group is not None else 0,
256
+ intermediate_size=intermediate_size_per_partition,
257
+ local_expert_offset=ep_rank * local_num_experts,
258
+ local_num_experts=local_num_experts,
259
+ routed_scaling_factor=None,
260
+ routing_method_type=routing_method_type,
261
+ enable_pdl=None,
262
+ output=None,
263
+ tune_max_num_tokens=8192,
264
+ ).to(x.dtype)
265
+
266
+ return out
vllm/utils/flashinfer.py
+2
-2
@@ -129,12 +129,11 @@ scaled_fp4_grouped_quantize = _lazy_import_wrapper(
129
"flashinfer", "scaled_fp4_grouped_quantize"
130
)
131
nvfp4_block_scale_interleave = _lazy_import_wrapper(
132
- "flashinfer", "nvfp4_block_scale_interleave"
132
+ "flashinfer.fp4_quantization", "block_scale_interleave"
133
)
134
trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
135
"flashinfer", "trtllm_fp4_block_scale_moe"
136
)
137
-
137
# Special case for autotune since it returns a context manager
138
autotune = _lazy_import_wrapper(
139
"flashinfer.autotuner",
@@ -196,6 +195,7 @@ def has_flashinfer_trtllm_fused_moe() -> bool:
195
("flashinfer.fused_moe", "trtllm_fp8_block_scale_moe"),
196
("flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe"),
197
("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),
198
+ ("flashinfer.fused_moe", "trtllm_mxint4_block_scale_moe"),
199
]
200
for module_name, attr_name in required_functions:
201
mod = _get_submodule(module_name)