Triton & Gluon Kernel Development

Source Code Locations

Triton 源码位于此 skill 安装目录下的 repos/triton/。 实际路径取决于所用工具:

• Cursor: ~/.cursor/skills/triton-skill/repos/triton/
• Claude Code: ~/.claude/skills/triton-skill/repos/triton/
• Codex: ~/.agents/skills/triton-skill/repos/triton/

TRITON_REPO: 下文示例用 ~/.cursor/skills/triton-skill/repos/triton/ 作占位符，替换为实际路径。

如果该路径不存在，在项目目录下运行 bash update-repos.sh triton。

Triton Tutorials (入门到进阶)

TRITON_REPO/python/tutorials/
├── 01-vector-add.py             # Triton 基础: @triton.jit, program_id, load/store
├── 02-fused-softmax.py          # 内核融合, reduction, tl.max/tl.sum/tl.exp
├── 03-matrix-multiplication.py  # Block matmul, L2 cache, @triton.autotune
├── 04-low-memory-dropout.py     # 并行 PRNG, tl.rand, seed-based dropout
├── 05-layer-norm.py             # Backward pass, atomic ops, tl.atomic_cas
├── 06-fused-attention.py        # Flash Attention v2, causal mask, FP8, warp spec
├── 07-extern-functions.py       # libdevice 外部函数调用
├── 08-grouped-gemm.py           # Group GEMM, TMA, tensor descriptors
├── 09-persistent-matmul.py      # 持久化内核, TMA, warp specialization, FP8
├── 10-block-scaled-matmul.py    # FP4/FP8, MXFP4, tl.dot_scaled
└── 11-programmatic-dependent-launch.py  # PDL, gdc_wait, gdc_launch_dependents

Gluon Tutorials (底层 GPU 编程)

TRITON_REPO/python/tutorials/gluon/
├── 01-intro.py                  # Gluon vs Triton, tile-based SPMD, @gluon.jit
├── 02-layouts.py                # BlockedLayout, size_per_thread, warps_per_cta
├── 03-async-copy.py             # cp.async, 流水线, shared memory
├── 04-tma.py                    # Tensor Memory Accelerator, tensor desc, mbarrier
├── 05-wgmma.py                  # Warp-Group MMA, Hopper Tensor Core, async MMA
├── 06-tcgen05.py                # Blackwell Tensor Core, Tensor Memory, tcgen05_mma
├── 07-persistence.py            # 持久化内核, work assignment, 多级流水线
├── 08-warp-specialization.py    # Warp 特化, 任务重叠
├── 09-tma-gather-scatter.py     # Native TMA Gather/Scatter (Blackwell)
├── 10-tcgen05-copy.py           # tcgen05_copy, shared→tensor memory
├── 11-tcgen05-mma-scaled.py     # tcgen05_mma_scaled, nvfp4/mxfp4/mxfp8
└── 12-cluster-launch-control.py # CLC, 动态 work distribution

Gluon Examples (完整实现)

TRITON_REPO/python/examples/gluon/
└── 01-attention-forward.py      # Flash Attention forward (Blackwell)
                                 # 完整的 producer/consumer, TMA, tcgen05_mma

Triton Kernels (生产级参考实现)

TRITON_REPO/python/triton_kernels/triton_kernels/
├── matmul.py                    # 矩阵乘法 API (融合激活/MoE/ragged)
├── matmul_details/
│   ├── _matmul.py               # Dense GEMM kernel (TMA, mxfp4/8)
│   ├── _p_matmul.py             # Persistent GEMM kernel (ragged TMA)
│   └── _common.py               # 偏移计算, XCD swizzle
├── reduce.py                    # Reduction kernel (mask/scale/mxfp/flexpoint)
├── topk.py                      # Top-K selection (forward/backward, bitmatrix)
├── swiglu.py                    # SwiGLU activation kernel
├── compaction.py                # Masked compaction kernel
├── numerics.py                  # FP8/MXFP 数值配置
├── numerics_details/
│   ├── mxfp.py                  # MXFP 量化/反量化
│   └── flexpoint.py             # Flexpoint 缩放
├── tensor.py                    # Tensor/Layout 抽象 (TMA descriptors)
├── tensor_details/
│   └── layout_details/          # Blackwell/Hopper/CDNA4 MX 布局
├── distributed.py               # 分布式 MoE, SymmetricMemory
├── testing.py                   # 测试工具 (assert_close, compute_sanitizer)
└── roofline.py                  # Roofline 性能分析

Triton 语言源码

TRITON_REPO/python/triton/language/           # tl.* 操作的定义和语义
TRITON_REPO/python/triton/experimental/gluon/ # gluon.* 操作的定义
TRITON_REPO/python/triton/runtime/            # JIT 编译, 缓存, 解释器
TRITON_REPO/python/triton/compiler/           # 代码生成
TRITON_REPO/python/triton/tools/              # Tensor descriptor 工具

C++ 编译器 (IR 定义和 Passes)

TRITON_REPO/include/triton/
├── Dialect/
│   ├── Triton/           # Triton IR dialect 定义 (.td, .h)
│   ├── TritonGPU/        # TritonGPU dialect (layouts, encodings)
│   ├── TritonNvidiaGPU/  # NVIDIA 特定 ops (wgmma, tma, tcgen05)
│   └── Gluon/            # Gluon dialect
├── Conversion/           # IR lowering passes (TritonGPU -> LLVM)
├── Analysis/             # Alias, Allocation, AxisInfo, Membar
└── Tools/                # 工具类

TRITON_REPO/lib/
├── Dialect/
│   ├── Triton/           # Triton ops 实现, canonicalize
│   ├── TritonGPU/        # GPU layout 优化 passes
│   ├── TritonNvidiaGPU/  # NVIDIA lowering
│   └── Gluon/            # Gluon ops 实现
├── Conversion/           # Lowering pass 实现 (TritonGPU -> LLVM IR)
├── Analysis/             # 分析 pass 实现
└── Target/               # 代码生成目标

Search Strategy

用 Grep 工具搜索，不要整文件加载。

先确定 TRITON_REPO 的实际路径，然后用绝对路径搜索。

Triton API 用法

# 设置路径变量（替换为实际路径）
TRITON_REPO="$HOME/.cursor/skills/triton-skill/repos/triton"

# 查找 tl.dot 的使用方式
rg "tl\.dot" $TRITON_REPO/python/tutorials/

# 查找 autotune 配置示例
rg "@triton.autotune" $TRITON_REPO/python/tutorials/

# 查找 tensor descriptor 创建
rg "TensorDescriptor" $TRITON_REPO/python/tutorials/

# 查找特定 tl 操作的定义
rg "def (load|store|dot)" $TRITON_REPO/python/triton/language/

Gluon API 用法

# 查找 gluon.jit 使用
rg "@gluon.jit" $TRITON_REPO/python/tutorials/gluon/

# 查找 wgmma 用法
rg "wgmma" $TRITON_REPO/python/tutorials/gluon/05-wgmma.py

# 查找 tcgen05 用法 (Blackwell)
rg "tcgen05" $TRITON_REPO/python/tutorials/gluon/

# 查找 TMA 异步拷贝模式
rg "async_copy" $TRITON_REPO/python/tutorials/gluon/

# 查找 mbarrier 使用
rg "mbarrier" $TRITON_REPO/python/tutorials/gluon/

编译器 IR 和 Passes

# 查找 Triton IR op 定义 (TableGen)
rg "def.*Op" $TRITON_REPO/include/triton/Dialect/Triton/IR/

# 查找 TritonGPU layout encoding
rg "Encoding" $TRITON_REPO/include/triton/Dialect/TritonGPU/IR/

# 查找 NVIDIA 特定 ops (wgmma, tma)
rg "wgmma\|tma\|tcgen05" $TRITON_REPO/include/triton/Dialect/TritonNvidiaGPU/

# 查找 lowering pass 实现
rg "Pattern\|Rewrite" $TRITON_REPO/lib/Conversion/TritonGPUToLLVM/

# 查找 Gluon dialect ops
rg "def.*Op" $TRITON_REPO/include/triton/Dialect/Gluon/

# 查找特定 pass (如 coalesce, pipeline, prefetch)
rg "coalesce\|pipeline\|prefetch" $TRITON_REPO/lib/Dialect/TritonGPU/Transforms/

生产内核参考

# 查找 matmul 内核的 TMA 使用
rg "tma" $TRITON_REPO/python/triton_kernels/triton_kernels/matmul_details/

# 查找 MXFP 量化实现
rg "mxfp" $TRITON_REPO/python/triton_kernels/triton_kernels/numerics_details/

# 查找 persistent kernel 模式
rg "persistent" $TRITON_REPO/python/triton_kernels/triton_kernels/matmul_details/_p_matmul.py

# 查找 layout swizzle
rg "swizzle" $TRITON_REPO/python/triton_kernels/triton_kernels/tensor_details/layout_details/

When to Use Each Source

Need	Source	Path
Triton 基础语法和模式	Tutorials 01-05	python/tutorials/01-*.py ~ 05-*.py
矩阵乘法优化	Tutorial 03, 09, 10	python/tutorials/03-*.py, 09-*.py, 10-*.py
Attention 内核	Tutorial 06, Gluon example	python/tutorials/06-*.py, python/examples/gluon/
Gluon 入门	Gluon tutorials 01-02	python/tutorials/gluon/01-intro.py, 02-layouts.py
TMA 异步拷贝	Gluon tutorials 03-04	python/tutorials/gluon/03-*.py, 04-*.py
WGMMA (Hopper)	Gluon tutorial 05	python/tutorials/gluon/05-wgmma.py
tcgen05 (Blackwell)	Gluon tutorials 06, 10, 11	python/tutorials/gluon/06-*.py, 10-*.py, 11-*.py
持久化内核模式	Gluon tutorial 07	python/tutorials/gluon/07-persistence.py
Warp 特化模式	Gluon tutorial 08	python/tutorials/gluon/08-warp-specialization.py
FP4/FP8/MXFP 量化	Tutorial 10, triton_kernels	python/tutorials/10-*.py, triton_kernels/numerics_details/
生产级 GEMM	triton_kernels matmul	triton_kernels/triton_kernels/matmul_details/
MoE / Ragged tensor	triton_kernels	triton_kernels/triton_kernels/distributed.py, tensor_details/ragged_tensor.py
Top-K / SwiGLU 内核	triton_kernels	triton_kernels/triton_kernels/topk.py, swiglu.py
Roofline 性能分析	triton_kernels	triton_kernels/triton_kernels/roofline.py
tl.* 操作语义/签名	Language source	python/triton/language/
布局和 swizzle	triton_kernels layouts	triton_kernels/triton_kernels/tensor_details/layout_details/
Triton IR op 定义	include	include/triton/Dialect/Triton/IR/
GPU layout encoding	include	include/triton/Dialect/TritonGPU/IR/
NVIDIA ops (wgmma/tma)	include	include/triton/Dialect/TritonNvidiaGPU/
Gluon dialect 定义	include	include/triton/Dialect/Gluon/
编译 pass 实现	lib	lib/Dialect/TritonGPU/Transforms/
IR lowering (GPU->LLVM)	lib	lib/Conversion/TritonGPUToLLVM/

Triton Kernel 编写模式

基本模式

import triton
import triton.language as tl

@triton.jit
def kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = x * 2  # 计算
    tl.store(y_ptr + offsets, y, mask=mask)

# 启动
grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)
kernel[grid](x, y, n, BLOCK_SIZE=1024)

Autotune 模式

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    ...

参考 python/tutorials/03-matrix-multiplication.py 获取完整 autotune matmul 示例。

Gluon 基本模式

from triton.experimental import gluon

@gluon.jit
def kernel(x: gluon.tensor[M, N, tl.float16]):
    # 直接使用 tile 操作
    y = x + 1.0
    return y

参考 python/tutorials/gluon/01-intro.py 获取 Gluon 入门示例。

常见问题排查

问题	可能原因	查找参考
tl.dot 结果错误	输入类型不匹配 (需要 float16/bfloat16)	rg "tl.dot" tutorials/03-*.py
CUDA OOM	BLOCK_SIZE 过大, num_stages 过多	rg "num_stages" tutorials/09-*.py
autotune 无效	key 参数未对齐实际变化维度	rg "key=" tutorials/03-*.py
TMA descriptor 错误	tensor 不连续或维度不匹配	rg "TensorDescriptor" tutorials/gluon/04-*.py
wgmma 精度问题	需要 float32 累加器	rg "accumulator" tutorials/gluon/05-wgmma.py
性能低	未用 persistent kernel 或 warp spec	tutorials/gluon/07-persistence.py, 08-warp-specialization.py

Triton 编译和调试

# 查看生成的 PTX
TRITON_PRINT_AUTOTUNING=1 python your_script.py

# 启用 IR dump
MLIR_ENABLE_DUMP=1 python your_script.py

# 使用 Triton 的 compute-sanitizer
from triton_kernels.testing import compute_sanitizer
# 参考: python/triton_kernels/triton_kernels/testing.py

# 性能分析
from triton_kernels.roofline import compute_roofline
# 参考: python/triton_kernels/triton_kernels/roofline.py

更新 Triton 源码

# 在 cursor-gpu-skills 项目目录下
bash update-repos.sh triton

Additional References

• Triton 官方文档: https://triton-lang.org
• Triton Language API: TRITON_REPO/python/triton/language/
• Gluon Experimental API: TRITON_REPO/python/triton/experimental/gluon/