SuperAPI AI 工具集

tilekernels-gpu-kernels

通用与其它

查看项目

更新于 10-15 21:05

技能介绍

用途:TileKernels 是面向大语言模型运算的高性能 GPU 内核库,用 TileLang(一种可自动优化的 Python GPU 内核 DSL)编写,目标为 NVIDIA SM90/SM100(Hopper/Blackwell)架构,力求接近硬件性能上限。

主要能力:包含多个模块。moe 提供 top-k 专家选择与打分、token 到专家映射、融合扩展/归约、权重归一化;quant 提供按 token、按块、按通道的 FP8 量化以及融合 SwiGLU + FP8 量化;transpose 提供批量矩阵转置;engram 提供融合 RMSNorm 的前向、反向与权重梯度归约;mhc 提供 Sinkhorn 归一化与 mix 拆分/应用;modeling 提供基于 torch.autograd.Function 的高层封装(如 EngramGateLayer、MHCPipeline);torch 提供 PyTorch 参考实现用于校验;testing 提供测试与基准工具。

使用方式:需 Python 3.10+、PyTorch 2.10+、TileLang 0.1.9+、CUDA Toolkit 13.1+ 及 SM90/SM100 GPU。开发安装用 pip install -e ".[dev]",发布安装用 pip install tile-kernels。典型用法是导入对应模块函数,传入 CUDA 张量并获取结果,例如 topk_gating(logits, top_k=8)、per_token_cast_fp8(x)、batched_transpose(x)、engram_gate_forward(hidden, weight)、sinkhorn_normalize(conn_weights, num_iters=20)。测试可用 pytest tests/transpose/test_transpose.py -n 4,加 --run-benchmark 输出性能数据,全量压力测试用 TK_FULL_TEST=1 pytest -n 4 --count 2。

适用场景:MoE 前向分发、激活与权重量化、批量转置、engram 门控训练、Manifold HyperConnection 等 LLM 内核级优化与验证。

限制:仅支持 SM90/SM100,不能在 Ampere/Ada 等旧 GPU 上运行;需匹配 TileLang 与 CUDA 版本;内核针对特定 tile 尺寸调优,可能遇到共享内存不足;基准测试需显式开启。其余细节暂未提供。