Skip to content

TensorPlay v1.0.0

Choose a tag to compare

View all tags
@github-actions github-actions released this 10 Sep 20:18
· 3571 commits to main since this release
v1.0.0
25b8c31

TensorPlay 1.0.0 Release Notes

亮点 Highlights

  • autograd.Function 完整实现：ctx-first backward 契约、grad_fn/ctx
    双侧 hooks 与引擎协同逐项核验，并新增原生 forward-mode AD。
  • tensorplay.library 自定义算子体系：覆盖 custom_op/triton_op/
    register_fake/register_vmap/opcheck 等全量公开面，调度层净开销 2.3µs；
    tile-lang 一等支持、tvm-ffi 零拷贝直通。
  • tensorplay.compile（Stax 编译器）收口：CUDA 减归约/逐元素链完成
    11 组负载验证（RTX 4090 D）。
  • CUDA graphs 完整实现：条件节点（if/while 图）、多设备并发捕获、
    make_graphed_callables、共享内存池，外加单调用回放快路径。
  • RNG 逐位稳定：CPU 原生随机引擎 + CUDA 设备端 Philox4_32_10，
    同种子逐值一致，state 字节布局固定。
  • 构建/CI/发布全链路收口：scikit-build-core PEP 517、统一版本链路
    与 CI 形态、CPU 走 PyPI、CUDA 走自建 PEP 503 索引，并新增 nightly 预览通道。

不兼容变更 Backwards Incompatible Changes

CUDA graphs

旧 _C.cuda_graph_* 自由函数绑定全部移除

原生层重写为 tensorplay::cuda::graph::CUDAGraph 类，
Python 层改用 tensorplay/cuda/graphs.py 的原生类封装：capture_end() 即完成
instantiate，replay() 直接持 exec 指针发射，支持 capture_error_mode、
graph_pool_handle() 共享池与 debug_dump()。依赖旧整型句柄接口的代码需迁移。

图 API

Graph.output() 改为单例替换语义

先擦旧 output 再建新。此前 _interpret 取首个 output 而 recompile 取末个，
双 output 时可能返回不同值；现在全链路一致。

tensorplay/fx.py 移除

公共图 API 迁移为 tensorplay/graph.py 门面（实现仍在 tensorplay/_stax/graph.py）。

构建

自建构建脚本删除，统一 PEP 517

rebuild.py（Windows 本地构建编排）与 release.py（cibuildwheel 发布编排）移除。
构建一律走 pip install . / python -m build --wheel（scikit-build-core），
python setup.py install 不再支持。USE_CUDA=OFF pip install .、
CMAKE_CUDA_ARCHITECTURES="70;75;86" pip install . 等环境变量直通生效。

版本号改为构建期生成

tensorplay/version.py 不再硬编码，由 tools/generate_tensorplay_version.py 在构建时
从 TENSORPLAY_BUILD_VERSION/BUILD_NUMBER → PKG-INFO → version.txt + git SHA
三级生成。main 分支版本形如 1.0.0a0+git<sha7>。

弃用 Deprecations

无。

新特性 New Features

Autograd

  • autograd.Function 完整实现：ctx-first backward、grad_fn/ctx 双侧 hooks、
    to_save/metadata/next_functions/mark_dirty/once_differentiable/
    InplaceFunction/NestedIOFunction、梯度数校验含"多余 None 截断"规则。
  • 新增 custom_function_apply 单入口融合：节点创建/unpack_input/forward/
    setup_context/wrap_outputs 一次 C 穿越完成。
  • 引擎 InputMetadata 原生化：Edge 记录 shape/dtype/device，Engine 在 C++ 层
    零填充缺失梯度槽。
  • saved_tensors_hooks（pack 在 save 时、unpack 随 ctx 快照）。
  • forward-mode AD 原生核心：tensorplay/autograd/forward_ad.py
    （level 栈/DualTensor/种子算子 JVP）。

自定义算子

  • tensorplay/library.py：custom_op/triton_op 装饰器、
    register_kernel/register_fake/register_autograd/register_vmap/register_autocast、
    Library（DEF/IMPL/FRAGMENT）、infer_schema、opcheck 四项检测、
    set_kernel_enabled/get_kernel、顶层 define/impl/impl_abstract。
  • tile-lang 一等支持：tile_lang_op + wrap_tilelang（同契约超集扩展）。
  • tvm-ffi 直通：TP Tensor 实现 DLPack 协议，tvm_ffi JIT/AOT 两路
    tvm::ffi::TensorView 参数零拷贝直通。
  • Stax 原生图新增 custom_op 节点：编译图中原生执行（非解释器），
    autograd 可穿透编译图求梯度。
  • Composite 分发键：形状/视图组合算子批
    在 CUDA 上不再 Kernel not found。

编译器（tensorplay.compile / Stax）

  • 控制流静态特化：if x.shape[0] > 2、range(x.ndim)、len(x) 切片等元数据
    驱动的 Python 控制流可 fullgraph 编译（数据依赖分支仍拒绝）。
  • 守卫版符号形状：被读形状的参数按尺寸精确重特化，未触碰参数保持 dynamic 通配。
  • AOT 双 partitioner：partition_default（标签切分）与 partition_min_cut
    （自研 Edmonds-Karp 最大流/最小割，支持 memory_budget）。
  • Pass 体系：PassManager 至不动点执行，
    内置 DeadCodeElimination/ConstFold/ShapeProp。
  • tensorplay.graph 门面：语义化节点命名、erase_node/
    replace_all_uses_with、is_leaf_module 钩子、concrete_args、
    create_feature_extractor、to_dot()/draw() 可视化。
  • Triton 代码生成：四类 launcher（single/dims/split/pw）+ 静态 fast-launch 路径。

CUDA

  • CUDA graphs 条件节点（if/while 图，CUDA ≥ 12.4）、多设备并发捕获、
    make_graphed_callables/make_graphed_autograd_function。
  • 共享内存池：graph_pool_handle()，多 graph 共享私有池，引用计数释放，
    池内相邻空闲块合并防碎片。
  • stage_and_launch 回放快路径：全部输入单次 Python→C++ 调用内暂存并发射。
  • 原生 memory_stats(device)：segments/free_blocks/graph_pools/capturing 等，
    接入 tp.cuda.memory_stats*。
  • OOM 防御梯队：分配失败先排空跨流 pending 块，再 flush 缓存重试，
    报文附 reserved/allocated/free 统计。

内核

  • LayerNorm CUDA 前向/反向自定义 kernel（不再依赖 cuDNN）：Welford warp 归约、
    4 宽向量化、半精度 fp32 累加、确定性反向归约。
  • argmax packed warp-shuffle：(value, index) 打包 u64 单次整数 max，
    NaN/±0/首现 tie 位等价。
  • Loss 半精度扩展（nll_loss/mse_loss_backward → Float64/Half/BFloat16）、
    二元算子与半精度一元浮点 4 宽向量化快路径。

RNG

  • CPU：原生随机引擎 + 分布变换公式（Box-Muller 缓存、Fisher-Yates
    randperm、poisson 等），normal_fill 走 AVX2 向量化路径，大张量 randn 逐位一致。
  • CUDA：curand 设备端 Philox4_32_10 + host (seed, offset) 预留。
  • Python：get_rng_state/set_rng_state/fork_rng、Generator.get_state/set_state
    （5056 字节 POD 布局）。

改进 Improvements

  • CUDAStream 热路径：设备数进程级缓存、每线程当前流改扁平数组（消除哈希查找）。
  • 当前设备线程级缓存：currentDevice()/getCurrentCUDAStream(-1) 降为一次 TLS 读。
  • 自定义算子热路径：捕获守卫无 trace 时跳过 proxy 扫描、内核选择无锁化。
  • 分布算子 dtype 补齐：Half/BFloat16、random_/geometric_ 全整型谱 + Bool、randperm Int32。
  • Tracer：call_module 节点 + 限定模块路径消歧。

问题修复 Bug fixes

  • DLPack：CPU/-1 设备归一为 device_id=0（此前 -1 使 tvm_ffi 走 workspace 复制路径段错误）。
  • CUDA scatter 分发：if (Add) 改 if constexpr，assign 模式全 dtype 实例化。
  • PyNode backward 输入槽语义修正为前向输出数。
  • eager 自定义算子调用丢弃 kwargs 的 bug。
  • is_autocast_key 恒假 / is_autograd_key 上界吞新键两个区间谓词。
  • persistent split 代码生成修复（误发 classic tail、死 preamble、stride 整除越界等）。

性能 Performance

测量方法均为静默机器 min-of-window；硬件标注于各条目。

  • 自定义 Function 端到端 fwd+bwd 40-45µs（backward 腿 13.5µs）。
  • 自定义算子调度层净开销 2.3µs，端到端调用 ~2.9× 于 Python 装饰器路径。
  • tensorplay.compile（RTX 4090 D，iters=200）：11 组负载基准，
    dims/epilogue 链 2.25-2.51x，
    pw 链 0.98x（内核已验证，残差为 Python 启动路径）。
  • kernel_launch CPU 开销 45 → 9.7 µs/call；compiled 全链 51.3 → 20.2 µs/call。
  • sum full 16M：62.5 → 41.0 µs。

文档 Documentation

  • docs/source/library.md 增补 tvm-ffi 用法章节。
  • CONTRIBUTING.md 构建说明更新为 PEP 517 命令。
  • README（中英）新增 CUDA 索引与 nightly 通道安装说明。

开发者 Developers

  • 构建统一 scikit-build-core（MAX_JOBS 伞形旋钮、editable redirect 模式、
    sdist 裁剪），cmake/EnvVarForwarding.cmake 实现环境变量直通。
  • 版本链路 provider 化（PEP 440 校验 + sdist 一致性断言）。
  • CI 重排：可复用 _binary-* 工作流 + pull/trunk/publish/lint 编排，
    BUILD_ENVIRONMENT/SHA1 环境变量与显式 permissions。
  • 发布：CPU 轮子 PyPI trusted publishing；CUDA 轮子 GitHub Release +
    pypi-pages 分支静态 PEP 503 索引（Cloudflare Pages）；新增手动 nightly
    预览通道（X.Y.0.dev<日期>+cuXXX 版本规则）。
  • Release notes 机制：PR 按路径自动打 release notes: * 标签
    （.github/labeler.yml），发布文案存于 docs/release-notes/。

安装 Installation

# CPU 轮子（PyPI）
pip install tensorplay==1.0.0

# CUDA 13.0 轮子
pip install tensorplay==1.0.0 \
  --index-url https://download.tensorplay.cn/whl/cu130/ \
  --extra-index-url https://pypi.org/simple
Assets 46
Loading