大模型推理的算力瓶颈从来不在模型本身,而在底层 kernel。vLLM 社区长期依赖手写 CUDA kernel 来榨取 GPU 性能,但每换一代硬件就得重写一遍——H100 上的优化到了 B200 可能就不再是最优路径。Helion 的思路很直接:用 PyTorch 原生的方式写 kernel,让同一份代码在 H100 和 B200 上都能跑出接近手写 CUDA 的性能。最近 Helion kernel 正式接入 vLLM,用 Qwen3 模型做了 FP8 推理的端到端验证,结果值得关注。
为什么 vLLM 需要"可移植"的 kernel
vLLM 的 PagedAttention、量化算子等核心 kernel 都是 CUDA C++ 手写的。好处是性能极致,代价是维护成本极高:
- 硬件绑定:每个新 GPU 架构(Hopper、Blackwell)的 tensor core 指令集、内存层级都有变化,kernel 参数需要重新调优。
- 人才门槛:写高性能 CUDA kernel 需要同时懂 GPU 微架构和模型计算图,这类工程师稀缺。
- 迭代速度慢:一个 FP8 kernel 从原型到合入 vLLM main branch,往往要数月的 PR review 和性能回归测试。
Helion 提供的是另一种开发范式——在 PyTorch 的 Python 层面描述 kernel 的计算逻辑,由编译器自动映射到不同 GPU 的指令集。开发者不再需要为 H100 和 B200 分别维护两份 kernel 代码。
Helion 的 PyTorch-native kernel 开发模型
Helion 的核心设计原则:kernel 就是 PyTorch 函数,只不过运行在 GPU 的细粒度并行维度上。
一个典型的 Helion kernel 看起来像这样:
import helion
import torch
@helion.kernel
def fp8_matmul(
a: torch.Tensor, # shape [M, K], dtype float8_e4m3fn
b: torch.Tensor, # shape [K, N], dtype float8_e4m3fn
scale_a: torch.Tensor, # per-tensor or per-row scale
scale_b: torch.Tensor, # per-tensor or per-col scale
) -> torch.Tensor:
# Helion 自动推断 block 大小和并行策略
M, K = a.shape
K2, N = b.shape
assert K == K2
# 声明输出 block,Helion 会分配寄存器/共享内存
output = torch.empty(M, N, dtype=torch.float32, device=a.device)
# 用 helion 的 block 循环描述 tile 级并行
for tile_m in helion.range(M, block_size=128):
for tile_n in helion.range(N, block_size=128):
acc = torch.zeros([128, 128], dtype=torch.float32)
for tile_k in helion.range(K, block_size=64):
a_tile = a[tile_m, tile_k] # [128, 64] float8
b_tile = b[tile_k, tile_n] # [64, 128] float8
# FP8 matmul + scale,映射到 tensor core 指令
acc += helion.dot(a_tile, b_tile) * scale_a[tile_m] * scale_b[tile_n]
output[tile_m, tile_n] = acc
return output
关键点解读:
@helion.kernel装饰器告诉编译器这是一个 GPU kernel,不是普通 PyTorch 函数。helion.range的 block_size 参数控制 tile 大小,编译器会根据目标 GPU 的 tensor core 尺寸自动选择最优值——在 H100 上可能映射到 128×128 的 MMA,在 B200 上可能调整到匹配其更大的 tensor core cluster。helion.dot是 FP8 矩阵乘法的抽象,编译器将其翻译为目标架构对应的 tensor core 指令。- 整份代码只有 Python,没有一行 CUDA C++。
在 vLLM 中启用 Helion FP8 kernel 的实测
Helion 团队用 Qwen3 系列模型在 vLLM 上做了完整的 FP8 推理基准测试,覆盖 H100(Hopper)和 B200(Blackwell)两张卡。测试配置如下:
# 拉取集成 Helion kernel 的 vLLM 分支
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout helion-fp8-integration
# 安装 Helion(需要 Python 3.10+,CUDA 12.4+)
pip install helion
# 安装 vLLM 开发模式
pip install -e .
# 用 Qwen3-32B FP8 量化模型启动推理服务
python -m vllm.entrypoints.openai.api_server \
--model Qwen/Qwen3-32B-FP8 \
--quantization fp8 \
--kv-cache-dtype fp8_e4m3fn \
--max-model-len 8192 \
--gpu-memory-utilization 0.90 \
--enforce-eager \
--use-helion-kernels
--use-helion-kernels是本次集成新增的 flag,启用后 vLLM 的 attention 和量化 matmul kernel 会走 Helion 实现,而非原生 CUDA 路径。
实测数据要点(基于公开的 benchmark 报告):
| GPU | 模型 | Batch Size | Throughput (tokens/s) | vs 原生 CUDA kernel |
|---|---|---|---|---|
| H100 | Qwen3-32B-FP8 | 32 | ~4200 | 差距 < 5% |
| B200 | Qwen3-32B-FP8 | 32 | ~6800 | 差距 < 3% |
在 B200 上 Helion kernel 的吞吐甚至更接近手写 CUDA,说明编译器对 Blackwell 架构的 tile 调优做得不错。延迟方面,TTFT(首 token 延迟)和 TPOT(每 token 延迟)的差距也在可接受范围内。
自己写一个 Helion kernel 并接入 vLLM
如果你想为 vLLM 的某个算子写 Helion 替代 kernel,流程大致如下。以 FP8 的 RMSNorm 为例:
import helion
import torch
@helion.kernel
def fp8_rmsnorm(
x: torch.Tensor, # [batch, seq_len, hidden_dim] float8_e4m3fn
weight: torch.Tensor, # [hidden_dim] float32
scale: torch.Tensor, # per-tensor scale for x
eps: float = 1e-6,
) -> torch.Tensor:
batch, seq_len, hidden_dim = x.shape
output = torch.empty(batch, seq_len, hidden_dim, dtype=torch.float8_e5m2, device=x.device)
for tile_b in helion.range(batch, block_size=1):
for tile_s in helion.range(seq_len, block_size=32):
# 读取 float8 tile,反量化到 float32
x_tile = x[tile_b, tile_s].to(torch.float32) * scale
# 计算 RMS
variance = (x_tile * x_tile).sum(dim=-1) / hidden_dim
rms = torch.sqrt(variance + eps)
# 归一化 + 乘 weight
normed = x_tile / rms.unsqueeze(-1) * weight
# 量化回 float8
output[tile_b, tile_s] = normed.to(torch.float8_e5m2)
return output
然后在 vLLM 的模型代码中替换调用:
# 在 vllm/model_executor/models/qwen3.py 中
# 原来的调用:
# hidden_states = self.rms_norm(hidden_states, weight)
# 替换为:
if self.use_helion_kernels:
hidden_states = fp8_rmsnorm(
hidden_states, weight, self.input_scale, eps=self.eps
)
else:
hidden_states = self.rms_norm(hidden_states, weight)
注册到 vLLM 的 kernel 选择机制:
# vllm/attention/backends/helion.py
from helion import register_kernel_override
register_kernel_override(
op_name="rms_norm_fp8",
kernel_fn=fp8_rmsnorm,
supported_dtypes=[torch.float8_e4m3fn, torch.float8_e5m2],
supported_gpus=["hopper", "blackwell"],
)
这样当 --use-helion-kernels 启用时,vLLM 的 kernel dispatcher 会自动选择 Helion 实现。
落地建议与当前边界
适合采用的场景:
- 你需要同时支持 H100 和 B200(或未来更多架构),不想维护多份 CUDA kernel。
- 团队 CUDA 经验有限,但 PyTorch 熟练度高——Helion 的开发效率优势明显。
- FP8 推理是主力场景,Helion 目前对 FP8 matmul + attention 的覆盖最成熟。
需要注意的边界:
- 性能天花板:极端场景下(比如极小 batch 的 decode 阶段),手写 CUDA 仍然可以通过更精细的寄存器分配和 warp 调度拿到额外 5-10% 的优势。Helion 编译器在持续改进,但短期内不会完全取代手写 kernel。
- dtype 覆盖范围:目前 Helion 在 vLLM 中的集成主要面向 FP8(e4m3fn / e5m2),BF16/INT8 的 kernel 覆盖还在推进中。
- 调试体验:Helion kernel 的编译错误信息比原生 CUDA 更友好,但性能调优(比如 block_size 选择)仍需要一定的 GPU 微架构知识。
- 生态成熟度:Helion 本身还在快速迭代,API 可能有变动,生产环境采用前建议锁定版本并做好回归测试。
快速验证 checklist:
- 确认 GPU 驱动 ≥ CUDA 12.4,Helion 对低版本驱动支持有限。
- 用
helion.detect_gpu_arch()确认你的卡被识别为 hopper 或 blackwell。 - 先在单卡上跑 Qwen3-8B-FP8 的 smoke test,确认 kernel 加载无误。
- 对比
--use-helion-kernels开关的开/关吞吐差异,确认在你的 batch size 分布下差距可接受。 - 监控 GPU 功耗和内存占用——Helion kernel 的共享内存使用模式可能与原生 CUDA 不同。
Helion 在 vLLM 中的集成是一个信号:GPU kernel 的开发范式正在从"手写汇编级 CUDA"向"PyTorch-native + 编译器自动优化"迁移。对于大多数推理团队来说,5% 的性能差距换来的是跨架构的可移植性和 10 倍的开发效率提升,这笔账值得算。