[Kernel] Pack topk id/weights triton kernel (#42527)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

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),
         }

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)