在深度学习框架的工程实践中，Triton 凭借高效的编译能力，成为 GPU 算子性能优化的核心工具。但在实际落地场景里，小计算量算子的调用痛点愈发突出 ——Python 层的 Wrapper + JIT runtime 开销，常常抵消掉 Triton 本身的性能优势，甚至让小算子执行效率大幅下降。

近期，DeepSeek V4 开源模型在多种 AI 芯片适配过程中，众智 FlagOS 采用了 Triton 重写全部新增算子，并通过 C++ Wrapper 优化调用链路，实测结果显示，在 NV H20 上，FlagOS 技术加持下的 DeepSeek V4 端到端推理性能最高超过原生发布版本的 11.2%，有效验证了 C++ Wrapper 在实际大模型推理中的性能收益。

本文将结合行业实践与实测数据，探讨如何用 C++ Wrapper 砍掉 Triton 算子的运行时冗余开销，让轻量算子也能跑出极致性能。

计算量很小的算子开发，Wrapper+JIT runtime 开销大

在深度学习领域，尤其是在处理大规模数据集和复杂模型时，算子的性能至关重要。Triton 作为一款高性能的编译器，通过 JIT（Just-In-Time）编译技术，能够在运行时生成高效的机器码，从而大幅提升算子的执行速度。然而，对于那些计算量很小的算子来说，这种优势往往被 Wrapper+JIT runtime 的开销所抵消。

具体来说，对于计算量很小的算子，Python 端的 Wrapper 和 JIT Runtime 的开销可能远大于 GPU Kernel 本身的执行时间。例如，一个简单的向量加法，GPU 计算本身可能只需要2微秒，但加上 Python 调用开销后，总时间可能增加到 60-100 微秒。这种开销不仅降低了整体性能，还使得开发者在优化代码时面临更多的挑战。

Python 调用主要带来三方面的问题：

• 类型转换开销：Python 和 C++ 之间的参数传递需要大量的类型转换

• 内存管理开销：Python 的垃圾回收机制不可预测，容易引入额外延迟

• 重复编译开销：虽然 Triton 有缓存机制，但每次调用仍有查找和分发的开销

这种情况在以下场景中尤为突出：

• 小 Tensor 操作（如序列长度小于512的Attention）

• Transformer 中大量的 add、mul 等逐元素操作

• 通过 torch.ops 注册的算子（需要额外的 Dispatcher 开销）

 C++ Wrapper 如何为 Triton 算子“减负”？

FlagOS 系统软件栈核心组件 FlagGems 提供了 C++Wrapper 来将 Wrapper 和 JITRuntime 从 Python 移至 C++，大幅减少调用开销。从架构上看，FlagGems 的 C++ Wrapper 方案与纯 Python 方案的对比如下。

组件	纯 Python 方案	C++ Wrapper 方案
Wrapper	Python 函数	C++ 函数（C++端重新实现）
JIT Runtime	Python Triton JIT Runtime	C++ TritonJITFunction（管理 key→TritonKernel 缓存）
CompiledKernel	Python 对象（通过 C binding 调用 cubin）	C++ TritonKernel（cubin 的简单封装）
Compiler	Python triton.compile	standalone_compiler（独立编译脚本，仍复用 Python triton.compile）

具体来说，通过以下几种方式来降低调用开销：

• 减少类型转换：在 Python 中调用 C++ 代码时，需要进行大量的类型转换。通过 C++ wrapper，可以直接在 C++ 层面进行类型管理和内存分配，从而避免这些开销。

• 优化内存管理：Python 的内存管理机制相对复杂，且容易引入垃圾回收等不可预测的因素。相比之下，C++ 的内存管理更加直接和高效。通过 C++ wrapper，可以更精细地控制内存的分配和释放。

• 避免重复编译：在C++端维护key → TritonKernel的缓存，由TritonJITFunction 管理，避免重复的编译和查找开销。

下面是一个简化示例，展示如何在 C++ 层封装 Triton 算子的调用流程：

#include <unordered_map>
struct KernelKey {
    int M, N, K;
    bool operator==(const KernelKey& other) const {
        return M == other.M && N == other.N && K == other.K;
    }
};
namespace std {
template <>
struct hash<KernelKey> {
    size_t operator()(const KernelKey& k) const {
        return (k.M * 131 + k.N) * 131 + k.K;
    }
};
}
class MatMulWrapper {
public:
    void matmul(float* A, float* B, float* C, int M, int N, int K) {
        KernelKey key{M, N, K};
        // kernel cache（示意）
        if (kernel_cache_.find(key) == kernel_cache_.end()) {
            kernel_cache_[key] = compile_kernel();
        }
        // launch（示意）
        launch_kernel(kernel_cache_[key], A, B, C, M, N, K);
    }
private:
    void* compile_kernel() {
        // 实际实现中：
        // Triton kernel 通常通过 Python JIT 编译生成
        return nullptr;
    }
    void launch_kernel(void* kernel,
                       float* A, float* B, float* C,
                       int M, int N, int K) {
        // 调用 runtime（示意）
    }
private:
    std::unordered_map<KernelKey, void*> kernel_cache_;
};

在这个示例中，我们创建了一个 MatMulWrapper 类，对 Triton kernel 的调用流程进行了封装，并在 C++ 层维护 kernel cache，从而避免重复编译带来的开销。通过这种方式，可以在 C++ 层完成算子的调度与参数传递，从而减少 Python 调用和类型转换带来的额外开销。

实战案例：DeepSeek V4 多芯片适配中的 C++ Wrapper 应用

在实际大模型推理场景中，C++ Wrapper 的价值被进一步放大。以最近开源的 DeepSeek V4 模型为例，其原生版本采用 “CUDA + TileLang” 组合使用了约 67 个算子。为了使模型能够在多种 AI 芯片上高效运行，众智 FlagOS 采用 Triton / Triton-TLE 重写了全部新增算子，并基于 Triton-TLE 实现了 TopK Selector 等高难度算子，实现“算子全覆盖”。

然而，在端到端推理测试中，团队发现了一个典型瓶颈：Triton 算子 + Python Wrapper 的端到端吞吐比 DeepSeek 原生版本降低了 20.05%。分析表明，Transformer 推理场景下一次 forward pass 需要调用数十个算子，大量小 kernel 的频繁 launch 使 Python 解释器调度开销（参数序列化、GIL 竞争等）成为性能瓶颈。

为此，FlagOS 团队引入了 C++ Wrapper 技术，基于 libtriton_jit 库，将 Triton kernel 的完整调用链路（参数绑定、grid 计算、kernel launch）下沉到 C++ 层，通过 pybind11 暴露为 Python 可调用接口。具体解决了三个核心问题：

• 消除 Python GIL 竞争：kernel launch 不再持有 GIL，多线程调度不阻塞；

• 减少解释器开销：参数传递从 Python 对象序列化改为 C++ 原生类型直传；

• 支持设备信息缓存：避免每次调用重复查询设备属性。

在 DeepSeek V4 模型中，包括 fp8_matmul、act_quant、nonzero、copy_、to_copy 等 5 个重要算子开启了 C++ Wrapper。端到端测试结果（NV H20，DeepSeek V4 FP8）显示：C++ Wrapper 联合 Triton 的技术组合，让 DeepSeek-FlagOS 版本在端到端推理性能上，比 DeepSeek 原版提升约 11%。

这一案例充分说明：即使单算子性能已经全面超越 TileLang，若忽视 Python 层的调用开销，端到端性能仍可能大幅回退。而 C++ Wrapper 正是解决这一“最后一公里”问题的有效手段。

手把手教你使用 C++ Wrapper，附真实性能对比数据

使用 C++ wrapper 的核心目标，并不是用 C++ 重写 Triton kernel，而是在 C++ 层封装 Triton kernel 的调用路径，从而减少 Python 调度开销。

第一步，我们使用 Triton 实现一个简单的 add 算子；

# triton_kernels/add_kernel.py
import triton
import triton.language as tl
import torch
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n):
    pid = tl.program_id(0)
    if pid < n:
        x = tl.load(x_ptr + pid)
        y = tl.load(y_ptr + pid)
        tl.store(out_ptr + pid, x + y)
def triton_add(x: torch.Tensor, y: torch.Tensor):
    n = x.numel()
    out = torch.empty_like(x)
    grid = (n,)
    add_kernel[grid](x, y, out, n)
    return out

第二步，Python 端基准测试；

# benchmark.py
import torch
import time
from triton_kernels.add_kernel import triton_add
def benchmark(fn, *args, warmup=10, iters=100):
    for _ in range(warmup):
        fn(*args)
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(iters):
        fn(*args)
    torch.cuda.synchronize()
    end = time.perf_counter()
    return (end - start) * 1e6 / iters  # μs
if __name__ == "__main__":
    N = 128 * 1024
    x = torch.randn(N, device="cuda", dtype=torch.float32)
    y = torch.randn(N, device="cuda", dtype=torch.float32)
    print("torch (Python):", benchmark(lambda a, b: a + b, x, y))
    print("triton (Python):", benchmark(triton_add, x, y))

第三步，C++ Wrapper。在 C++ 中，我们通过 Python embedding 的方式调用 Triton kernel，从而将调用路径下沉到 C++；

// wrapper.cpp
#include <Python.h>
#include <iostream>
class TritonWrapper {
public:
    TritonWrapper() {
        Py_Initialize();
        PyRun_SimpleString("import sys");
        PyRun_SimpleString("sys.path.append('./triton_kernels')");
        module_ = PyImport_ImportModule("add_kernel");
        if (!module_) {
            PyErr_Print();
            throw std::runtime_error("Failed to load module");
        }
        func_ = PyObject_GetAttrString(module_, "triton_add");
    }
    ~TritonWrapper() {
        Py_Finalize();
    }
    void run(PyObject* x, PyObject* y) {
        PyObject* args = PyTuple_Pack(2, x, y);
        PyObject* result = PyObject_CallObject(func_, args);
        if (!result) {
            PyErr_Print();
        }
        Py_XDECREF(result);
        Py_DECREF(args);
    }
private:
    PyObject* module_;
    PyObject* func_;
};

第四步，C++ 侧性能测试；

// main.cpp
#include <chrono>
#include <iostream>
#include "wrapper.cpp"
int main() {
    TritonWrapper wrapper;
    const int iters = 100;
    // 这里为了演示调用路径，省略 Tensor 构造细节
    PyObject* x = Py_None;
    PyObject* y = Py_None;
    // warmup
    for (int i = 0; i < 10; ++i) {
        wrapper.run(x, y);
    }
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < iters; ++i) {
        wrapper.run(x, y);
    }
    auto end = std::chrono::high_resolution_clock::now();
    double us = std::chrono::duration<double, std::micro>(end - start).count() / iters;
    std::cout << "C++ wrapper time: " << us << " us\n";
}

性能对比结果如下：

Wrapper时间对比（单位：微秒）

算子	libtorch (C++)	torch (Python)	triton (Python)	triton_cpp (C++)
add（128k fp32）	14.44	14.44	61.11	14.44
sum_dim（16x4k fp32）	17.78	18.89	82.22	16.67

从数据可以看出：

• 纯 Python Triton wrapper 开销约为 C++ 方案的 4~5 倍；

• C++ Wrapper 方案的耗时与 libtorch（PyTorch原生C++接口）基本持平。

Kernel执行时间对比（单位：微秒）：

算子	PyTorch原生CUDA	Triton实现
add（128k fp32）	2.126	2.097
sum_dim（16x4k fp32）	8.805	4.538

在 sum_dim 算子上，本次测试中 FlagGems 的 Triton 实现，其 kernel 执行时间约为 PyTorch 原生 CUDA 实现的 1.94 倍加速（在相同测试环境与输入规模下）。

从整体性能上来看：

• 纯 Python Triton wrapper 开销是 libtorch 的 4~5 倍，torch_ops 注册路径会再增加约 1.5~2 倍开销

• C++ Wrapper 方案（triton_cpp）的 wrapper 时间与 libtorch 持平（约 15μs）

• Kernel 执行时间方面，FlagGems Triton 实现的 sum_dim 性能约是 PyTorch 原生的 1.94 倍

• 对于计算密集型算子，wrapper 开销可忽略；对于小算子，C++ Wrapper 至关重要

注：torch.op 的 unboxed 调用方式比 torch.ops.lib.op 效率更高，但实现较复杂

结语

通过本文的探讨，我们深入了解了在处理计算量很小的算子时，如何通过 C++ wrapper 来降低 Wrapper+JIT runtime 的开销。通过 C++ wrapper，我们可以直接在 C++ 层面进行类型管理和内存分配，从而避免 Python 调用带来的额外开销。实验数据也证明了这种方法的有效性，性能提升显著。对于大模型推理、端侧部署、高频小算子调用等场景，C++ Wrapper 能让 Triton 算子摆脱运行时束缚，真正释放硬件算力。如果你正在优化 Triton 算子的端到端性能，不妨从 C++ Wrapper 开始，一步到位解决小算子开销难题。

想深入了解 FlagGems 里的更多算子实现，可前往https://github.com/flagos-ai/FlagGems查看源码，也欢迎大家一起参加“FlagOS 开放计算全球挑战赛”算子开发和优化赛道（https://flagos.io/RaceDetail?id=296flq8k&lang=cn），使用 Triton 语言开发大模型常用算子，深入探索芯片体系结构，挑战算子极致性能！

FlagGems 是基于 Triton 的高性能通用算子库，作为 FlagOS 开源生态的核心组件，FlagGems 目前支持超过 400 个算子，基本覆盖大模型需求，82% 以上的算子性能可与 CUDA 原生算子平齐甚至实现超越，成为全球最大的 Triton 单一算子库。目前，FlagGems 已完成对英伟达、华为、摩尔线程、海光、天数等 28 种主流 AI 芯片的适配支持，在 40 个主流 AI 模型上的推理任务算子覆盖度达到 90%~100%，为开发者提供了极致的开发体验。

关于众智FlagOS社区

为解决不同 AI 芯片大规模落地应用，北京智源研究院联合众多科研机构、芯片企业、系统厂商、算法和软件相关单位等国内外机构共同发起并创立了众智 FlagOS 社区。成员单位包括北京智源研究院、中科院计算所、中科加禾、安谋科技、北京大学、北京师范大学、百度飞桨、硅基流动、寒武纪、海光信息、华为、基流科技、摩尔线程、沐曦股份、澎峰科技、清微智能、天数智芯、先进编译实验室、移动研究院、中国矿业大学(北京)等多家在 FlagOS 软件栈研发中做出卓越贡献的单位。

FlagOS 是一款专为异构 AI 芯片打造的开源、统一系统软件栈，支持 AI 模型一次开发即可无缝移植至各类硬件平台，大幅降低迁移与适配成本。它包括大型算子库、统一AI编译器、并行训推框架、统一通信库等核心开源项目，致力于构建「模型-系统-芯片」三层贯通的开放技术生态，通过“一次开发跨芯迁移”释放硬件计算潜力，打破不同芯片软件栈之间生态隔离。

官网：https://flagos.io

GitHub 项目地址：https://github.com/flagos-ai

GitCode 项目地址：https://gitcode.com/flagos-ai

SkillHub: https://skillhub.flagos.io