diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 6b29a6f69d..6262234679 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( ) from vllm.model_executor.layers.fused_moe.utils import ( fi_moe_largest_bucket, - trtllm_moe_pack_topk_ids_weights, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, @@ -199,11 +198,10 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular) import flashinfer from flashinfer.fused_moe import WeightLayout - # Pack topk ids and weights into format expected by the TRTLLM kernel. - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + topk_ids = topk_ids.to(dtype=torch.int32) result = flashinfer.fused_moe.trtllm_bf16_routed_moe( - topk_ids=packed_topk_ids, + topk_ids=(topk_ids, topk_weights), hidden_states=hidden_states, gemm1_weights=w1, gemm2_weights=w2, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 4cb6076c81..709ba17d42 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -17,7 +17,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( ) from vllm.model_executor.layers.fused_moe.utils import ( fi_moe_largest_bucket, - trtllm_moe_pack_topk_ids_weights, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, @@ -273,8 +272,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): import flashinfer from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout - # Pack topk ids and weights into format expected by the kernel. - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + topk_ids = topk_ids.to(dtype=torch.int32) if a1q_scale is None: raise RuntimeError( @@ -295,7 +293,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( - topk_ids=packed_topk_ids, + topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=hidden_states, hidden_states_scale=hidden_states_scale, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py index 1a947756ce..a3d50011d9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py @@ -38,9 +38,6 @@ from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import ( - trtllm_moe_pack_topk_ids_weights, -) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe @@ -206,7 +203,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular): hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - packed_topk_ids: torch.Tensor, + topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor], gemm1_lora_delta: torch.Tensor | None, global_num_experts: int, a1q_scale: torch.Tensor | None, @@ -250,10 +247,6 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular): intermediate_size = self.intermediate_size_per_partition K = output.size(1) - # Routing is computed outside the MoE; pack it into the - # (eid<<16)|w.bf16 format the routed API expects. - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) - # ---- Base-model fast path ---- # When no token in the batch selects a LoRA adapter, skip the LoRA machinery # and run the plain base MoE with do_finalize=True, which writes the finalized @@ -263,7 +256,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular): hidden_states=hidden_states, w1=w1, w2=w2, - packed_topk_ids=packed_topk_ids, + topk_ids_and_weights=(topk_ids, topk_weights), gemm1_lora_delta=None, # without LoRA, no delta global_num_experts=global_num_experts, a1q_scale=a1q_scale, @@ -337,7 +330,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular): hidden_states=hidden_states, w1=w1, w2=w2, - packed_topk_ids=packed_topk_ids, + topk_ids_and_weights=(topk_ids, topk_weights), gemm1_lora_delta=gemm1_lora_delta, global_num_experts=global_num_experts, a1q_scale=a1q_scale, @@ -518,7 +511,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase): hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - packed_topk_ids: torch.Tensor, + topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor], gemm1_lora_delta: torch.Tensor | None, global_num_experts: int, a1q_scale: torch.Tensor | None, @@ -533,7 +526,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase): # the caller's buffer via output= so it finalizes in place -- no copy. do_finalize = gemm1_lora_delta is None ret = flashinfer.fused_moe.trtllm_bf16_routed_moe( - topk_ids=packed_topk_ids, + topk_ids=topk_ids_and_weights, hidden_states=hidden_states, gemm1_weights=w1, gemm2_weights=w2, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index e28e401e4c..159c77848a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( ) from vllm.model_executor.layers.fused_moe.utils import ( fi_moe_largest_bucket, - trtllm_moe_pack_topk_ids_weights, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, @@ -363,9 +362,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula ) -> None: from flashinfer import trtllm_fp4_block_scale_routed_moe - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_tensor, + topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=x_quant, hidden_states_scale=x_scale, @@ -417,6 +415,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): + topk_ids = topk_ids.to(dtype=torch.int32) + topk = topk_ids.size(-1) local_num_experts = w1.size(0) local_expert_offset = self.moe_config.ep_rank * local_num_experts diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 688465092c..b5d5f3b96a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -20,10 +20,7 @@ from vllm.model_executor.layers.fused_moe.moe_output import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import ( - fi_moe_largest_bucket, - trtllm_moe_pack_topk_ids_weights, -) +from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, has_flashinfer_situ_activation, @@ -369,14 +366,11 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula ) else: block_scale, per_token_scale = a1q_scale, None - - # Pack topk ids and weights into format expected by the kernel. - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) output1_scale_gate_scalar = self.quant_config.g1_alphas # Invoke kernel. flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( - topk_ids=packed_tensor, + topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=hidden_states, hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape( @@ -434,6 +428,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula # Per-token defers input quant to _invoke_kernel, so a1q_scale is None. assert a1q_scale is not None or self.per_token_activation + # DeepEP produces int64 indexes. + topk_ids = topk_ids.to(dtype=torch.int32) + M = hidden_states.shape[0] chunk_size = self._get_chunk_size() diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 915ea1dfb4..48e217b1a9 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -503,35 +503,6 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int: return max(moe_config.max_num_tokens * moe_config.dp_size, 8192) -def trtllm_moe_pack_topk_ids_weights( - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - block_size: int = 1024, -) -> torch.Tensor: - assert topk_ids.shape == topk_weights.shape - assert topk_ids.is_contiguous() and topk_weights.is_contiguous() - - original_shape = topk_ids.shape - ids_flat = topk_ids.reshape(-1) - weights_flat = topk_weights.reshape(-1) - - n_elements = ids_flat.numel() - output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device) - - use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90) - grid = (triton.cdiv(n_elements, block_size),) - _pack_topk_ids_weights_kernel[grid]( - ids_flat, - weights_flat, - output, - n_elements, - BLOCK_SIZE=block_size, - USE_GDC=use_gdc, - launch_pdl=use_gdc, - ) - return output.reshape(original_shape) - - @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _swiglu_limit_torch( output: torch.Tensor,