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

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 12, 18:30 UTC

Add TRTLLM MoE NVFP4 kernel to CompressedTensorsW4A4MoeMethod (#28892)

Signed-off-by: mingyuanm <mingyuanm@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Mingyuan Ma committed Nov 21, 2025 at 08:54 UTC b4c8fbaae2592501f442817f86e32cfeb795d81f

3 files changed +358 -210

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+122 -20

@@ -8,6 +8,7 @@ from enum import Enum

8 import torch

9 from compressed_tensors import CompressionFormat

10 from compressed_tensors.quantization import ActivationOrdering, QuantizationStrategy

11 +from torch.nn.parameter import Parameter

13 import vllm.envs as envs

14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

@@ -50,9 +51,15 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (

51 from vllm.model_executor.layers.quantization.utils import replace_parameter

52 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

53 build_flashinfer_fp4_cutlass_moe_prepare_finalize,

54 + flashinfer_trtllm_fp4_moe,

55 + prepare_static_weights_for_trtllm_fp4_moe,

56 reorder_w1w3_to_w3w1,

57 select_nvfp4_gemm_impl,

58 )

59 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

60 + FlashinferMoeBackend,

61 + get_flashinfer_moe_backend,

62 +)

63 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

64 expert_weight_is_col_major,

65 requant_weight_ue8m0_inplace,

@@ -193,6 +200,13 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

200 self.allow_flashinfer = _nvfp4.allow_flashinfer

201 self.use_marlin = _nvfp4.use_marlin

202 self.group_size = 16

203 + self.flashinfer_moe_backend = None

204 + if self.allow_flashinfer:

205 + self.flashinfer_moe_backend = get_flashinfer_moe_backend()

206 + logger.info_once(

207 + f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"

208 + " for CompressedTensorsW4A4MoeMethod."

209 + )

210

211 def create_weights(

212 self,

@@ -344,21 +358,20 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

358 if self.use_marlin:

359 prepare_moe_fp4_layer_for_marlin(layer)

360 return

347 -

348 - # swizzle weight scales

349 - layer.w13_weight_scale = torch.nn.Parameter(

350 - swizzle_blockscale(layer.w13_weight_scale), requires_grad=False

351 - )

352 -

353 - layer.w2_weight_scale = torch.nn.Parameter(

354 - swizzle_blockscale(layer.w2_weight_scale), requires_grad=False

355 - )

356 -

361 # w13

358 - w13_input_global_scale = layer.w13_input_global_scale.max(dim=1).values.to(

359 - torch.float32

360 - )

361 -

362 + if (

363 + self.allow_flashinfer

364 + and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

365 + ):

366 + w13_input_global_scale = (

367 + layer.w13_input_global_scale.min()

368 + .to(torch.float32)

369 + .expand(layer.num_experts)

370 + )

371 + else:

372 + w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to(

373 + torch.float32

374 + )

375 layer.g1_alphas = torch.nn.Parameter(

376 ((1 / w13_input_global_scale) * layer.w13_weight_scale_2),

377 requires_grad=False,

@@ -369,22 +382,92 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

382 )

383

384 # w2

385 + if (

386 + self.allow_flashinfer

387 + and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

388 + ):

389 + w2_input_global_scale = (

390 + layer.w2_input_global_scale.min()

391 + .to(torch.float32)

392 + .expand(layer.num_experts)

393 + )

394 + else:

395 + w2_input_global_scale = layer.w2_input_global_scale

396 +

397 layer.g2_alphas = torch.nn.Parameter(

373 - ((1 / layer.w2_input_global_scale) * layer.w2_weight_scale_2).to(

374 - torch.float32

375 - ),

398 + ((1 / w2_input_global_scale) * layer.w2_weight_scale_2).to(torch.float32),

399 requires_grad=False,

400 )

401

402 layer.w2_input_scale_quant = torch.nn.Parameter(

380 - (layer.w2_input_global_scale), requires_grad=False

403 + (w2_input_global_scale), requires_grad=False

404 )

405

406 + # TensorRT-LLM specific processing

407 + if (

408 + self.allow_flashinfer

409 + and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

410 + ):

411 + # Prepare static weights for TRT-LLM kernel

412 + # alternate: prepare_static_weight_layouts_for_trtllm_moe

413 + (

414 + gemm1_weights_fp4_shuffled,

415 + gemm1_scales_fp4_shuffled,

416 + gemm2_weights_fp4_shuffled,

417 + gemm2_scales_fp4_shuffled,

418 + ) = prepare_static_weights_for_trtllm_fp4_moe(

419 + layer.w13_weight,

420 + layer.w2_weight,

421 + layer.w13_weight_scale,

422 + layer.w2_weight_scale,

423 + layer.w2_weight.size(-2), # hidden_size

424 + layer.w13_weight.size(-2) // 2, # intermediate_size

425 + layer.w13_weight.size(0), # num_experts

426 + )

427 + logger.debug_once("Finished shuffling weights for TRT-LLM MOE")

428 +

429 + layer.gemm1_weights_fp4_shuffled = Parameter(

430 + gemm1_weights_fp4_shuffled, requires_grad=False

431 + )

432 + layer.gemm2_weights_fp4_shuffled = Parameter(

433 + gemm2_weights_fp4_shuffled, requires_grad=False

434 + )

435 + layer.gemm1_scales_fp4_shuffled = Parameter(

436 + gemm1_scales_fp4_shuffled, requires_grad=False

437 + )

438 + layer.gemm2_scales_fp4_shuffled = Parameter(

439 + gemm2_scales_fp4_shuffled, requires_grad=False

440 + )

441 +

442 + # Additional parameter needed for TRT-LLM

443 + layer.g1_scale_c = Parameter(

444 + (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),

445 + requires_grad=False,

446 + )

447 +

448 + # Clean up weights that won't be used by TRT-LLM

449 + del layer.w2_weight

450 + del layer.w2_weight_scale

451 + del layer.w13_weight

452 + del layer.w13_weight_scale

453 + else:

454 + # swizzle weight scales

455 + layer.w13_weight_scale = torch.nn.Parameter(

456 + swizzle_blockscale(layer.w13_weight_scale), requires_grad=False

457 + )

458 +

459 + layer.w2_weight_scale = torch.nn.Parameter(

460 + swizzle_blockscale(layer.w2_weight_scale), requires_grad=False

461 + )

462 +

463 def maybe_make_prepare_finalize(

464 self,

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

466 ) -> mk.FusedMoEPrepareAndFinalize | None:

387 - if self.use_marlin:

467 + if self.use_marlin or (

468 + self.allow_flashinfer

469 + and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

470 + ):

471 return None

472 elif not self.allow_flashinfer:

473 return super().maybe_make_prepare_finalize(routing_tables)

@@ -411,7 +494,10 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

494 def get_fused_moe_quant_config(

495 self, layer: torch.nn.Module

496 ) -> FusedMoEQuantConfig | None:

414 - if self.use_marlin:

497 + if (

498 + self.use_marlin

499 + or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

500 + ):

501 return None

502

503 return nvfp4_moe_quant_config(

@@ -452,6 +538,22 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

538 )

539 assert activation == "silu", "Only SiLU activation is supported."

540

541 + if (

542 + self.allow_flashinfer

543 + and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

544 + ):

545 + return flashinfer_trtllm_fp4_moe(

546 + layer=layer,

547 + x=x,

548 + router_logits=router_logits,

549 + top_k=top_k,

550 + global_num_experts=global_num_experts,

551 + num_expert_group=num_expert_group,

552 + topk_group=topk_group,

553 + custom_routing_function=custom_routing_function,

554 + e_score_correction_bias=e_score_correction_bias,

555 + )

556 +

557 topk_weights, topk_ids, _ = FusedMoE.select_experts(

558 hidden_states=x,

559 router_logits=router_logits,

vllm/model_executor/layers/quantization/modelopt.py

+15 -190

@@ -15,7 +15,6 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

15 from vllm.logger import init_logger

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

17 FusedMoEQuantConfig,

18 - RoutingMethodType,

18 fp8_w8a8_moe_quant_config,

19 nvfp4_moe_quant_config,

20 )

@@ -38,6 +37,8 @@ from vllm.model_executor.layers.quantization.base_config import (

37 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

38 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

39 build_flashinfer_fp4_cutlass_moe_prepare_finalize,

40 + flashinfer_trtllm_fp4_moe,

41 + prepare_static_weights_for_trtllm_fp4_moe,

42 reorder_w1w3_to_w3w1,

43 select_nvfp4_gemm_impl,

44 )

@@ -1136,7 +1137,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1137 self.allow_flashinfer = _nvfp4.allow_flashinfer

1138 self.use_marlin = _nvfp4.use_marlin

1139 self.flashinfer_moe_backend = None

1139 - self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

1140 if self.allow_flashinfer:

1141 self.flashinfer_moe_backend = get_flashinfer_moe_backend()

1142 logger.info_once(

@@ -1303,138 +1303,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1303 )

1304 layer.register_parameter("w2_input_scale", w2_input_scale)

1305

1306 - def prepare_static_weights_for_trtllm_fp4_moe(

1307 - self,

1308 - # args_dequant,

1309 - # args,

1310 - gemm1_weights,

1311 - gemm2_weights,

1312 - gemm1_scales_linear_fp4_bytes,

1313 - gemm2_scales_linear_fp4_bytes,

1314 - hidden_size,

1315 - intermediate_size,

1316 - num_experts,

1317 - ):

1318 - from flashinfer import nvfp4_block_scale_interleave

1319 - from flashinfer.fused_moe.core import (

1320 - _maybe_get_cached_w3_w1_permute_indices,

1321 - get_w2_permute_indices_with_cache,

1322 - )

1323 -

1324 - """Prepare quantized weights for kernel (done offline with weights)."""

1325 - epilogue_tile_m = 128 # FIXME: this depends on the kernel internals

1326 -

1327 - # Convert quantized weights to proper formats

1328 - gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(

1329 - num_experts, 2 * intermediate_size, hidden_size // 2

1330 - ) # packed fp4

1331 - gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(

1332 - torch.float8_e4m3fn

1333 - ).reshape(

1334 - num_experts, 2 * intermediate_size, hidden_size // 16

1335 - ) # fp8 scaling factors

1336 -

1337 - gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(

1338 - num_experts, hidden_size, intermediate_size // 2

1339 - ) # packed fp4

1340 - gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(

1341 - torch.float8_e4m3fn

1342 - ).reshape(

1343 - num_experts, hidden_size, intermediate_size // 16

1344 - ) # fp8 scaling factors

1345 -

1346 - gemm1_weights_fp4_shuffled = []

1347 - gemm1_scales_fp4_shuffled = []

1348 - gemm2_weights_fp4_shuffled = []

1349 - gemm2_scales_fp4_shuffled = []

1350 - for i in range(num_experts):

1351 - # Calculate the permute indices for the following:

1352 - # 1. Reorder rows of W1 and scales for fused gated activation

1353 - # 2. Shuffle weights and scaling factors for transposed mma output

1354 - # for both w3_w1 and w2 weights and scale factors

1355 - permute_indices = _maybe_get_cached_w3_w1_permute_indices(

1356 - self._cache_permute_indices,

1357 - gemm1_weights_fp4[i].view(torch.uint8),

1358 - epilogue_tile_m,

1359 - )

1360 - gemm1_weights_fp4_shuffled.append(

1361 - gemm1_weights_fp4[i]

1362 - .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]

1363 - .contiguous()

1364 - )

1365 -

1366 - permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(

1367 - self._cache_permute_indices,

1368 - gemm1_scales_linear_fp4[i].view(torch.uint8),

1369 - epilogue_tile_m,

1370 - num_elts_per_sf=16,

1371 - )

1372 - gemm1_scales_fp4_shuffled.append(

1373 - nvfp4_block_scale_interleave(

1374 - gemm1_scales_linear_fp4[i]

1375 - .view(torch.uint8)[

1376 - permute_sf_indices.to(gemm1_scales_linear_fp4.device)

1377 - ]

1378 - .contiguous()

1379 - )

1380 - )

1381 -

1382 - permute_indices = get_w2_permute_indices_with_cache(

1383 - self._cache_permute_indices,

1384 - gemm2_weights_fp4[i].view(torch.uint8),

1385 - epilogue_tile_m,

1386 - )

1387 - gemm2_weights_fp4_shuffled.append(

1388 - gemm2_weights_fp4[i]

1389 - .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]

1390 - .contiguous()

1391 - )

1392 -

1393 - permute_sf_indices = get_w2_permute_indices_with_cache(

1394 - self._cache_permute_indices,

1395 - gemm2_scales_linear_fp4[i].view(torch.uint8),

1396 - epilogue_tile_m,

1397 - num_elts_per_sf=16,

1398 - )

1399 - gemm2_scales_fp4_shuffled.append(

1400 - nvfp4_block_scale_interleave(

1401 - gemm2_scales_linear_fp4[i]

1402 - .view(torch.uint8)[

1403 - permute_sf_indices.to(gemm2_scales_linear_fp4.device)

1404 - ]

1405 - .contiguous()

1406 - )

1407 - )

1408 -

1409 - # Stack weights for all experts

1410 - gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)

1411 - gemm1_scales_fp4_shuffled = (

1412 - torch.stack(gemm1_scales_fp4_shuffled)

1413 - .view(torch.float8_e4m3fn)

1414 - .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)

1415 - )

1416 -

1417 - gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)

1418 - gemm2_scales_fp4_shuffled = (

1419 - torch.stack(gemm2_scales_fp4_shuffled)

1420 - .view(torch.float8_e4m3fn)

1421 - .reshape(num_experts, hidden_size, intermediate_size // 16)

1422 - )

1423 - return (

1424 - gemm1_weights_fp4_shuffled,

1425 - gemm1_scales_fp4_shuffled,

1426 - gemm2_weights_fp4_shuffled,

1427 - gemm2_scales_fp4_shuffled,

1428 - )

1429 -

1306 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

1307 # GEMM 1 processing

1308 gemm1_weight = layer.w13_weight.data

1309 gemm1_weight_scale = layer.w13_weight_scale.data

1310

1435 - if (

1436 - self.allow_flashinfer

1437 - and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1311 + if self.allow_flashinfer and (

1312 + self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1313 + or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1314 ):

1315 gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(

1316 gemm1_weight, gemm1_weight_scale, dim=-2

@@ -1508,7 +1384,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1384 gemm1_scales_fp4_shuffled,

1385 gemm2_weights_fp4_shuffled,

1386 gemm2_scales_fp4_shuffled,

1511 - ) = self.prepare_static_weights_for_trtllm_fp4_moe(

1387 + ) = prepare_static_weights_for_trtllm_fp4_moe(

1388 layer.w13_weight,

1389 layer.w2_weight,

1390 layer.w13_weight_scale,

@@ -1614,68 +1490,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1490 self.allow_flashinfer

1491 and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1492 ):

1617 - import flashinfer

1618 -

1619 - from vllm.model_executor.models.llama4 import Llama4MoE

1620 -

1621 - a1_gscale = layer.w13_input_scale_quant

1622 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = (

1623 - flashinfer.fp4_quantize(

1624 - x,

1625 - a1_gscale,

1626 - is_sf_swizzled_layout=False,

1627 - )

1628 - )

1629 - use_llama4_routing = (

1630 - custom_routing_function is Llama4MoE.custom_routing_function

1631 - )

1632 - routing_method_type = layer.routing_method_type

1633 - if use_llama4_routing:

1634 - routing_method_type = RoutingMethodType.Llama4

1635 - router_logits = (

1636 - router_logits.to(torch.float32)

1637 - if routing_method_type == RoutingMethodType.DeepSeekV3

1638 - else router_logits

1639 - )

1640 - routing_bias = e_score_correction_bias

1641 - if routing_bias is not None:

1642 - routing_bias = routing_bias.to(torch.bfloat16)

1643 - out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

1644 - routing_logits=router_logits,

1645 - routing_bias=routing_bias,

1646 - hidden_states=hidden_states_fp4,

1647 - hidden_states_scale=hidden_states_scale_linear_fp4.view(

1648 - torch.float8_e4m3fn

1649 - ).flatten(),

1650 - gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,

1651 - gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(

1652 - torch.float8_e4m3fn

1653 - ),

1654 - gemm1_bias=None,

1655 - gemm1_alpha=None,

1656 - gemm1_beta=None,

1657 - gemm1_clamp_limit=None,

1658 - gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,

1659 - gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(

1660 - torch.float8_e4m3fn

1661 - ),

1662 - gemm2_bias=None,

1663 - output1_scale_scalar=layer.g1_scale_c.data,

1664 - output1_scale_gate_scalar=layer.g1_alphas.data,

1665 - output2_scale_scalar=layer.g2_alphas.data,

1666 - num_experts=global_num_experts,

1493 + return flashinfer_trtllm_fp4_moe(

1494 + layer=layer,

1495 + x=x,

1496 + router_logits=router_logits,

1497 top_k=top_k,

1668 - n_group=num_expert_group,

1498 + global_num_experts=global_num_experts,

1499 + num_expert_group=num_expert_group,

1500 topk_group=topk_group,

1670 - intermediate_size=layer.intermediate_size_per_partition,

1671 - local_expert_offset=layer.ep_rank * layer.local_num_experts,

1672 - local_num_experts=layer.local_num_experts,

1673 - routed_scaling_factor=1.0,

1674 - tile_tokens_dim=None,

1675 - routing_method_type=routing_method_type,

1676 - do_finalize=True,

1677 - )[0]

1678 - return out

1501 + custom_routing_function=custom_routing_function,

1502 + e_score_correction_bias=e_score_correction_bias,

1503 + )

1504

1505 topk_weights, topk_ids, _ = FusedMoE.select_experts(

1506 hidden_states=x,

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+221

@@ -9,6 +9,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

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

10 FusedMoEConfig,

11 FusedMoEQuantConfig,

12 + RoutingMethodType,

13 )

14 from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (

15 FlashInferCuteDSLExperts,

@@ -110,3 +111,223 @@ def select_nvfp4_gemm_impl(

111 "CutlassExpertsFp4 doesn't support DP. Use flashinfer CUTLASS "

112 "Fused MoE backend instead (set VLLM_USE_FLASHINFER_MOE_FP4=1)"

113 )

114 +

115 +

116 +def prepare_static_weights_for_trtllm_fp4_moe(

117 + # args_dequant,

118 + # args,

119 + gemm1_weights,

120 + gemm2_weights,

121 + gemm1_scales_linear_fp4_bytes,

122 + gemm2_scales_linear_fp4_bytes,

123 + hidden_size,

124 + intermediate_size,

125 + num_experts,

126 +):

127 + from flashinfer import nvfp4_block_scale_interleave

128 + from flashinfer.fused_moe.core import (

129 + _maybe_get_cached_w3_w1_permute_indices,

130 + get_w2_permute_indices_with_cache,

131 + )

132 +

133 + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

134 + """Prepare quantized weights for kernel (done offline with weights)."""

135 + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals

136 +

137 + # Convert quantized weights to proper formats

138 + gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(

139 + num_experts, 2 * intermediate_size, hidden_size // 2

140 + ) # packed fp4

141 + gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(

142 + torch.float8_e4m3fn

143 + ).reshape(

144 + num_experts, 2 * intermediate_size, hidden_size // 16

145 + ) # fp8 scaling factors

146 +

147 + gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(

148 + num_experts, hidden_size, intermediate_size // 2

149 + ) # packed fp4

150 + gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view(

151 + torch.float8_e4m3fn

152 + ).reshape(num_experts, hidden_size, intermediate_size // 16) # fp8 scaling factors

153 +

154 + gemm1_weights_fp4_shuffled = []

155 + gemm1_scales_fp4_shuffled = []

156 + gemm2_weights_fp4_shuffled = []

157 + gemm2_scales_fp4_shuffled = []

158 + for i in range(num_experts):

159 + # Calculate the permute indices for the following:

160 + # 1. Reorder rows of W1 and scales for fused gated activation

161 + # 2. Shuffle weights and scaling factors for transposed mma output

162 + # for both w3_w1 and w2 weights and scale factors

163 + permute_indices = _maybe_get_cached_w3_w1_permute_indices(

164 + _cache_permute_indices,

165 + gemm1_weights_fp4[i].view(torch.uint8),

166 + epilogue_tile_m,

167 + )

168 + gemm1_weights_fp4_shuffled.append(

169 + gemm1_weights_fp4[i]

170 + .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]

171 + .contiguous()

172 + )

173 +

174 + permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(

175 + _cache_permute_indices,

176 + gemm1_scales_linear_fp4[i].view(torch.uint8),

177 + epilogue_tile_m,

178 + num_elts_per_sf=16,

179 + )

180 + gemm1_scales_fp4_shuffled.append(

181 + nvfp4_block_scale_interleave(

182 + gemm1_scales_linear_fp4[i]

183 + .view(torch.uint8)[

184 + permute_sf_indices.to(gemm1_scales_linear_fp4.device)

185 + ]

186 + .contiguous()

187 + )

188 + )

189 +

190 + permute_indices = get_w2_permute_indices_with_cache(

191 + _cache_permute_indices,

192 + gemm2_weights_fp4[i].view(torch.uint8),

193 + epilogue_tile_m,

194 + )

195 + gemm2_weights_fp4_shuffled.append(

196 + gemm2_weights_fp4[i]

197 + .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]

198 + .contiguous()

199 + )

200 +

201 + permute_sf_indices = get_w2_permute_indices_with_cache(

202 + _cache_permute_indices,

203 + gemm2_scales_linear_fp4[i].view(torch.uint8),

204 + epilogue_tile_m,

205 + num_elts_per_sf=16,

206 + )

207 + gemm2_scales_fp4_shuffled.append(

208 + nvfp4_block_scale_interleave(

209 + gemm2_scales_linear_fp4[i]

210 + .view(torch.uint8)[

211 + permute_sf_indices.to(gemm2_scales_linear_fp4.device)

212 + ]

213 + .contiguous()

214 + )

215 + )

216 +

217 + # Stack weights for all experts

218 + gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)

219 + gemm1_scales_fp4_shuffled = (

220 + torch.stack(gemm1_scales_fp4_shuffled)

221 + .view(torch.float8_e4m3fn)

222 + .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)

223 + )

224 +

225 + gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)

226 + gemm2_scales_fp4_shuffled = (

227 + torch.stack(gemm2_scales_fp4_shuffled)

228 + .view(torch.float8_e4m3fn)

229 + .reshape(num_experts, hidden_size, intermediate_size // 16)

230 + )

231 + return (

232 + gemm1_weights_fp4_shuffled,

233 + gemm1_scales_fp4_shuffled,

234 + gemm2_weights_fp4_shuffled,

235 + gemm2_scales_fp4_shuffled,

236 + )

237 +

238 +

239 +def flashinfer_trtllm_fp4_moe(

240 + layer: torch.nn.Module,

241 + x: torch.Tensor,

242 + router_logits: torch.Tensor,

243 + top_k: int,

244 + global_num_experts: int,

245 + num_expert_group: int | None,

246 + topk_group: int | None,

247 + custom_routing_function: object | None,

248 + e_score_correction_bias: torch.Tensor | None,

249 +) -> torch.Tensor:

250 + """

251 + Apply FlashInfer TensorRT-LLM FP4 MoE kernel.

252 +

253 + Args:

254 + layer: The MoE layer with weights and scales

255 + x: Input tensor

256 + router_logits: Router logits for expert selection

257 + top_k: Number of experts to select per token

258 + global_num_experts: Total number of experts across all ranks

259 + num_expert_group: Number of expert groups (for grouped routing)

260 + topk_group: Top-k within each group

261 + custom_routing_function: Custom routing function (e.g., Llama4)

262 + e_score_correction_bias: Optional routing bias correction

263 +

264 + Returns:

265 + Output tensor from the MoE layer

266 + """

267 + import flashinfer

268 +

269 + from vllm.model_executor.models.llama4 import Llama4MoE

270 +

271 + # Quantize input to FP4

272 + a1_gscale = layer.w13_input_scale_quant

273 + (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

274 + x,

275 + a1_gscale,

276 + is_sf_swizzled_layout=False,

277 + )

278 +

279 + # Determine routing method type

280 + use_llama4_routing = custom_routing_function is Llama4MoE.custom_routing_function

281 + routing_method_type = layer.routing_method_type

282 + if use_llama4_routing:

283 + routing_method_type = flashinfer.RoutingMethodType.Llama4

284 +

285 + # Prepare routing bias

286 + routing_bias = e_score_correction_bias

287 + if routing_bias is not None:

288 + routing_bias = routing_bias.to(torch.bfloat16)

289 +

290 + router_logits = (

291 + router_logits.to(torch.float32)

292 + if routing_method_type == RoutingMethodType.DeepSeekV3

293 + else router_logits

294 + )

295 +

296 + # Call TRT-LLM FP4 block-scale MoE kernel

297 + out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

298 + routing_logits=router_logits,

299 + routing_bias=routing_bias,

300 + hidden_states=hidden_states_fp4,

301 + hidden_states_scale=hidden_states_scale_linear_fp4.view(

302 + torch.float8_e4m3fn

303 + ).flatten(),

304 + gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,

305 + gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(

306 + torch.float8_e4m3fn

307 + ),

308 + gemm1_bias=None,

309 + gemm1_alpha=None,

310 + gemm1_beta=None,

311 + gemm1_clamp_limit=None,

312 + gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,

313 + gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(

314 + torch.float8_e4m3fn

315 + ),

316 + gemm2_bias=None,

317 + output1_scale_scalar=layer.g1_scale_c.data,

318 + output1_scale_gate_scalar=layer.g1_alphas.data,

319 + output2_scale_scalar=layer.g2_alphas.data,

320 + num_experts=global_num_experts,

321 + top_k=top_k,

322 + n_group=num_expert_group if num_expert_group is not None else 0,

323 + topk_group=topk_group if topk_group is not None else 0,

324 + intermediate_size=layer.intermediate_size_per_partition,

325 + local_expert_offset=layer.ep_rank * layer.local_num_experts,

326 + local_num_experts=layer.local_num_experts,

327 + routed_scaling_factor=None,

328 + tile_tokens_dim=None,

329 + routing_method_type=routing_method_type,

330 + do_finalize=True,

331 + )[0]

332 +

333 + return out