Helion on TPU: Towards Hardware Heterogeneous Kernel Authoring

Helion on TPU：迈向异构硬件内核开发 – PyTorch

特色项目

TL;DR

Helion 是 PyTorch 用于编写高性能可移植机器学习内核的高级领域特定语言（DSL）。与 Google 合作，我们构建了 TPU 后端，将 Helion 内核编译为 Pallas，为编写高性能 TPU 内核提供了 PyTorch 友好的方式。在闪存注意力工作负载中，Helion 生成的内核在 TPU v7 上实现了 838 TFLOPs（约单个张量核心的 79% MFU）。针对不同输入形状，Helion 通过不同代码生成策略进行自动调优，选择最佳流水线方案，最大程度利用 TPU 的可用 VMEM 和计算资源。

引言

随着 ML 计算平台的演进，TPU 作为 GPU 的补充变得越来越重要。Google 最新的 TPU v7（Ironwood）在潜在更低的总拥有成本（TCO）下实现了与 NVIDIA B200 相当的性能，使 TPU 成为大规模训练和推理任务的有吸引力的选择。然而，传统上编写 TPU 内核需要掌握 Pallas——一个学习曲线陡峭且代码复杂度高的低级 DSL。Helion 桥接了这一鸿沟。作为 PyTorch 的可移植 DSL，Helion 允许用户编写熟悉的 PyTorch 风格代码，并将其编译为优化的 TPU 代码。结合其自动调优器带来的性能优势，Helion 正逐渐成为编写 TPU 内核的有吸引力的选择。具体而言，Helion TPU 针对以下三个主要使用场景进行优化：

• 需要自动调优探索配置空间的性能关键型场景
• 希望快速上手 TPU 内核开发的非 Pallas 专家
• 希望在 TPU 和 GPU 上维护相同内核集的跨硬件用户

本文首先简要概述 TPU 与 GPU 相比的硬件特性和编程模型，然后演示 Helion 如何为不同输入形状生成具有理想流水线特性的高性能 Pallas 代码。

TPU 入门

TPU 是专为机器学习工作负载设计和优化的高度专业化的加速器。TPU 的架构和编程模型与 GPU 存在显著差异。最显著的区别在于，TPU 是一个顺序机器，具有宽向量寄存器和计算单元。这与通过大规模并行执行（CUDA 核心）和专用张量单元（张量核心）实现性能的 GPU 形成对比。

| TPU（Pallas） | GPU（CUDA） | |---------------|-------------| | 线程模型 | 顺序执行 | 并行 SIMT（+ 张量核心） | | 工作者数量 | 几个大型工作者 | 许多小型工作者 | | 内存层次 | 显式内存空间（持久内存 vs 刮板内存），流水线需要异步内存拷贝 | 隐式缓存，硬件管理 |

因此，TPU 的内存层次结构是内核开发者必须深入理解的关键，这样才能编写出能够协调数据从片外 HBM 加载到快速片内 VMEM 的时机和方式的内核。高性能的 Pallas 内核会将这些 HBM<->VMEM 内存传输与矩阵（MXU）和向量计算单元中的浮点运算重叠执行。

尽管存在架构差异，当前一代 TPU 和 GPU 在原始性能上高度可比。TPU7x 和 NVIDIA B200 在 BF16 计算 TFLOPS 和 HBM 带宽上非常接近——这两个指标是现代 ML 工作负载最重要的硬件指标。

Helion 的 Pallas 代码生成

为了从TPU中提取最大性能，Helion的Pallas代码生成器旨在最大化软件流水线，确保内存传输和计算尽可能重叠。本节说明了Helion生成流水线内核的三重策略：

• 外层循环：由pallas提供的流水线设备调用（pallas_call / emit_pipeline）
• 内层循环：在以下两种方式之间自动调优：由pallas提供的流水线设备端循环（emit_pipeline）和尽可能将所有值预取到VMEM（unroll）
• 自动调优的流水线缓冲区大小

示例：add

作为简单示例，考虑以下用于添加两个张量的helion内核。

@helion.kernel
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    for tile in hl.tile(out.size()):
        out[tile] = x[tile] + y[tile]
    return out

Helion编译器会将其转换为两个函数：一个在主机端启动的函数，它对输入进行分块并以流水线方式调用设备函数，以及一个在VMEM驻留分块上运行的设备函数：

def _helion_add(x, y, out):
    out[:] = x[:] + y[:]

def add(x: torch.Tensor, y: torch.Tensor):
    _BLOCK_SIZE_0 = <autotuner-selected value>
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_add, 
        ((x.shape[0] + _BLOCK_SIZE_0 - 1) // _BLOCK_SIZE_0,), # grid size
        x, y, out, 
        _block_spec_info=[_BLOCK_SIZE_0, ...], ...
    )
    return out

在生成的代码中：

• Helion源代码中的hl.tile循环在主机端变为网格。启动器（包装pallas_call）对每个分块调用_helion_add一次，每次调用自动流水线化——当一个分块正在计算时，下一个分块的数据正在从HBM加载到VMEM。
• 设备函数_helion_add很简单：它接收VMEM引用（而非HBM指针），因此内核主体只是一个简单的加法。
• _BLOCK_SIZE_0（分块/缓冲区大小）由自动调优器选择，自动调优器会探索不同大小以找到目标硬件上内存传输和计算的最佳重叠。

这产生了如下面所示的流水线执行。

示例：Flash Attention

注意力是现代语言模型中的关键操作。生产实现遵循“Flash Attention”模式——一种内存高效的技巧，通过以分块方式计算注意力来避免生成完整的S×S注意力矩阵。下图展示了Helion中Flash Attention内核的结构：

B, H, S, D = 8, 32, 8192, 256 # batch, head, sequence length, head dimension
@helion.kernel
def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    out = torch.empty(...)
    for tile_b, tile_q in hl.tile(B * H, S):
        this_q = q[tile_b, tile_q, :]
        acc = ...
        for tile_kv in hl.tile(S):
            this_k = k[tile_b, tile_kv, :]
            this_v = v[tile_b, tile_kv, :]
            <qk matmul, online softmax, v matmul, update acc>
        out[tile_b, tile_q] = acc
    return out

与之前讨论的“add”示例相比，Flash Attention内核包含一个额外的内层循环，该循环在整个K和V序列上执行分块访问。我们在内层循环中如何对内存和计算进行流水线处理，是该内核性能的关键。

在 Helion 中，编译器会根据两种不同的策略自动调整将此内核转换为 Pallas 的方式。这一过程由 pallas_loop_type 自动调优配置进行控制。

当使用默认的 pallas_loop_type == emit_pipeline 配置时，Helion 会依赖 Pallas 的设备端 emit_pipeline API 对内层循环体函数进行流水线处理，这与主机端逻辑使用 pallas_call 对设备函数调用进行流水线处理的方式类似：

def _helion_attention(q_VMEM, k_HBM, v_HBM, out_VMEM):
    acc = ...
    this_q = q_VMEM[:, :, :]
    def _inner_pipeline_body(k_VMEM, v_VMEM):
        this_k = k_VMEM[:, :, :]
        this_v = v_VMEM[:, :, :]
        <matmul, online softmax, matmul, update acc>
    pallas.tpu.emit_pipeline(_inner_pipeline_body, k_HBM, v_HBM, _block_spec_info=[BLOCK_SIZE_KV ,...], ... )
    out_VMEM = acc

def attention(q: torch.Tensor, k: torch.Tensor, v:torch.Tensor):
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_attention, 
        q, k, v, out, 
        _block_spec_info=[BLOCK_SIZE_Q ,...], ...
    )
    return out

该生成的内核遵循嵌套流水线结构：

• （外层流水线）主机使用 pallas_call 调用 _helion_attention 。q 的 HBM 引用被分块处理，每次调用 _helion_attention 会接收 q 的 VMEM 分块。对于 k 和 v，_helion_attention 直接接收 HBM 引用。

• （内层流水线）在 _helion_attention 中，设备使用 emit_pipeline 调用 _inner_pipeline_body ，该函数接收 k 和 v 的 VMEM 分块。

这导致了如图所示的流水线执行方式：

该流水线一个显而易见的效率低下之处在于计算单元存在气泡 – 每当有新的 Q 分块到来时，在获取第 0 个 KV 分块期间，计算单元没有可用工作。这归因于我们为每个新的 Q 分块都从 HBM 重新加载 KV 分块到 VMEM。

Helion 提供了另一种 pallas_loop_type == unroll 配置，可以避免这种气泡现象。使用 unroll 配置时，我们将内层 for 循环转换为简单的 Python for 循环：

def _helion_attention(q_VMEM, k_VMEM_FULL, v_VMEM_FULL, out_VMEM):
    acc = ...
    this_q = q_VMEM[:, :, :]
    for offset in range(0, k_VMEM_FULL.size(1) , BLOCK_SIZE_KV):
        this_k = k_VMEM_FULL[:, pallas.dslice(offset, BLOCK_SIZE_KV), :]
        this_v = v_VMEM_FULL[:, pallas.dslice(offset, BLOCK_SIZE_KV), :]
        <matmul, online softmax, matmul, update acc>
    out_VMEM = acc

def attention(q: torch.Tensor, k: torch.Tensor, v:torch.Tensor):
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_attention, 
        q, k, v, out, 
        _block_spec_info=[BLOCK_SIZE_Q, None, None], ...
    )
    return out

（"unroll" 这个名称反映了 Pallas 设备函数被 JAX 的 JIT 追踪的事实 – Python for 循环在追踪时会被展开为一个扁平的操作序列。）

在该版本的生成内核中：

• K 和 V 被预先完全加载：主机将 K 和 V 的块规格设为 None，指示 pallas_call 将它们完全加载到 VMEM。完整的 VMEM 引用在所有设备函数调用中持续存在。

• 内层循环在本地进行切片：每次迭代使用 pallas.dslice 从已驻留的 VMEM 缓冲区中选择相关的 KV 分块。内层循环期间不会发生 HBM 通信。

这导致了如图所示的不同流水线方案：

在该工作流程中，计算流水线中不再存在气泡。这种优化的代价是需要更多的虚拟内存（VMEM）使用，因为必须完整存储所有K和V序列。这意味着尽管性能更优，但这种转换并非总能实现。VMEM使用量与输入序列长度呈线性关系（而非与分片大小相关），这在处理长序列时会带来显著限制。这两种策略的性能差异非常显著——下表展示了在B=8、H=32、D=256的工作负载下的结果：

S = 8k

S = 32k

emit_pipeline TFLOPs

653

695

unroll TFLOPs

892

OOM

Helion的优势在于其自动调优能力，能够选择最佳的自动调优配置。因此在处理短序列时，它会充分利用可用的VMEM并生成无计算气泡的流水线代码。对于长序列，它会回退到emit_pipeline方法，该方法可以扩展到任意长度的上下文。下图展示了该注意力内核在不同序列长度下的性能表现，与各种其他Pallas注意力实现进行对比：

Helion的自动调优器能够根据输入长度生成不同的循环和流水线策略，这使其即使与高度优化的实现（如Tokamax）相比也具有显著优势。

更广泛的内核基准测试

我们在多种内核上对Helion进行了基准测试，结果可在我们的仪表板上查看。下表将Helion与TorchTPU即时执行和torch.compile（使用XLA）在不同内核上的表现进行了对比。Helion相比即时执行的几何平均加速比为1.55倍，相比编译执行为1.12倍。

kernel

shape

torch_tpu eager (ms)

torch.compile(tpu) (ms)

Helion (ms)

Helion vs torch_tpu eager

Helion vs torch.compile

attention

[8,32,8192,256]

87.77

88.28

19.72

4.45×

4.48×

softmax

[65536,2560]

0.712

0.743

0.477

1.49×

1.56×

batch_softmax

[64,2048,4096]

1.888

1.373

0.982

1.92×

1.40×

softmax_two_pass

[8192,8192]

0.386

0.417

0.334

1.16×

1.25×

bmm

[64,2048,2048,2048]

3.211

1.860

1.527

2.10×

1.22×

rms_norm-bwd

1.792

0.789

0.661

2.71x

1.19x

epilogue_subtiling

[4096,4096,4096]

0.850

0.462

2.04x

1.11x

matmul_layernorm

0.535

0.523

0.489

1.10×

1.07×

welford

[524288,512]

1.330

1.357

1.316

1.01x

1.03x

swiglu

[16,16384,4096]

3.510

2.244

2.295

1.53×

0.98×

matmul

[8192,8192,8192]

1.552

1.597

0.97×

0.96×

geglu

[16,8192,8192]

3.779

2.240

2.424

0.92×

cross_entropy

[128,2048]

0.363

0.264

0.320

1.13×

0.82×

broadcast_matmul

1.817

1.440

1.806

1.01×

0.80×

layer_norm

[16384,16384]

1.253

0.779

1.126

1.11×

0.69×

rms_norm

1.194

0.419

0.617

1.94×

0.68×

Helion在涉及融合或优化模式的内核上表现最为突出，这些模式对XLA的自动发现能力构成挑战——Flash Attention就是典型例子。对于矩阵乘法和层归一化等标准操作，XLA编译器已能生成高质量代码，Helion的性能与之相当。

后续计划

Helion在TPU上的开发正在进行中。以下是我们正在推进的非穷举事项列表：

• 扩展内核覆盖范围：在TPU上实现更多Helion示例
• 进一步提升性能
• 更好地支持锯齿状和稀疏操作
• 支持分布式TPU计算

入门指南

Helion 是一个开源项目，可在 GitHub 上获取。其 TPU 后端依赖于 TorchTPU，该库预计今年晚些时候将公开发布。当 TorchTPU 发布后，我们鼓励您在 TPU 上尝试使用 Helion，并分享您的反馈意见。相关资源：

• Helion GitHub 代码仓库
• Helion 官方文档
• Helion TPU 示例代码
• 性能监控仪表盘

致谢

本项目得以实现，得益于同行们宝贵的协作和技术见解。特别感谢 Google 的 Joe Pamer、Robert Hundt、Claudio Basile 和 Adam Paszke，以及 Meta 的 Jana van Greunen、Gregory Chanan、Peng Wu 和 Zongwei Zhou，感谢他们为本项目提供的反馈和支持。

/post-content

/inner-wrap