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 14, 05:30 UTC
[gpt-oss] Cache permute indices for faster MXFP4 MoE layer loading (#24154)
Signed-off-by: Wei Wei <wwei6@meta.com>
Wei committed
Sep 9, 2025 at 21:27 UTC
0efdb5c3bad240121d91083189565df0b82502c2
2 files changed
+146
-35
tests/kernels/moe/test_mxfp4_moe.py
+86
-17
@@ -24,6 +24,8 @@ if TRTLLM_GEN_MXFP4_AVAILABLE:
24
next_positive_power_of_2,
25
reorder_rows_for_gated_act_gemm, shuffle_matrix_a,
26
shuffle_matrix_sf_a, trtllm_fp4_block_scale_moe)
27
+ from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
28
+ from flashinfer.fused_moe.core import _maybe_get_cached_w2_permute_indices
29
30
31
@dataclass
@@ -204,6 +206,7 @@ def tg_mxfp4_moe(
206
alpha,
207
beta,
208
limit,
209
+ transpose_optimized: bool = False,
210
) -> torch.Tensor:
211
sf_block_size = 32
212
assert (w13_weight.dim() == 3 and w13_weight.shape[0] == num_experts
@@ -267,22 +270,85 @@ def tg_mxfp4_moe(
270
gemm1_bias_shuffled = []
271
gemm2_bias_shuffled = []
272
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
270
- for i in range(num_experts):
271
- gemm1_weights_shuffled.append(
272
- shuffle_matrix_a(w13_weight[i].view(torch.uint8), epilogue_tile_m))
273
- gemm1_scales_shuffled.append(
274
- shuffle_matrix_sf_a(w13_weight_scale[i].view(torch.uint8),
275
- epilogue_tile_m))
276
-
277
- gemm2_weights_shuffled.append(
278
- shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m))
279
- gemm2_scales_shuffled.append(
280
- shuffle_matrix_sf_a(w2_weight_scale[i].view(torch.uint8),
281
- epilogue_tile_m))
282
- gemm1_bias_shuffled.append(
283
- shuffle_matrix_a(w13_bias[i].reshape(-1, 1), epilogue_tile_m))
284
- gemm2_bias_shuffled.append(
285
- shuffle_matrix_a(w2_bias[i].reshape(-1, 1), epilogue_tile_m))
273
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
274
+ if transpose_optimized:
275
+ for i in range(num_experts):
276
+ # w13 weight shuffling
277
+ permute_indices = _maybe_get_cached_w2_permute_indices(
278
+ _cache_permute_indices,
279
+ w13_weight[i].view(torch.uint8),
280
+ epilogue_tile_m,
281
+ )
282
+ gemm1_weights_shuffled.append(w13_weight[i].view(
283
+ torch.uint8)[permute_indices.to(
284
+ w13_weight.device)].contiguous())
285
+ # w13 scale shuffling
286
+ permute_sf_indices = _maybe_get_cached_w2_permute_indices(
287
+ _cache_permute_indices,
288
+ w13_weight_scale[i].view(torch.uint8),
289
+ epilogue_tile_m,
290
+ num_elts_per_sf=16,
291
+ )
292
+ gemm1_scales_shuffled.append(
293
+ nvfp4_block_scale_interleave(w13_weight_scale[i].view(
294
+ torch.uint8)[permute_sf_indices.to(
295
+ w13_weight_scale.device)].contiguous()))
296
+ # w13 bias shuffling
297
+ permute_bias_indices = _maybe_get_cached_w2_permute_indices(
298
+ _cache_permute_indices,
299
+ w13_bias[i].clone().reshape(-1, 1),
300
+ epilogue_tile_m,
301
+ )
302
+ gemm1_bias_shuffled.append(w13_bias[i].clone().reshape(
303
+ -1, 1)[permute_bias_indices.to(w13_bias.device)].contiguous())
304
+ # w2 weight shuffling
305
+ permute_indices = _maybe_get_cached_w2_permute_indices(
306
+ _cache_permute_indices,
307
+ w2_weight[i].view(torch.uint8),
308
+ epilogue_tile_m,
309
+ )
310
+ gemm2_weights_shuffled.append(w2_weight[i].view(
311
+ torch.uint8)[permute_indices.to(
312
+ w2_weight.device)].contiguous())
313
+ # w2 scale shuffling
314
+ permute_sf_indices = _maybe_get_cached_w2_permute_indices(
315
+ _cache_permute_indices,
316
+ w2_weight_scale[i].view(torch.uint8),
317
+ epilogue_tile_m,
318
+ num_elts_per_sf=16,
319
+ )
320
+ gemm2_scales_shuffled.append(
321
+ nvfp4_block_scale_interleave(w2_weight_scale[i].view(
322
+ torch.uint8)[permute_sf_indices.to(
323
+ w2_weight_scale.device)].contiguous()))
324
+ # w2 bias shuffling
325
+ permute_indices = _maybe_get_cached_w2_permute_indices(
326
+ _cache_permute_indices,
327
+ w2_bias[i].clone().reshape(-1, 1),
328
+ epilogue_tile_m,
329
+ )
330
+ gemm2_bias_shuffled.append(w2_bias[i].clone().reshape(
331
+ -1, 1)[permute_indices.to(w2_bias.device)].contiguous())
332
+
333
+ else:
334
+ for i in range(num_experts):
335
+ gemm1_weights_shuffled.append(
336
+ shuffle_matrix_a(w13_weight[i].view(torch.uint8),
337
+ epilogue_tile_m))
338
+ gemm1_scales_shuffled.append(
339
+ shuffle_matrix_sf_a(w13_weight_scale[i].view(torch.uint8),
340
+ epilogue_tile_m))
341
+
342
+ gemm2_weights_shuffled.append(
343
+ shuffle_matrix_a(w2_weight[i].view(torch.uint8),
344
+ epilogue_tile_m))
345
+ gemm2_scales_shuffled.append(
346
+ shuffle_matrix_sf_a(w2_weight_scale[i].view(torch.uint8),
347
+ epilogue_tile_m))
348
+ gemm1_bias_shuffled.append(
349
+ shuffle_matrix_a(w13_bias[i].reshape(-1, 1), epilogue_tile_m))
350
+ gemm2_bias_shuffled.append(
351
+ shuffle_matrix_a(w2_bias[i].reshape(-1, 1), epilogue_tile_m))
352
353
w13_weight = torch.stack(gemm1_weights_shuffled)
354
w13_weight_scale = torch.stack(gemm1_scales_shuffled).reshape(
@@ -356,6 +422,7 @@ def check_accuracy(a, b, atol, rtol, percent):
422
@pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None),
423
(1.702, 1.0, 7.0)])
424
@pytest.mark.parametrize("act_type", ['mxfp8', 'bf16'])
425
+@pytest.mark.parametrize("transpose_optimized", [False, True])
426
@pytest.mark.skipif(
427
not TRTLLM_GEN_MXFP4_AVAILABLE,
428
reason="nvidia gpu and compute capability sm100 is required for this test")
@@ -369,6 +436,7 @@ def test_trtllm_gen_mxfp4_fused_moe(
436
beta: float,
437
limit: Optional[float],
438
act_type: str,
439
+ transpose_optimized: bool,
440
):
441
seed = 42
442
torch.manual_seed(seed)
@@ -470,6 +538,7 @@ def test_trtllm_gen_mxfp4_fused_moe(
538
act_type,
539
alpha=alpha,
540
beta=beta,
473
- limit=limit)
541
+ limit=limit,
542
+ transpose_optimized=transpose_optimized)
543
# relatively loose check since the mxfp4 quantization is less accurate
544
check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)
vllm/model_executor/layers/quantization/mxfp4.py
+60
-18
@@ -122,6 +122,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
122
"MXFP4 MoE is enabled on Blackwell but FlashInfer "
123
"is not available. This may result in degraded performance. "
124
"Please `pip install vllm[flashinfer]` for best results.")
125
+ self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
126
127
def _should_use_marlin(self):
128
if envs.VLLM_MXFP4_USE_MARLIN is not None:
@@ -266,7 +267,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
267
if self.use_marlin:
268
prepare_moe_fp4_layer_for_marlin(layer)
269
elif should_use_flashinfer_mxfp4():
269
- from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
270
+ from flashinfer.fp4_quantization import (
271
+ nvfp4_block_scale_interleave)
272
+ from flashinfer.fused_moe.core import (
273
+ _maybe_get_cached_w2_permute_indices)
274
layer.gemm1_alpha = Parameter(torch.tensor(
275
[1.702] * self.num_experts, dtype=torch.float32).cuda(),
276
requires_grad=False)
@@ -343,25 +347,63 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
347
gemm2_bias_shuffled = []
348
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
349
for i in range(self.num_experts):
346
- gemm1_weights_mxfp4_shuffled.append(
347
- shuffle_matrix_a(w13_weight[i].view(torch.uint8),
348
- epilogue_tile_m))
350
+ # w13 weight shuffling
351
+ permute_indices = _maybe_get_cached_w2_permute_indices(
352
+ self._cache_permute_indices,
353
+ w13_weight[i].view(torch.uint8),
354
+ epilogue_tile_m,
355
+ )
356
+ gemm1_weights_mxfp4_shuffled.append(w13_weight[i].view(
357
+ torch.uint8)[permute_indices.to(
358
+ w13_weight.device)].contiguous())
359
+ # w13 scale shuffling
360
+ permute_sf_indices = _maybe_get_cached_w2_permute_indices(
361
+ self._cache_permute_indices,
362
+ w13_weight_scale[i].view(torch.uint8),
363
+ epilogue_tile_m,
364
+ num_elts_per_sf=16,
365
+ )
366
gemm1_scales_mxfp4_shuffled.append(
350
- shuffle_matrix_sf_a(w13_weight_scale[i].view(torch.uint8),
351
- epilogue_tile_m))
352
- gemm1_bias_shuffled.append(
353
- shuffle_matrix_a(w13_bias[i].clone().reshape(-1, 1),
354
- epilogue_tile_m))
355
-
356
- gemm2_weights_mxfp4_shuffled.append(
357
- shuffle_matrix_a(w2_weight[i].view(torch.uint8),
358
- epilogue_tile_m))
367
+ nvfp4_block_scale_interleave(w13_weight_scale[i].view(
368
+ torch.uint8)[permute_sf_indices.to(
369
+ w13_weight_scale.device)].contiguous()))
370
+ # w13 bias shuffling
371
+ permute_bias_indices = _maybe_get_cached_w2_permute_indices(
372
+ self._cache_permute_indices,
373
+ w13_bias[i].clone().reshape(-1, 1),
374
+ epilogue_tile_m,
375
+ )
376
+ gemm1_bias_shuffled.append(w13_bias[i].clone().reshape(
377
+ -1,
378
+ 1)[permute_bias_indices.to(w13_bias.device)].contiguous())
379
+ # w2 weight shuffling
380
+ permute_indices = _maybe_get_cached_w2_permute_indices(
381
+ self._cache_permute_indices,
382
+ w2_weight[i].view(torch.uint8),
383
+ epilogue_tile_m,
384
+ )
385
+ gemm2_weights_mxfp4_shuffled.append(w2_weight[i].view(
386
+ torch.uint8)[permute_indices.to(
387
+ w2_weight.device)].contiguous())
388
+ # w2 scale shuffling
389
+ permute_sf_indices = _maybe_get_cached_w2_permute_indices(
390
+ self._cache_permute_indices,
391
+ w2_weight_scale[i].view(torch.uint8),
392
+ epilogue_tile_m,
393
+ num_elts_per_sf=16,
394
+ )
395
gemm2_scales_mxfp4_shuffled.append(
360
- shuffle_matrix_sf_a(w2_weight_scale[i].view(torch.uint8),
361
- epilogue_tile_m))
362
- gemm2_bias_shuffled.append(
363
- shuffle_matrix_a(w2_bias[i].clone().reshape(-1, 1),
364
- epilogue_tile_m))
396
+ nvfp4_block_scale_interleave(w2_weight_scale[i].view(
397
+ torch.uint8)[permute_sf_indices.to(
398
+ w2_weight_scale.device)].contiguous()))
399
+ # w2 bias shuffling
400
+ permute_indices = _maybe_get_cached_w2_permute_indices(
401
+ self._cache_permute_indices,
402
+ w2_bias[i].clone().reshape(-1, 1),
403
+ epilogue_tile_m,
404
+ )
405
+ gemm2_bias_shuffled.append(w2_bias[i].clone().reshape(
406
+ -1, 1)[permute_indices.to(w2_bias.device)].contiguous())
407
408
w13_weight = torch.stack(gemm1_weights_mxfp4_shuffled)
409
w13_weight_scale = torch.stack(