这篇文章解决什么问题

这一节只做一件事：把后续系列会用到的两套 Triton 环境准备好，并完成最小验证。

这里的“两套 Triton”分别指：

  • OpenAI Triton：用于编写和调优 GPU kernel。
  • NVIDIA Triton Inference Server：用于部署模型服务。

两者名字相同，但用途完全不同。上一篇已经说明了职责边界，这一篇只关心环境是否可用。

需要先说明一个前提：本文默认命令在一台远端 Linux GPU 服务器上执行。本地 Windows 机器可以作为编辑器、SSH 客户端或文件管理端，但不承担 Triton kernel 编译执行，也不承担 Triton Server 的 GPU 推理任务。

如果当前机器没有 NVIDIA GPU，不需要在本地强行复现本文命令。把环境准备集中在真正的远端执行机上，会比在 Windows 本地做兼容性折腾更稳妥。

环境目标与验收标准

这篇文章的验收标准很简单：

  • OpenAI Triton 能够成功编译并执行一个最小向量加法 kernel。
  • NVIDIA Triton Inference Server 能够成功加载一个 ONNX 模型，并返回一次推理结果。

这两个结果一旦成立，后续关于编程模型、性能分析、服务调度和压测的内容就有了统一的实验起点。

环境边界

本文不讨论以下内容：

  • Windows 本地 CUDA 开发环境搭建。
  • 多机集群、Kubernetes、生产级部署编排。
  • 驱动升级、CUDA 多版本共存策略。

本文只覆盖一台远端 Linux GPU 服务器上的单机准备流程，目标是让系列内容能稳定继续，而不是一次性解决所有部署变体。

执行前先确认三件事

在安装任何组件之前，先确认下面三项。

1. 驱动与 GPU 是否可见

1
nvidia-smi

至少需要确认两点：

  • 系统能识别出 GPU 型号。
  • 驱动状态正常，没有出现设备不可见或驱动未加载的情况。

如果这里就失败，后续 OpenAI Triton 和 NVIDIA Triton 都没有继续安装的意义。

2. Python 与隔离环境可用

建议在远端服务器上使用独立 Python 环境，不要直接把实验依赖装到系统 Python。

1
2
3
4
python3 --version
python3 -m venv .venv
source .venv/bin/activate
python -V

如果团队已经统一使用 conda 或 uv，也可以沿用现有方式。重点不是工具本身，而是保证后续依赖可回收、可重建。

3. Docker 能正常访问 GPU

NVIDIA Triton Inference Server 通常通过容器运行，因此还需要确认 Docker 与 NVIDIA Container Toolkit 正常。

1
2
docker --version
docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

如果第二条命令无法看到 GPU，问题不在 Triton 本身，而在容器 GPU 运行时没有配置完成。这个问题必须先解决。

第一部分：准备 OpenAI Triton

版本选择原则

OpenAI Triton 与 PyTorch、CUDA 运行时之间存在版本耦合。为了让后续示例保持一致，这里采用一组明确版本：

  • torch==2.2.1
  • triton==2.3.0
  • CUDA 11.8 对应的 PyTorch 轮子

这不是唯一可用组合，但它有两个优点：

  • 依赖关系简单，适合系列文章统一复现。
  • 后续最小 kernel 示例与 PyTorch 集成示例都可以直接沿用。

安装命令

1
2
pip install torch==2.2.1 --index-url https://download.pytorch.org/whl/cu118
pip install triton==2.3.0

安装完成后，先确认 Python 侧依赖已就绪。

1
python -c "import torch, triton; print(torch.__version__, triton.__version__, torch.cuda.is_available())"

预期结果不是某个固定字符串，而是满足下面三个条件：

  • torch 成功导入。
  • triton 成功导入。
  • torch.cuda.is_available() 返回 True。

最小验证：向量加法 kernel

环境验证不需要一上来就写复杂算子。一个能跑通的向量加法 kernel 已经足够证明三件事：

  • Triton JIT 能正常工作。
  • GPU kernel 能被成功发射。
  • 结果能与 PyTorch 基线对齐。

将下面代码保存为 smoke_test_openai.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import torch
import triton
import triton.language as tl


@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    out_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def triton_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n = x.numel()
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)
    return out


if __name__ == "__main__":
    x = torch.rand(2**20, device="cuda", dtype=torch.float32)
    y = torch.rand(2**20, device="cuda", dtype=torch.float32)

    out_triton = triton_add(x, y)
    out_torch = x + y

    assert torch.allclose(out_triton, out_torch, atol=1e-5)
    print("OpenAI Triton smoke test passed")
    print("GPU:", torch.cuda.get_device_name(0))
    print("triton:", triton.__version__)
    print("torch:", torch.__version__)

执行命令：

1
python smoke_test_openai.py

如果这一步通过，说明当前远端服务器已经具备后续 OpenAI Triton 示例运行条件。

这一阶段最常见的错误

No module named 'triton'

通常是因为当前 shell 没有进入正确的 Python 环境，或者 pip 装到了别的解释器上。优先检查：

1
2
3
which python
which pip
python -m pip show triton

torch.cuda.is_available() 为 False

这通常不是 Triton 代码问题，而是 CUDA 运行环境没有准备好，或者当前 Python 环境装的是 CPU 版本 PyTorch。

kernel 启动时报 CUDA 相关错误

优先核对驱动状态、PyTorch 安装源和 GPU 可见性，不要过早怀疑 kernel 代码本身。

第二部分：准备 NVIDIA Triton Inference Server

OpenAI Triton 解决的是 kernel 执行问题，NVIDIA Triton 解决的是模型服务问题。后续服务端章节会统一使用一个最小 ONNX 模型，因此这里先把模型与服务闭环准备好。

第一步：导出一个最小 ONNX 模型

后续示例统一使用一个两层 MLP。它不追求业务意义，只追求结构简单、输入输出稳定、便于验证服务是否工作。

将下面代码保存为 export_model.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import os
import torch
import torch.nn as nn


class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))


os.makedirs("models/simple_mlp/1", exist_ok=True)

model = SimpleMLP().eval()
dummy = torch.randn(1, 128)

torch.onnx.export(
    model,
    dummy,
    "models/simple_mlp/1/model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
    opset_version=13,
)

print("export done: models/simple_mlp/1/model.onnx")

执行：

1
python export_model.py

执行完成后，目录结构应至少包含：

1
2
3
4
models/
  simple_mlp/
    1/
      model.onnx

第二步：补上 config.pbtxt

NVIDIA Triton 不会直接读取一个随意摆放的模型文件。模型必须放在标准目录结构里，并配套配置文件。

在 models/simple_mlp/ 下创建 config.pbtxt：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
name: "simple_mlp"
backend: "onnxruntime"
max_batch_size: 32

input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [ 128 ]
  }
]

output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [ 10 ]
  }
]

此时模型目录应为：

1
2
3
4
5
models/
  simple_mlp/
    config.pbtxt
    1/
      model.onnx

第三步：启动 Triton Server

拉取镜像：

1
docker pull nvcr.io/nvidia/tritonserver:23.12-py3

启动服务：

1
2
3
4
5
6
7
docker run --gpus all --rm \
  -p 8000:8000 \
  -p 8001:8001 \
  -p 8002:8002 \
  -v $(pwd)/models:/models \
  nvcr.io/nvidia/tritonserver:23.12-py3 \
  tritonserver --model-repository=/models

启动成功后，日志里通常会出现三类服务已启动的信息：

  • HTTP 服务
  • gRPC 服务
  • Metrics 服务

可以在另一个终端做一次最小健康检查：

1
curl -s http://localhost:8000/v2/health/ready

如果返回空 JSON，说明服务已经就绪。

第四步：发送一次最小推理请求

为了避免只验证“服务活着”，还需要验证“模型可推理”。先安装客户端：

1
pip install tritonclient[http]==2.40.0

保存为 infer_client.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import numpy as np
import tritonclient.http as httpclient


client = httpclient.InferenceServerClient(url="localhost:8000")

assert client.is_server_ready()
assert client.is_model_ready("simple_mlp")

batch_input = np.random.randn(2, 128).astype(np.float32)

infer_input = httpclient.InferInput("input", batch_input.shape, "FP32")
infer_input.set_data_from_numpy(batch_input)

infer_output = httpclient.InferRequestedOutput("output")

response = client.infer(
    model_name="simple_mlp",
    inputs=[infer_input],
    outputs=[infer_output],
)

result = response.as_numpy("output")
assert result.shape == (2, 10)

print("NVIDIA Triton smoke test passed")
print("input shape:", batch_input.shape)
print("output shape:", result.shape)

执行：

1
python infer_client.py

如果结果形状是 (2, 10)，说明这个最小服务闭环已经成立。

为什么这里只做最小闭环

这一节故意没有展开 dynamic batching、instance group、TensorRT backend 或多模型管理，原因有两个：

  • 环境篇的目标是确认“后续实验能否继续”，而不是提前引入过多变量。
  • 一旦环境、模型和请求三者的最小闭环成立，后续性能分析才有可比较的基线。

工程上，先建立一个稳定起点，再逐步增加复杂度，通常比一开始就堆满配置项更可靠。

常见误区

误区一：在 Windows 本地直接复现全部步骤

如果真实运行环境是远端 Linux GPU 服务器，那么本地 Windows 机器只需要承担编辑、同步和远程连接职责。强行把所有组件搬到本地，不会让后续实验更容易，反而会额外引入驱动、路径和容器兼容性问题。

误区二：把 OpenAI Triton 和 NVIDIA Triton 混成一套安装流程

它们的用途不同、运行方式不同、验证方式也不同。一个偏 Python 包与 kernel 执行，一个偏容器化服务部署。环境准备必须分开验证。

误区三：服务 ready 就等于模型可用

服务进程可启动，只能说明 Triton Server 本身启动成功；只有客户端请求能够返回正确形状，才能说明模型加载、配置与推理路径都正常。

结论

这一节的目标不是把 Triton 全部讲透，而是建立一个稳定起点。

  • 对 OpenAI Triton，最小验收是向量加法 kernel 成功运行。
  • 对 NVIDIA Triton Inference Server，最小验收是 ONNX 模型成功加载并返回一次推理结果。
  • 对当前系列而言，真正需要保证的是远端 Linux GPU 服务器上的实验环境可重复，而不是本地 Windows 机器是否具备执行能力。

完成这一节后，后续就可以分别进入两条主线：

  • 沿着 OpenAI Triton 继续看编程模型与性能思维。
  • 沿着 NVIDIA Triton 继续看服务配置、调度和压测闭环。

系列导航