跳到正文
PyTorch Blog· Sean Chen (Red Hat) and Shangdi Yu (PyTorch, Meta Platforms)·· 15 天前AI 评分58

用 Helion 为 vLLM 构建高性能可移植线性后端

Building a High-Performance and Portable vLLM Linear Backend with Helion

AI 导读

PyTorch 团队将 Helion 集成进 vLLM 的线性后端，用单一 GEMM 实现覆盖 Standard GEMM、Split-K 和 Swap-AB 三种变体，由按形状的自动调优选择最佳变体与配置。

正文 · AI 翻译

TL;DR

我们将 Helion 集成到 vLLM 的 线性后端 中，以探索自动调优的高层内核 DSL 如何在降低内核实现复杂度的同时提升 LLM 推理性能。单个 Helion 通用矩阵乘法（GEMM）实现即可覆盖多种算法变体，包括 Standard GEMM、Split-K 和 Swap-AB，并通过按形状自动调优自动选择最佳变体和配置。

在 NVIDIA Hopper GPU 上，Helion 线性后端将按形状调优与混合调度相结合，在评估的模型上均优于 vLLM 默认的 CUTLASS 和 DeepGEMM 后端，带来一致的端到端性能提升，部分工作负载吞吐量提升超过 10%。

vLLM 与 Helion 简要背景

vLLM 是一个面向大语言模型（LLM）的高性能推理与服务框架。对于量化线性层，如 FP8、INT8、INT4 和 NVFP4，vLLM 提供了专门的 线性后端，集成了来自 CUTLASS、DeepGEMM 和 FlashInfer 等库的硬件优化内核实现。这些后端使 vLLM 能够针对不同的量化格式和硬件平台使用优化内核。

Helion 是一个 PyTorch 原生的硬件无关内核 DSL，旨在使用 tile 编程模型编写高性能内核。其高层抽象允许开发者用简洁的 Pythonic 代码编写一次内核，而 Helion 会为不同硬件平台上的不同工作负载生成专用代码。Helion 无需为每个工作负载和硬件目标手动特化内核，而是依赖提前（AOT）自动调优器系统地探索广阔的搜索空间，从底层内存布局和内核调度到高层算法选择，并选出性能最佳的配置。这种高层抽象与系统化自动调优的结合，使单个内核实现能够在多样化的负载和硬件上得到优化。Helion 还支持 LLM 引导的搜索，以进一步提升调优效率。

Helion 的应用机会与挑战

本节总结了在 vLLM 等推理引擎中采用 Helion 内核的关键机会与挑战，从高层概述了 Helion 的价值主张与权衡。

机会

性能：我们 之前的工作 展示了 Helion 通过细粒度调优，在多样化的工作负载模式和硬件平台上实现推理内核 SOTA 性能的潜力。 

系统化调优框架：与更开放的智能体式方法（迭代生成、分析并改进内核实现）相比，Helion 将内核调优形式化为定义明确的搜索空间上的结构化数值优化问题。这使调优过程更加稳健可靠，同时在调优效率上保持竞争力，且不依赖性能分析来指导搜索。Helion 进一步支持 LLM 引导的搜索，将 LLM 的推理能力与系统化数值优化相结合，以高效地探索搜索空间。

可移植性与抽象：Helion 不仅仅是一个可移植的 DSL。它可以在针对不同工作负载模式和硬件平台进行优化的同时，保持单一内核实现。 

客户端内核优化：vLLM 等推理引擎中的默认内核通常针对通用工作负载和常用模型进行优化。借助 Helion，用户无需专门的内核专业知识或大量工程投入，即可针对其特定部署进一步调优内核性能。

挑战

Helion 在集成到推理引擎方面仍面临一些挑战，Helion 团队正在持续努力解决和缓解这些问题。

提前内核调优开销：Helion 将内核调优自动化和系统化，但细粒度调优仍可能耗时数小时。这一开销主要来自调优策略的粒度，而非 Helion 本身。若使用其他内核 DSL 实现同等程度的工作负载特定优化，将产生相当——甚至可能更高——的调优成本。

启动时开销：vLLM 启动期间的 CUDA Graph 捕获会触发 Helion JIT 编译，从而增加冷启动延迟。通过缓存编译产物，这一开销在热启动时基本可以消除。

推理运行时开销：在 CUDA Graph 捕获范围之外，推理期间 Helion 内核的分发和启动可能引入额外的 CPU 开销，可能抵消细粒度调优带来的性能收益。在实践中，Helion 在 CUDA Graphs 下执行时最为有效，可避免这一 CPU 开销。

维护开销：为热门模型提供预调优配置会带来持续的上游维护负担。大型配置文件难以维护，且通过单元测试和 CI 进行详尽验证也不切实际。 

权衡三角

使用 Helion 进行细粒度内核调优需要在性能、易用性和可维护性之间进行权衡。这一权衡并非 Helion 所特有，而是工作负载特定内核优化的固有特性：优化其中两项往往以牺牲第三项为代价。

图 1：性能-易用性-可维护性权衡三角

更高的性能通常需要更细粒度的配置调优。这要么增加客户端自动调优开销，要么需要维护者在上游提供并维护更多预调优配置。因此，目标是为目标用例找到恰当的平衡。 

Helion 线性后端 

范围

本工作为 vLLM 添加了 Helion 线性后端，并聚焦于使用 Helion 的 Triton 后端的 NVIDIA Hopper GPU。FP8 和 INT8 是在 Hopper 上实现高效推理的主要量化格式，因此我们针对以下格式的量化 GEMM：

  • FP8_Dynamic：FP8 逐 token 激活和逐通道权重缩放。
  • W8A8_INT8：INT8 逐 token 激活和逐通道权重缩放。
  • Block_FP8：FP8 采用 1×128 激活缩放和 128×128 权重缩放。 

初步结果 显示，使用 Helion 的 CuteDSL 后端在 NVIDIA Blackwell GPU 上取得了具有竞争力的 GEMM 性能。随着 CuteDSL 后端逐渐成熟，这项工作可以扩展以支持 Blackwell。

内核实现

若干 GEMM 算法变体可以针对特定输入形状提升性能。具体来说：

  • Split-K：将 K 维度划分到多个线程块上，以在 M 和/或 N 维度太小时增加并行度。
  • Swap-AB：将 A@B 重写为 (B.T@A.T).T，通过启用更有利的分块和更好的 GPU 利用率，来提升 M 维度较小的形状的性能。

传统上，内核作者需要实现多个 GEMM 变体，将标准 GEMM 与这些专用变体进行基准测试，并开发启发式规则来根据输入形状决定分派到哪个实现。例如，vLLM 当前的 Block_FP8 线性后端在 Hopper 上当 M < 32 时会分派到 Swap-AB 变体。这类启发式规则通常被设计为在一般情况下表现良好，但可能并非对每个模型或工作负载都是最优的。 

借助 Helion，单个统一的 GEMM 实现即可覆盖全部三种变体——Standard、Split-K 和 Swap-AB。无需实现单独的内核并手动设计分派启发式规则，算法选择被暴露为可调参数。随后，AOT 自动调优器可以为每个工作负载自动选择最佳变体以及底层内核配置。

为简单起见，我们下面使用一个基础的矩阵乘法内核来说明该方法。本工作中使用的量化 GEMM 内核遵循相同的结构，只是增加了额外的量化和缩放逻辑。

def matmul(
    out: torch.Tensor,  # [M, N]
    a: torch.Tensor,    # [M, K]
    b: torch.Tensor,    # [K, N]
) -> None:
    M, K = a.shape
    N = b.shape[1]
    hl.specialize(K)
    hl.specialize(N)

    out_dtype = out.dtype
    acc_dtype = torch.float32

    split_k = hl.register_tunable(
        "split_k", PowerOfTwoFragment(1, 256)
    )
    k_block_size = helion.next_power_of_2(helion.cdiv(K, split_k))
    if split_k > 1:
        out.zero_()

    swap_ab = hl.register_tunable(
        "swap_ab", BooleanFragment()
    )

    for tile_m, tile_n, outer_k in hl.tile(
        [M, N, K],
        block_size=[None, None, k_block_size]
    ):
        acc = hl.zeros([tile_m, tile_n], acc_dtype)
        acc_t = acc.t()
        for tile_k in hl.tile(outer_k.begin, outer_k.end):
            if swap_ab:
                a_blk = hl.load(a, [tile_m.index[None, :], tile_k.index[:, None]])
                b_blk = hl.load(b, [tile_k.index[None, :], tile_n.index[:, None]])
                acc_t = hl.dot(
                    b_blk,
                    a_blk,
                    acc=acc_t,
                    out_dtype=acc_dtype,
                )
            else:
                acc = hl.dot(
                    a[tile_m, tile_k],
                    b[tile_k, tile_n],
                    acc=acc,
                    out_dtype=acc_dtype,
                )

        if swap_ab:
            out_blk = acc_t.t().to(out_dtype)
        else:
            out_blk = acc.to(out_dtype)

        if split_k == 1:
            out[tile_m, tile_n] = out_blk
        else:
            hl.atomic_add(out, [tile_m, tile_n], out_blk)

该实现将两种算法选择都暴露为可调参数：split_k 被约束为最大 256 的 2 的幂值，而 swap_ab 是一个布尔参数。自动调优器会探索受约束的搜索空间，对不同的组合进行基准测试，并为每个工作负载选择最佳变体。

混合分派

图 2：基于运行时 num_tokens 和 CUDA Graph 覆盖范围的 Helion 线性后端混合分派策略。

我们采用基于运行时 num_tokens 和 CUDA Graph 覆盖范围的混合分派策略。对于最大到 max_helion_size 的小形状，线性后端在 CUDA Graph 重放时分派到 Helion。对于超过 max_helion_size 的较大形状，则回退到默认内核（CUTLASS/DeepGEMM）。

这种混合分派策略带来三个实际好处：

  • 消除 Helion 运行时开销。Helion 内核仅通过 CUDA Graph 重放执行，避免了内核分派和启动带来的额外 CPU 开销。
  • 减少内核调优开销。细粒度的 Helion 调优仅限于主导解码工作负载的小 num_tokens 范围，显著减少了需要调优的输入形状数量，同时保留了获得有意义端到端性能提升的机会。
  • 减少配置维护开销。更小的已调优配置集合使预调优配置更易于验证、发布和长期维护。

内核自动调优

我们使用 vLLM 中提供的实用脚本对 Helion 线性内核进行自动调优：

HELION_AUTOTUNER=LLMSeededLFBOTreeSearch \
HELION_BENCHMARK_CUDAGRAPH =1 \
python scripts/autotune_helion_kernels.py \
  --kernels scaled_mm block_scaled_mm \
  --autotune-effort "full"

以下各节描述了本工作中使用的关键调优策略和设置。

按形状配置调优

为了与 CUTLASS、DeepGEMM 和 FlashInfer 等高度优化的 GEMM 库竞争，我们在 Helion 调度范围内针对每个输入形状单独调优 Helion 内核。

如混合调度策略所述，我们为本工作设置 max_helion_size = 32。vLLM 为以下 num_tokens 值捕获 CUDA Graph：

[1, 2, 4] + range(8, 256, 8) + range(256, max_graph_size + 1, 16)

在 max_helion_size = 32 下，Helion 内核因此针对以下情况进行自动调优：

num_tokens = [1, 2, 4, 8, 16, 24, 32]

每个 num_tokens 值都针对模型所使用的相应 GEMM 形状单独调优。

为自动调优器基准测试启用 CUDA Graph

由于内核调度和启动带来的额外 CPU 开销，Helion 内核仅在 CUDA Graph 执行下使用。因此，启用 CUDA Graph 进行基准测试使自动调优器能够在更接近实际推理执行的条件下评估配置。

Helion 通过 HELION_BENCHMARK_CUDAGRAPH环境变量公开此功能，并在本工作中开启。

使用 LLM 引导搜索

本工作中的内核配置使用 LLMSeededLFBOTreeSearch 自动调优器生成，该调优器使用 LLM 识别一组有前景的候选配置，作为后续 LFBOTreeSearch 的种子。在本工作中，我们使用 Claude Opus 4.8。

从更高质量的候选配置开始搜索有助于自动调优器发现性能更好的配置。它还可能通过将搜索引向搜索空间中更有前景的区域来减少整体自动调优时间。

性能评估

我们在内核和端到端服务两个层面评估 Helion 线性后端，涵盖一系列稠密模型和量化格式。为了解性能如何随模型规模扩展，我们对 Qwen3-1.7B、Qwen3-4B、Qwen3-8B、Qwen3-14B 和 Qwen3-32B 进行基准测试。我们还加入了最新的 Qwen3.8-27B 模型，以在更新的模型架构上评估该后端。

对于每个模型，我们在 NVIDIA Hopper GPU 上评估三种常用的 8 位量化格式。例如，对于 Qwen3.8-27B，我们基准测试： 

所有基准测试均在 NVIDIA H100 80GB HBM3 GPU 上执行。

内核级评估

我们首先单独评估 Helion 量化 GEMM 内核，以了解细粒度调优独立于推理栈其余部分的性能收益。每个 Helion 内核都与 Hopper 上默认 vLLM 线性后端所使用的相应内核进行比较：

  • FP8_Dynamic：Helion 对比 CUTLASS
  • W8A8_INT8：Helion 对比 CUTLASS
  • Block_FP8：Helion 对比 FlashInfer/DeepGEMM

图 3：在所有端到端基准测试模型使用的输入形状上，Helion 量化 GEMM 内核相对于默认 vLLM 内核库的加速分布。菱形表示几何平均加速比。

Helion 在 FP8_Dynamic 上相对 CUTLASS 取得几何平均 1.110× 加速，W8A8_INT8 相对 CUTLASS 为 1.178×，Block_FP8 相对 FlashInfer 为 1.149×，Block_FP8 相对 DeepGEMM 为 1.177×。分布还表明，性能在不同输入形状间存在差异，凸显了细粒度调优的价值，而非依赖单一 kernel 配置或通用调度启发式。

端到端评估

接下来我们评估 kernel 层面的改进是否能转化为端到端服务性能。

服务器设置

我们使用以下命令启动 vLLM 服务器：

vllm serve \
    --model "$MODEL" \
    --max-num-seqs 32 \
    --tensor-parallel-size 1 \
    --no-enable-prefix-caching \
    --linear-backend helion

相关选项如下：

  • --max-num-seqs 32：Helion linear 后端目前使用 32 num_tokens 的混合调度阈值。因此我们将评估重点放在最大 32 的 batch size 上，此时 Helion kernel 处于活跃状态，可以直接影响端到端性能。
  • --no-enable-prefix-caching：有意禁用前缀缓存，以避免其对基准测试结果的影响。
  • --linear-backend helion：启用 Helion linear 后端。省略此标志则使用默认 linear 后端。

基准测试设置

我们使用 ShareGPT 数据集运行端到端服务基准测试，配置如下：

vllm bench serve \
    --backend vllm \
    --model "${MODEL}" \
    --endpoint /v1/completions \
    --dataset-name sharegpt \
    --dataset-path "${DATASET}" \
    --max-concurrency "${BATCH_SIZE}" \
    --num-warmups "${NUM_WARMUPS}" \
    --num-prompts "${PROMPTS}" \
    --ignore-eos

每个工作负载都分别使用默认和 Helion linear 后端进行基准测试。在 Hopper 上用作基线的默认 vLLM linear 后端为：

  • FP8_Dynamic：CutlassFP8ScaledMMLinearKernel
  • W8A8_INT8：CutlassInt8ScaledMMLinearKernel
  • Block_FP8：FlashInferFp8DeepGEMMDynamicBlockScaledKernel

端到端基准测试结果

下图展示了 Helion linear 后端在不同模型规模、batch size 和量化格式下，相对相应默认后端的端到端吞吐量加速。

图 4：Helion linear 后端端到端吞吐量加速。

总体而言，Helion linear 后端在所评估的模型和量化格式上带来了一致的端到端性能提升，部分工作负载的吞吐量提升超过 10%。

迈向实用的采用模式

本工作提出的 Helion linear 后端目前已在我们的 vLLM fork 中提供，并已准备好投入生产使用。该 fork 还包含自动调优工具和说明，供最终用户为更多模型和工作负载生成优化配置。上游采用仍面临的一个挑战是，发布和验证大量预调优配置所带来的维护开销。

对于诸如本工作所研究的 GEMM 内核这类延迟敏感型内核，我们正在探索一种模式：Helion 内核与集成框架在上游维护，并带有用于功能测试和 CI 的默认配置，而针对特定工作负载的自动调优则交由最终用户完成。本工作中使用的同一套自动化工具允许用户在其目标模型和硬件上生成优化配置，然后再进行部署。这种模式更看重性能和可维护性，而非开箱即用的易用性。我们认为，对于延迟敏感型内核而言，这种权衡是切实可行的，因为性能收益足以证明额外的离线优化工作是值得的，尤其是对于愿意在生产部署前进行这项投入的模型服务提供商而言。

然而，细粒度调优并不是 Helion 内核唯一实用的采用策略。对于较小的辅助内核，例如量化、激活和归一化内核，宁可牺牲一些峰值性能，换用一组规模小得多、能跨不同形状通用的配置，可能更为可取。我们的初步实验表明，对于这些内核，仅用六个配置就能实现有意义的性能提升。这为辅助内核的上游采用提供了另一条路径，同时大幅降低了调优和维护开销。

我们鼓励感兴趣的用户试用当前实现，并在vLLM Helion 线性后端 RFC中分享他们的经验。来自真实部署的反馈将帮助我们评估这些方法，并引导集成工作走向更广泛的采用。

未来工作

未来工作将聚焦于扩展 Helion 内核在各类推理工作负载和硬件平台上的覆盖范围。

硬件覆盖。Helion 团队正在继续扩展和优化跨硬件平台的后端支持，包括面向 NVIDIA Blackwell GPU 的 CuteDSL 后端，以及针对 AMD GPU 和 TPU 的持续性能优化工作。随着这些后端逐渐成熟，我们计划在这些硬件目标上扩展并重新评估 Helion 线性后端。

模型覆盖。本工作主要针对稠密模型，其中线性后端占推理延迟的很大一部分。对于 MoE 模型，MoE 后端成为更重要的优化目标。我们正在开发 Helion MoE 后端，以将 Helion 的内核优化能力带给这些工作负载，并扩大模型覆盖范围。

结论

Helion 的高层抽象使得用单一内核实现来表达和维护，同时针对不同工作负载和硬件目标进行优化成为可能。正如本工作中的量化 GEMM 内核所展示的那样，即使是 Standard GEMM、Split-K 和 Swap-AB 等算法变体，也可以统一到一个实现中，并作为可调选项暴露给自动调优器。结合按形状的细粒度内核调优和混合调度，Helion 线性后端在 Hopper GPU 上于所评估的模型中均优于默认的 CUTLASS、DeepGEMM 和 FlashInfer 后端，带来一致的端到端性能提升，某些工作负载的吞吐量提升超过 10%。

然而，这些性能提升伴随着权衡。针对延迟关键内核的细粒度调优可提升性能，但会增加 AOT 调优工作量；而提供预调优配置则提升了开箱即用性，代价是持续的维护成本。这反映了本文通篇讨论的更广泛的性能-可用性-可维护性权衡。

目前，我们在 vLLM fork 中提供预调优配置。对于上游采用，我们正在探索一种不同的模式：在上游维护 Helion 内核和集成框架，同时将特定工作负载的自动调优委托给最终用户。这更倾向于性能和可维护性，而非开箱即用性，但 Helion 的自动调优框架使额外的优化步骤对性能敏感型部署而言变得切实可行。

致谢

这项工作得到了 Red Hat 的 OCTO 和 vLLM 团队以及 Meta 的 Helion 团队众多贡献者的支持。我们特别感谢我们的同事：Richard Zou 和 Jongsok Choi 在整个工作过程中提供的反馈和支持。 

文章 使用 Helion 构建高性能且可移植的 vLLM 线性后端 首次发表于 PyTorch。

来源：PyTorch Blog · pytorch.org

Meta / Llama推理能力开源生态部署工程教程实践#教程/实践#部署/工程#推理#开源生态#Meta