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 69e5b7fe4f..7308ed195c 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 @@ -343,6 +343,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula # the TRTLLM C++ kernel supports. "routing_method_type": RoutingMethodType.Renormalize, "do_finalize": True, + "enable_pdl": True, "output": output, "tune_max_num_tokens": max(self.max_capture_size, 1), } diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b37d17bdac..b254dc7b7a 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -376,17 +376,62 @@ def disable_inplace() -> bool: return is_torch_equal_or_newer("2.9") -@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) +@triton.jit +def _pack_topk_ids_weights_kernel( + topk_ids_ptr, + topk_weights_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + USE_GDC: tl.constexpr, + launch_pdl: tl.constexpr, # triton metadata +): + pid = tl.program_id(axis=0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + if USE_GDC: + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32) + expert_id_shifted = expert_id << 16 + + weight = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0) + weight_bf16 = weight.to(tl.bfloat16) + weight_int16 = weight_bf16.to(tl.int16, bitcast=True) + + weight_int32 = weight_int16.to(tl.int32) & 0xFFFF + + packed = expert_id_shifted | weight_int32 + tl.store(output_ptr + offsets, packed, mask=mask) + + def trtllm_moe_pack_topk_ids_weights( - topk_ids: torch.Tensor, topk_weights: torch.Tensor + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + block_size: int = 1024, ) -> torch.Tensor: - """ - Pack topk_ids and topk_weights into a single int32 tensor. - Format: (expert_id << 16) | weight_bf16.view(int16) - """ - return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view( - torch.int16 + 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)