torchada兼容层实战

一、torchada 是什么？

torchada（torch adapter）是摩尔线程开源的 运行时兼容层：只要在 Python 代码开头 import torchada，原本调用 torch.cuda.* 的 CUDA 代码就会自动映射到 MUSA 上执行，几乎不需要改一行业务代码。

这与我们在上一课学过的 MUSIFY 完全不同 —— MUSIFY 是 源码转换工具，一次性把 .cu → .mu、cudaMemcpy → musaMemcpy；而 torchada 工作在 Python 运行时：它通过 monkey-patch / 符号别名让 torch.cuda 家族 API、CUDAPluggableAllocator、flash_attn 等第三方扩展直接在 MUSA 上工作。

二、安装与第一个示例

pip install torchada
# 或从源码
git clone https://github.com/MooreThreads/torchada.git
cd torchada
pip install -e .

最小示例 —— 下面这段代码原本是 CUDA 上跑的，现在只多了一行 import torchada：

import torchada   # <-- 关键：必须在 import torch 之前或紧挨着
import torch

# 下面这些都是"CUDA 风格"的 API
x = torch.randn(1024, 1024, device="cuda")
y = torch.randn(1024, 1024, device="cuda")
z = x @ y

print(z.device)                 # musa:0 — 被 torchada 接管
print(torch.cuda.is_available())     # True
print(torch.cuda.device_count())     # 与 MUSA 设备数一致
print(torch.cuda.get_device_name(0)) # MTT S4000 / S5000 / …

注意：device="cuda" 没有改。torchada 把 torch.cuda 整个子模块的入口都重定向到了 torch.musa，包括张量创建、设备查询、事件、流、内存管理。

三、torch.cuda.memory 补丁

vLLM、DeepSpeed、bitsandbytes 等框架都会深度使用 torch.cuda.memory 下的 API：

• caching_allocator_alloc / delete
• memory_stats、memory_allocated、memory_reserved
• CUDAPluggableAllocator（自定义分配器 ABI）

torchada 在 commit #51 里完整补齐了这些 API，并为 CUDAPluggableAllocator 加了 alias，使其指向 MUSA 的等价实现。用户代码无需感知。

import torchada
import torch

print(torch.cuda.memory_allocated() / 1024**2, "MiB")  # 已分配
print(torch.cuda.memory_reserved()  / 1024**2, "MiB")  # 已保留

# 自定义分配器（CUDAPluggableAllocator）用法同 CUDA 版本
alloc = torch.cuda.memory.CUDAPluggableAllocator(
    path_to_so="/path/to/my_alloc.so",
    alloc_fn_name="my_alloc",
    free_fn_name="my_free",
)
torch.cuda.memory.change_current_allocator(alloc)

四、与推理框架的集成

4.1 vLLM-MUSA

vLLM 假设所有代码都跑在 torch.cuda 上。torchada 让 vLLM-MUSA 几乎不用 fork vLLM 的核心调度代码 —— 只在插件初始化阶段 import torchada，之后 vLLM 对外暴露的 torch.cuda.current_device() 实际调用的是 torch.musa.current_device()。

4.2 SGLang-MUSA 与 flash_attn

SGLang 依赖 flash_attn，而 flash_attn 在启动时会做几件事：① 加载 CUDA 符号；② 校验 GPU 能力；③ 注册自定义算子。torchada PR #49 通过"hook import 机制 + CUDA 符号别名"解决了所有这些问题，让 flash_attn 不用改一行代码就能在 MUSA 上加载。

# 正确顺序（在任何 SGLang / flash_attn 之前）
import torchada
import torch

import sglang
import flash_attn  # torchada 已经处理好符号

4.3 bitsandbytes / 第三方量化库

大部分量化库都依赖 torch.cuda.current_stream()、torch.cuda.Event 等 API，加上 torchada 后基本可以直接用。遇到 attempted to use cuda kernel on musa tensor 类错误，说明该库在底层调了 .cu 手写 kernel —— 此时需要该库提供 MUSA build 或走 MUSIFY 迁移源码。

五、torchada 的工作原理（进阶）

虽然日常使用只需要一行 import，了解它的底层有助于排障：

六、torchada 的边界

torchada 只处理 Python 层的 API 重定向，这些场景需要额外手段：

• 纯 C++ / CUDA 项目：只有 .cu、.cpp，没有 Python 入口 → 用 MUSIFY 做源码迁移。
• 自定义 CUDA kernel：项目里有手写 .cu 并 jit 编译 → 仍需把 .cu 改写为 .mu。
• 依赖 PTX 内联汇编：PTX 无法直接转 MUSA 汇编，必须手工重写。
• CUDA Graph 特定 API：部分子集目前由 torch_musa 支持，使用前查阅 release note。

七、torchada vs MUSIFY 选型指南

附录：2026-08-10 更新章节 —— torchada 升级到 0.1.79

0. 三个 PR 的共同主线：基线切换前的铺路

单看这三个 PR，会觉得它们很琐碎：一个补 API 映射，一个加环境变量，一个修属性访问。但把它们放在一起，主线非常清晰——PyTorch 生态正在从 2.9 向 2.11 迁移，torchada 必须在基线切换发生前把坑提前填平。

理解这一点很重要，因为它决定了你应该怎么对待这次升级。如果你的环境还稳稳地跑在 torch 2.9.1.post1 上，那么三项改动里只有 TORCHADA_EXCLUDE_DIRS 对你有即时价值（它解决的是构建期头文件冲突，与 torch 版本无关）；另外两项是为未来准备的。反之，如果你已经在测试 torch 2.11，那么这三项你一项都逃不掉。

另一个值得先建立的认知是：torchada 的工作层次比本课正文介绍的要更深。正文重点讲的是 Python 层的 monkey-patch 与符号别名，而这三个 PR 里有两个都涉及 C++ ABI 层的映射。现代推理框架为了避免每个 PyTorch 版本都重编一遍 C++ 扩展，大量使用 libtorch 的稳定 ABI 接口，而这些接口的符号名里硬编码着 cuda。这意味着 torchada 不能只在 Python 层做文章，它必须在链接层也提供对应的 MUSA 符号，否则编译好的扩展根本加载不了。下面逐个 PR 看具体做法。

1. PR #97：vLLM 0.26 兼容与 stable ABI 映射

PR #97（2026-07-30 合入）做了两件事。

第一件是 API 映射：torch.accelerator.get_memory_info 映射到 torch.musa.mem_get_info。torch.accelerator 是 PyTorch 新引入的设备无关加速器接口，上游 vLLM 已开始使用它来查询显存；torchada 补上这条映射，MUSA 设备才能正确回报显存信息，否则调度器会拿到错误的可用显存并做出错误的批处理决策。

第二件更底层：为 vLLM csrc/libtorch_stable/ 中的 cudaDeviceProp、cudaStream_t、cublasHandle_t、aoti_torch_get_current_cuda_stream 提供 MUSA stable-ABI 映射。这里的关键词是 stable ABI——vLLM 为了避免每个 PyTorch 版本都重编译 C++ 扩展，使用了 libtorch 的稳定 ABI 接口，而这些接口的符号名里带着 cuda。torchada 需要在 ABI 层面把它们映射到 MUSA 对应实现，编译好的 vLLM 扩展才能在 MUSA 上加载。

2. PR #96：TORCHADA_EXCLUDE_DIRS 与按 torch 版本分支的 ABI 策略

PR #96（2026-08-03 合入）引入了一个新的环境变量 TORCHADA_EXCLUDE_DIRS，用于保护依赖方的头文件不被 torchada 的映射逻辑改写。它接受两种形式的值：完整路径，或者目录名（例如 torch_musa）。

为什么需要这个开关？torchada 的工作方式是在构建期把嵌套的 torch/cuda.h 之类的包含改写为 MUSA 对应头文件。但如果被改写的目标是 torch_musa 自身的头文件——那本来就是 MUSA 原生实现——改写就会造成重复定义或符号错乱。TORCHADA_EXCLUDE_DIRS 让你显式声明"这些目录不要动"。

# 按目录名排除（最常用）
export TORCHADA_EXCLUDE_DIRS="torch_musa"

# 按完整路径排除，多个用冒号分隔
export TORCHADA_EXCLUDE_DIRS="/usr/local/lib/python3.10/site-packages/torch_musa:/opt/mydeps/include"

# 构建第三方扩展时带上该变量
TORCHADA_EXCLUDE_DIRS="torch_musa" pip install -e ./my-cuda-extension

同一 PR 还实现了按已安装 torch 版本自动选择 stable-ABI 策略的分支逻辑：

此外该 PR 对齐了 MUSA stream 与 BLAS handle 的 shim，并在端到端 pytest 中增加了对 vLLM weak_ref_tensor 的断言——这是一个很好的信号，说明 torchada 的回归测试已经覆盖到 vLLM 的具体用法，而不只是通用 API。

3. PR #98：torch 2.11 下的 torch.cuda.streams shim

PR #98（2026-08-07 合入）修复了一个非常具体但极易撞上的问题。PyTorch 2.11 的 Dynamo guard 会访问 torch.cuda.streams.Stream；而 torch_musa 2.11 的实现位于 torch_musa.core.stream，并不导出 torch_musa.streams。结果是编译版 vLLM 在启动时直接抛 AttributeError。

4. Step-by-Step：升级、验证与回滚

# Step 1 — 记录现状
pip show torchada | grep -E "Version|Location"
python -c "import torch; print(torch.__version__)"
pip freeze > /tmp/before_torchada.txt

# Step 2 — 升级
pip install "torchada==0.1.79"

# Step 3 — 基础验证：一行 import 后 torch.cuda 家族是否可用
python - <<'PY'
import torchada, torch
print("torchada ok")
print("cuda available:", torch.cuda.is_available())
print("device count:", torch.cuda.device_count())
print("device name:", torch.cuda.get_device_name(0))
PY

# Step 4 — 验证 PR #97 的显存接口映射
python - <<'PY'
import torchada, torch
try:
    free, total = torch.cuda.mem_get_info()
    print("mem_get_info ok: free=%.2fGB total=%.2fGB" % (free/2**30, total/2**30))
except Exception as e:
    print("mem_get_info FAILED:", e)
PY

# Step 5 — 验证 PR #98 的 streams shim（torch 2.11 环境）
python - <<'PY'
import torchada, sys
try:
    import torch.cuda.streams as s
    print("torch.cuda.streams ok:", s.Stream)
    print("registered in sys.modules:", "torch.cuda.streams" in sys.modules)
except Exception as e:
    print("streams shim FAILED:", e)
PY

# Step 6 — 构建第三方扩展时排除 torch_musa 目录
TORCHADA_EXCLUDE_DIRS="torch_musa" python setup.py build_ext --inplace

# Step 7 — 回滚
pip install "torchada==0.1.76" --force-reinstall
# 或整体还原：pip install -r /tmp/before_torchada.txt

5. 版本路线提示

6. 何时该升、何时该等

把前面的内容收敛成一个可操作的建议。torchada 与 MATE 不同：它是兼容层而非计算层，升级它不会改变数值结果，只会改变"能不能跑起来"。这个性质决定了它的升级风险模型与算子库完全不同——失败是显式的、立即可见的，而不是隐蔽的精度漂移。

应该立即升级的三种情况：① 你在 torch 2.11 环境下遇到了 torch.cuda.streams 相关的 AttributeError，这是 PR #98 正对的问题；② 你在构建第三方 CUDA 扩展时遇到头文件重复定义或符号错乱，TORCHADA_EXCLUDE_DIRS 是目前唯一干净的解法；③ 你在跟进 vLLM 0.26，需要 torch.accelerator 接口与 stable ABI 映射。

可以暂缓的情况：你的生产环境固定在 torch 2.9.1.post1 + vLLM-MUSA v0.24.0-dev，且没有遇到上述任何一种报错。这种情况下升级的收益接近于零，而任何依赖变动都有非零风险。但建议你在测试环境提前把 torch 2.11 + torchada 0.1.79 的组合跑通，因为基线切换是早晚的事，提前把问题暴露在测试环境总比到时候仓促切换强。

一个容易被忽略的依赖链提醒：torchada 的版本不是孤立的。它同时受三方约束：你安装的 torch/torch_musa 版本、你使用的 vLLM-MUSA 分支的 requirements、以及你自己构建的第三方扩展。升级前请把三者的要求写在一张表里对照，而不是只看其中一个。实践中最常见的翻车方式是：为了解决某个报错单独把 torchada 升到最新，结果它按 torch 2.11 的策略选择了原生 stable ABI，而环境里实际装的是 torch 2.9——于是一个报错变成了另一个更难查的报错。好消息是 PR #96 已经把这个选择自动化了，但你仍需要确认它判断到的 torch 版本确实是你预期的那个——尤其在同一台机器上装了多个虚拟环境时。

本章节小结：torchada 0.1.79 的三项兼容 shim 分别解决三类问题——#97 补齐 torch.accelerator.get_memory_info → torch.musa.mem_get_info 映射，并为 vLLM stable-ABI 中的 cudaDeviceProp/cudaStream_t/cublasHandle_t/aoti_torch_get_current_cuda_stream 提供 MUSA 映射；#96 引入 TORCHADA_EXCLUDE_DIRS 保护依赖头文件，并按 torch 2.9（backport）与 torch 2.11+（原生 stable ABI）分支选择策略；#98 修复 torch 2.11 Dynamo guard 访问 torch.cuda.streams 导致的 AttributeError，采用属性与 sys.modules 双重映射。整体指向同一个方向：torch 2.11 基线正在到来，但尚未合并，生产环境暂缓切换。