执行摘要

• 一句话：Triton内核打包topk id/权重，支持GDC
• 推荐动作：此PR对于了解Triton内核封装和GDC/PDL在vLLM中的应用有参考价值。建议关注MoE路径性能回归测试。整体改动小，可快速合并。

功能与动机

原始的torch.compile实现无法利用GDC/PDL且性能不足。通过手写Triton内核可以获得更好的控制，并启用Blackwell GPU上的高级调度特性，从而降低MoE路由的延迟。

实现拆解

1. 新增Triton内核：在vllm/model_executor/layers/fused_moe/utils.py中编写_pack_topk_ids_weights_kernel，使用Triton JIT编译，将专家ID和权重打包为单个int32张量。利用GDC（全局依赖链）和PDL（程序描述符列表）在支持的硬件上优化执行顺序。
2. 重写打包函数：将trtllm_moe_pack_topk_ids_weights从@torch.compile装饰器改为调用上述内核，并增加block_size参数和连续性断言。动态检测CUDA计算能力以决定是否启用GDC/PDL（目前针对SM 90+）。
3. 启用PDL：在vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py的apply方法中向kwargs添加"enable_pdl": True，使flashinfer的routed MoE内核也能受益于PDL。
   此PR未包含直接测试文件，但内核通过已有的MoE测试路径间接验证。

关键文件：

• vllm/model_executor/layers/fused_moe/utils.py（模块 MoE工具；类别 source；类型 core-logic；符号 _pack_topk_ids_weights_kernel, trtllm_moe_pack_topk_ids_weights）: 核心变更文件：用Triton内核替换torch.compile打包，引入GDC/PDL支持。
• vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py（模块 MoE专家；类别 source；类型 configuration）: 为flashinfer的routed MoE调用添加enable_pdl=True，以匹配内核GDC/PDL支持。

关键符号：_pack_topk_ids_weights_kernel, trtllm_moe_pack_topk_ids_weights

关键源码片段

vllm/model_executor/layers/fused_moe/utils.py

核心变更文件：用Triton内核替换torch.compile打包，引入GDC/PDL支持。

@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, not used directly
):
    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()
​
    # 加载 topk_ids 并左移 16 位
    expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32)
    expert_id_shifted = expert_id << 16
​
    # 加载 topk_weights，转为 bfloat16，再 bitcast 为 int16
    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)
​
    # 转为 int32 并掩码低 16 位，然后与 expert_id 按位或
    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,
    block_size: int = 1024,
) -> torch.Tensor:
    """将 topk_ids 和 topk_weights 打包成单个 int32 张量。
    格式: (expert_id << 16) | weight_bf16.view(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)
​
    # 仅在 CUDA 且计算能力 >= 90 时启用 GDC/PDL
    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)

评论区精华

Gemini Code Assist 自动生成两条评论：

• 建议移除注释掉的旧代码（但最终patch中已无残留，作者标记为Done）。

• 建议将GDC/PDL的能力检查从SM 90改为SM 100，作者回复“you are wrong”并保持原样，表明SM 90同样支持或Triton已正确处理。
  最终由zyongye审批通过。

• GDC能力检查应使用SM 100而非SM 90 (correctness): 开发者回复'you are wrong'并保留SM 90，表明SM 90同样支持或Triton已正确处理。

风险与影响

• 风险：如果GDC/PDL在不支持的硬件上启用，可能引发运行时错误。但作者坚持SM 90检测正确，且Triton的GDC接口在SM 90+上可用。此外，Triton内核仅在CUDA平台生效，非NVIDIA GPU使用use_gdc=False，兼容性良好。输出格式与旧实现一致（int32打包），回归风险低。缺少直接单元测试，依赖集成测试覆盖。
• 影响：影响所有使用trtllm_moe_pack_topk_ids_weights的MoE模型（如DeepSeek、Mixtral等），预期降低路由开销。同时trtllm_mxfp4_moe受益于PDL启用。无外部接口变化。
• 风险标记：缺少测试覆盖, GDC/PDL兼容性依赖硬件

关联脉络

• PR #42497 [Perf] Wire silu_and_mul_per_block_quant into TritonFP8MoE (MiniMax-M2): 同样针对MoE内核路径的性能优化，使用Triton内核融合操作。
• PR #40131 [Bugfix] moe lora align kernel grid: MoE内核错误修复，相同功能领域。