Tlaloc：Kotlin 编译期自动微分与 GPU 推理服务

核心能力与快速开始

★ 0 · Kotlin

能做什么

编译期自动微分：K2 编译器插件把 grad { }、grad2/grad3、valueAndGrad、jvp、vjp、jacobian、jacobianReverse、hessian 等变换在编译时转为梯度代码，支持 if/when/for 循环微分、变换嵌套以及 customVjp/customJvp 自定义求导规则；可将导数打印为 Kotlin 源码。类型化张量：在类型中携带秩、dtype 与轴名称，形状与可微性错误在调用处按文件、行、列报告。算子覆盖逐元素、广播、归约、形状操作、matmul、conv2d、池化、softmax、embedding、损失、batch norm，以及 lgamma/digamma/polygamma/integral 等特殊函数，随机数为与 JAX 逐位一致的 stateless threefry-2x32，另支持 rank-2 CSR 稀疏与稀疏×稠密 matmul（仅 CPU）。训练：:nn 提供不可变层与纯优化器，层包括 Dense、Conv2d、MaxPool、AvgPool、BatchNorm、Dropout、Embedding、EmbeddingBag、GRU、Flatten，优化器包括 SGD、Momentum、RMSprop、Adam，并有学习率调度、梯度裁剪、bf16 混合精度与可逐位恢复的检查点（单个 safetensors 文件）。推理服务：可从 Hugging Face 检查点导出由 StableHLO prefill/decode 程序、safetensors 权重与 manifest 组成的服务制品，由 NVIDIA Triton（libtriton_tlaloc.so 后端）、vLLM 平台插件或通过 ctypes 调用 PJRT C API 的 Python 运行时加载，均不依赖 JVM；服务特性包括 paged attention、滑动窗口 KV cache、分块与批量 prefill、原地 KV cache 更新、连续批处理解码、CUDA 共享内存输入输出与无状态模型动态批处理。已附带 TinyLlama-1.1B、Qwen3-0.6B、Muse Glimmer 30B（纯文本）等模型示例。

适用场景

在 Kotlin 项目中编写可微数组程序并自动获得梯度；训练神经网络模型（如 MNIST 示例、GPU 训练示例）；把训练好的模型或 Hugging Face 检查点导出为 StableHLO 服务制品，在 Triton 或 vLLM 上做 LLM 推理（triton-llm 示例可导出 Qwen3-0.6B 并启动 Triton 流式对话）；可微物理模拟等需要穿过循环求导的场景（differentiable-physics 示例对 38 步 Euler 积分器求导）；通过类型中的轴名称在编译期发现形状错误（named-indices 示例）；查看并复用可读的导数 Kotlin 源码（readable-gradients 示例）。

项目形态

开发框架

部署方式

以 Maven Central 上的库和 Gradle 插件形式使用（插件仅发布在 Maven Central，需在 pluginManagement 中添加 mavenCentral()）；构建时可运行 ./gradlew publishToMavenLocal 后配合 mavenLocal()。GPU 执行、Triton 服务需相应硬件与驱动；具体部署拓扑与运维细节未说明。

主要依赖

Kotlin 2.4.20 – 2.4.29（其他版本编译期拒绝，Kotlin 2.3.x 需用 0.1.0-alpha01）；构建使用 grad { } 的代码需 JDK 25，运行需 JDK 21，GPU 执行需 JDK 25；GPU 需 NVIDIA 驱动与 PJRT CUDA 插件（文档举例在虚拟环境中 pip install "jax[cuda12]"）；Triton 服务需 Docker 与 nvcr.io/nvidia/tritonserver:25.11-py3 镜像。

使用限制

处于 alpha 阶段（0.1.0-alpha02），名称与签名可能在不经弃用周期的情况下变更；CPU 解释器面向正确性而非速度，性能数据来自 GPU 路径；项目本身不生成 GPU 代码，由 XLA 与 IREE 生成，KPTX（PTX 内核 DSL）可替换已识别算子但默认不启用任何内核；每个 Tlaloc 版本只支持一个 Kotlin 版本（插件依赖编译器内部实现）；无 Python API，其他语言需使用 StableHLO 制品；尚未编写或从未运行的能力包括 TPU、多主机训练、单个 Triton 模型使用多 GPU、抢占运行中的序列、Muse Glimmer 图像输入；平台支持上 Windows 未测试，Linux x86_64 的 GPU 与 Triton 服务为“预期可用但未测试”，Google TPU 为“已编写但从未运行”；全部 GPU 结果来自单台 NVIDIA GB10（Blackwell，aarch64），CI 无 GPU。可选依赖 Symja 为 LGPL-3.0（上游仓库 license.txt 为 GPL-3.0），仅用于对循环次数非常量场景求导且默认不在依赖图中。文档未说明任何 API 成本或商用许可限制，许可证为 Apache-2.0。

许可证

Apache-2.0

最近推送

2026-09-27 16:16

资料检查：2026-09-28 07:00 · 卡片整理：2026-09-30 07:00
卡片由 AI 根据仓库资料整理，未进行安装或运行验证 · 本次更新未完成，保留上次成功内容