Triton 自定义算子
Custom Triton Kernels
解决 PyTorch 原生算子无法高效处理分页内存访问的问题，通过 Triton JIT 编写自定义 GPU 内核
子问题
1.分页 KV Cache 的分散写入
2.Triton 内核的线程映射
3.与 PyTorch 张量的互操作
4.内核自动调优
5.CUDA Graph 固定 grid 下的 padding token 跳过机制
6.Triton kernel 与全局 Context 的元数据传递模式
各项目的解法1 solutions
Signals
横向对比
| 维度 | nano-vllm |
|---|---|
| 内核粒度 | 一 token 一 program，D=num_heads*head_dim 展平为一维向量操作 |
| 编程模型 | Triton JIT + tl.constexpr 编译期特化，10 行核心逻辑 |
| 地址映射 | slot_mapping 间接寻址，-1 哨兵支持 CUDA Graph padding |
| 与框架集成 | Python wrapper 做 stride 校验，全局 Context dataclass 传递元数据 |
| 自动调优 | 未使用 triton.autotune，D 为 constexpr 由编译器自动向量化 |
最佳实践
1.用 slot_mapping 实现逻辑地址到物理地址的映射
2.Triton 内核粒度对齐到 head_dim 以最大化并行度
3.用 tl.constexpr 声明编译期常量让 Triton 自动向量化
4.Python wrapper 中用 assert 校验 stride 确保指针算术正确
5.用 -1 哨兵值替代 mask 张量节省显存并兼容 CUDA Graph