thu-ml/TurboDiffusion

一、项目定位

TurboDiffusion 是清华大学机器学习组（THU-ML）开源的端到端视频扩散模型（DiT）推理加速框架。其核心目标是通过算法与系统协同设计，在单张消费级或数据中心 GPU 上实现 100-200 倍的生成加速，同时保持视频质量。项目目前处于早期快速迭代阶段（README 明确指出“checkpoints and paper are not finalized”）。 从架构上看，该项目并非通用的深度学习编译器，而是深度绑定特定模型（如 Wan2.1/2.2、LTX-2）的深度优化方案。其加速管线由三部分构成：1) 算法层：基于 rCM（一致性模型）的步数蒸馏，将采样步数压缩至 1-4 步；2) 注意力层：集成 SageAttention 和 SLA（Sparse-Linear Attention）以缓解长序列（如高分辨率多帧视频）的 Attention 瓶颈；3) 算子层：深度定制的 CUDA/CUTLASS Kernel，包括量化、RMSNorm、LayerNorm 和 GEMM。 针对你的硬约束（仅有 Mac + Colab T4）的切入警告： 这是一个重度依赖高算力 GPU 的项目。根据 setup.py，其自定义 CUDA 扩展 turbo_diffusion_ops 强依赖 CUTLASS，且硬编码了编译目标为 sm_80 到 sm_120a（A100 到 RTX 5090/Blackwell）。这意味着： 1. 你的 Mac 无法编译和运行其核心算子。 2. 你的 Colab T4（Turing 架构，sm_75）默认无法编译该项目。你必须先修改 setup.py 增加 sm_75 的支持，且由于 CUTLASS 某些新特性可能不支持 Turing，你可能需要降级或修改 Kernel 代码才能在 T4 上跑通。 3. 项目依赖 triton>=3.3.0 和 flash-attn，在 Apple Silicon 上极难原生运行。 切入建议：发挥你的 PyTorch 融合与分布式经验。在 Mac 上，你可以通过 Mock 掉 turbo_diffusion_ops 模块，专注于 turbodiffusion/rcm（蒸馏逻辑）、turbodiffusion/inference（推理管线）的纯 Python 逻辑重构，或者研究其最近提交的 FSDP 训练代码。Kernel 级别的优化（如 DiT 算子融合）只能盲写，依赖 CI 或极短的 T4 时间验证。

解决什么问题、给谁用

解决高分辨率（480p/720p）视频扩散模型（如 Wan2.1/2.2）在端到端生成时极度缓慢（如 4549秒）以及显存占用过高的问题。通过步数蒸馏（rCM）、稀疏注意力（SLA）和底层算子量化（Quant Linear），使得在单张 24GB/32GB 显存的消费级 GPU（如 RTX 4090/5090）上实现秒级（如 1.9秒）生成成为可能。

AI 视频生成领域的研究人员、AIGC 应用开发者、以及需要在有限算力（单卡）下部署实时/准实时视频生成服务的模型部署工程师。

同类项目与差别

• TensorRT / TensorRT-LLM：TRT 是通用的底层推理引擎，而 TurboDiffusion 是针对特定 DiT 模型（Wan系列）的端到端框架，包含了模型层的算法修改（rCM 蒸馏、SLA 稀疏注意力），并非纯粹的计算图编译。
• FlashAttention / SageAttention：TurboDiffusion 是这些底层 Attention 算子库的消费者。它将 SageAttention 包装进其 SLA 模块中，结合自身的量化和 Norm 算子，形成完整的 DiT 加速管线。
• DeepSpeed / Megatron-LM：传统框架侧重于大规模集群训练，而 TurboDiffusion 当前的核心亮点是单卡极致推理加速（尽管近期 commit 显示正在加入 FSDP 训练支持）。

核心能力

阶段：早期/实验阶段 (Early / WIP)。README 明确声明 "The checkpoints and paper are not finalized, and will be updated later"，且近期有大量关于引入新模型 (TurboT2AV/LTX-2) 和修复核心逻辑的提交。

技术栈：Python >= 3.9；PyTorch >= 2.7.0；Triton >= 3.3.0；CUDA C++17；CUTLASS；Ninja / Setuptools (构建系统)；flash-attn；einops

规模：包含 2 个主要模块目录（turbodiffusion 157 个文件，turbot2va 139 个文件）。近期提交极其活跃，LTX-2 相关目录有 58 次提交热点，rcm 目录有 15 次。目前 Release 数量为 0。CUDA 算子硬编码支持 5 种架构（sm_80, sm_89, sm_90, sm_100, sm_120a）。

二、架构与代码地图

TurboDiffusion 的架构设计并非一个通用的大模型推理框架（如 vLLM 或 TensorRT-LLM），而是一个深度绑定特定视频生成模型（Wan2.1/2.2, LTX-2）的算法-系统协同优化（Algorithm-System Co-design）工程。对于只有 Mac 和极短 T4 验证时间的你来说，理解其分层逻辑至关重要，因为你必须避开底层强依赖高算力 GPU 的 C++ 扩展，将精力集中在纯 Python 逻辑、Triton 算子以及分布式训练调度的重构上。 从上至下，项目可划分为四个核心层： 1. 接入与推理层 (Inference & Serving Layer)：包含 turbodiffusion/inference 和 turbodiffusion/serve，以及 turbot2va 下的 pipeline。这里是端到端生成的入口，负责解析用户 Prompt、加载 VAE/Text Encoder、初始化 DiT 模型，并执行 1-4 步的极速采样循环。此层完全是纯 Python 逻辑，是你在 Mac 上进行重构、API 封装或端侧部署（如导出 ONNX/CoreML）的绝佳切入点。 2. 算法与模型层 (Algorithm & Model Layer)：包含 turbodiffusion/rcm 和 turbot2va/LTX-2/packages/ltx-core。这是项目的核心价值所在。算法上，它实现了基于 rCM（一致性模型）的步数蒸馏（Distillation）；模型上，它硬编码了 Wan2.1/2.2 和 LTX-2 的网络结构。最近的提交热点（FSDP 鲁棒性优化、离散时间 CM）都在这里。你可以在 Mac 上阅读并优化其 PyTorch 算子融合逻辑，或研究其 JVP（雅可比向量乘积）的实现。 3. 算子与硬件加速层 (Operator & Hardware Layer)：包含 turbodiffusion/ops（CUDA/CUTLASS）和 turbodiffusion/SLA（稀疏线性注意力），以及散落在各处的 Triton kernels（如 fast_norm_kernels.py）。这是你的硬约束红线区。ops 目录下的量化 GEMM 和 Norm 算子硬编码了 sm_80 到 sm_120a，在 Mac 和 T4 上均无法直接编译。你的策略应该是：在 Mac 上通过 Python Mock 掉 turbodiffusion.ops，或者专注于优化那些对硬件要求较低的 Triton 算子（T4 可跑）。 4. 基础设施层 (Infrastructure Layer)：包含 turbodiffusion/imaginaire。这是一个从 NVIDIA 借用/魔改的底层库，负责分布式训练状态管理（FSDP/DDP）、Hydra 懒加载配置解析（LazyConfig）、Checkpoint 读写和回调机制（Callbacks）。

一次调用怎么流过这些模块

以一次典型的 Wan2.1 T2V 推理请求（wan2.1_t2v_infer.py）为例，数据流穿透了从 Python 调度到底层 CUDA 算子的全过程。由于你只有 Mac，理解这个流程有助于你找到可以剥离或 Mock 的环节。 1. 环境与配置初始化：脚本启动后，首先解析命令行参数（如 --resolution, --num_steps, --quant_linear）。如果启用了量化，系统会通过 turbodiffusion/inference/modify_model.py（需验证具体路径，通常伴随模型加载）或直接在模型初始化时，将标准的 torch.nn.Linear 替换为 turbodiffusion.ops 中的自定义量化算子。 2. 模型加载与文本编码：系统加载 Wan2.1 的 VAE 和 umT5 Text Encoder。用户输入的 Prompt 经过 umT5 编码，生成高维的文本条件特征（Conditional Embeddings）。 3. 噪声初始化与采样循环 (核心加速区)： - 根据设定的分辨率和帧数，在 GPU 上初始化纯噪声 Latent Tensor。 - 进入 rCM（一致性模型）的极速采样循环（通常只需 1-4 步，由 --num_steps 控制）。 - 在每一步中，Latent Tensor 和文本特征被送入 Wan_Networks（如 wan2pt1.py）定义的 DiT Block 中。 4. 注意力机制路由 (Attention Dispatch)： - 在 DiT Block 内部，当执行到 Attention 层时，数据流会根据 --attention_type 发生分叉。 - 如果是 original，调用 turbodiffusion/rcm/utils/attention.py，该模块会根据 GPU 架构（SM80/90/100）动态路由到 Flash Attention 3、SDPA 或 xformers。 - 如果是 sla 或 sagesla，数据流进入 turbodiffusion/SLA 模块，执行稀疏线性注意力计算，大幅降低长视频序列的显存占用。 5. 底层算子执行 (Mac 阻断区)： - 在前向传播中，所有的 Linear 层计算会下沉到 turbodiffusion/ops/gemm/gemm.cu（CUTLASS W8A8 GEMM）。 - 所有的 Norm 操作下沉到 rmsnorm.cu 或 layernorm.cu。 - 注意：在 Mac 上，这一步会因为缺少 turbo_diffusion_ops 扩展而崩溃。你必须在第 1 步拦截量化替换逻辑，强制使用原生 PyTorch 算子才能跑通纯逻辑。 6. VAE 解码与输出：经过 1-4 步去噪后，最终的 Latent Tensor 被送入 VAE Decoder，还原为像素级的视频帧，最后通过 imageio 或 ffmpeg 保存为 MP4 文件。

关键类型与函数

扩展点

• turbodiffusion/rcm/networks/：新增模型：在此目录下新建模型定义文件（如 new_model.py），实现 DiT 结构，并在 turbodiffusion/rcm/configs/ 中注册对应的 Hydra 配置文件。
• turbodiffusion/rcm/utils/attention.py：新增 Attention 后端：修改 _get_sdpa_config 函数，根据 device_cc (Compute Capability) 增加新的路由逻辑（如接入新的 Triton Attention 实现）。
• turbodiffusion/rcm/callbacks/：新增训练调度策略/监控：继承 imaginaire.utils.callback.Callback，实现 on_training_step_start 等钩子，并在配置中注册。
• turbot2va/LTX-2/packages/ltx-distillation/src/ltx_distillation/fast_norm_kernels.py：新增 Triton Kernel：在此处编写 @triton.jit 函数，并务必在 Python 包装函数中提供纯 PyTorch 的 fallback 路径，以保证跨平台兼容性。

最近在动的地方

• turbot2va/LTX-2：高达 58 次近期提交，表明团队正在全力集成和优化 LTX-2（音视频双向生成模型），这是目前最活跃的开发分支。
• turbodiffusion/rcm：15 次提交，涉及离散时间 CM（discrete-time CM）的引入和 FSDP 鲁棒性的优化，是算法迭代的核心区。
• turbodiffusion/rcm/utils/fsdp_helper.py：对应提交 'improve FSDP robustness'，说明分布式训练的显存/通信调度是当前的痛点。
• turbodiffusion/rcm/samplers/：对应提交 'simplify timestep sampler'，采样逻辑正在被重构以适应更极端的少步数（1-4步）生成。
• README.md：频繁更新加速指标和模型链接，反映出项目处于早期快速迭代、频繁发版的阶段。

建议阅读顺序

1. README.md
2. setup.py
3. turbodiffusion/inference/wan2.1_t2v_infer.py
4. turbodiffusion/rcm/utils/attention.py
5. turbodiffusion/rcm/utils/flash_attention_jvp_triton.py
6. turbot2va/LTX-2/packages/ltx-distillation/src/ltx_distillation/fast_norm_kernels.py
7. turbodiffusion/ops/core.py
8. turbodiffusion/rcm/trainers/trainer_distillation.py

三、本地跑起来（没有 GPU 的 Mac）

安装

1. git clone https://github.com/thu-ml/TurboDiffusion.git
2. cd TurboDiffusion
3. git submodule update --init --recursive
4. sed -i '' 's/ext_modules=ext_modules/ext_modules=\[\]/g' setup.py
5. sed -i '' '/triton>=3.3.0/d' pyproject.toml
6. sed -i '' '/flash-attn/d' pyproject.toml
7. conda create -n turbodiffusion python=3.12
8. conda activate turbodiffusion
9. pip install -e . --no-build-isolation

哪些路径能真跑

在 Mac (CPU/MPS) 上真正可执行的路径包括：1) turbodiffusion/inference/ 和 turbodiffusion/serve/ 下的纯 Python 推理管线与 TUI 交互逻辑；2) turbodiffusion/rcm/ 下的配置解析（Hydra）、调度器（如 lr_scheduler.py）和纯 PyTorch 模型结构定义；3) turbot2va/LTX-2/packages/ltx-core/src/ltx_core/components/schedulers.py 等纯数学张量计算。只能读代码或依赖 Colab T4 验证的路径包括：1) turbodiffusion/ops/ 下的所有 CUDA/CUTLASS 算子（量化 GEMM、RMSNorm）；2) turbodiffusion/rcm/utils/flash_attention_jvp_triton.py 等 Triton 算子；3) turbot2va/LTX-2/packages/ltx-distillation/src/ltx_distillation/fast_norm_kernels.py（虽有 fallback 逻辑，但强依赖 GPU 环境）。

最小可运行

1. export PYTHONPATH=turbodiffusion
2. turbodiffusion-serve --help
3. wget https://huggingface.co/TurboDiffusion/TurboWan2.1-T2V-1.3B-480P/resolve/main/TurboWan2.1-T2V-1.3B-480P.pth
4. python turbodiffusion/inference/wan2.1_t2v_infer.py --model Wan2.1-1.3B --dit_path TurboWan2.1-T2V-1.3B-480P.pth --resolution 480p --prompt "A test video" --num_samples 1 --num_steps 1 --attention_type original --default_norm

测试

项目未在根目录提供统一的测试入口，但存在局部测试目录，如 turbot2va/LTX-2/packages/ltx-core/tests/ 和 turbot2va/LTX-2/packages/ltx-distillation/tests/。推测使用 pytest 作为测试框架。在 Mac 上，可以通过 pytest turbot2va/LTX-2/packages/ltx-core/tests/ 只运行不依赖 GPU 的纯逻辑子集（如 test_transformer_fusion_helpers.py），大概耗时在几秒到十几秒内（需验证）。

调试

• 日志开关：项目依赖 loguru（见 pyproject.toml），可通过设置环境变量 LOGURU_LEVEL=DEBUG 开启底层详细日志。
• 断点建议：在 Mac 上重构时，建议在 turbodiffusion/inference/wan2.1_t2v_infer.py 的 Pipeline 初始化处打断点，观察模型结构和权重加载逻辑。
• Profiling 入口：项目内置了 turbodiffusion/imaginaire/utils/profiling.py，可用于追踪 PyTorch 算子耗时。
• 环境变量：执行任何脚本前，必须设置 export PYTHONPATH=turbodiffusion，否则会报模块找不到的错误（README 明确要求）。

CI

仓库目录树中未见 .github/workflows/ 目录，说明目前开源版本没有配置公开的 GitHub Actions CI（需验证是否在内部 GitLab 运行）。PR 提交目前主要依赖人工 Review，不会被自动化的 Lint 或单元测试卡住。

坑

• CUDA 架构硬编码：setup.py 中的 cc_flag 写死了 sm_80 到 sm_120a。在 Colab T4 上编译会直接报错，必须手动修改 setup.py 添加 -gencode arch=compute_75,code=sm_75。
• 量化算子硬件限制：README 明确指出量化权重（-quant）适用于 RTX 4090/5090。在 T4 上强行使用 --quant_linear 会因 CUTLASS 算子不兼容而崩溃，必须下载未量化权重并移除该参数。
• Triton 依赖与 Fallback 陷阱：虽然 fast_norm_kernels.py 写了 if not x.is_cuda ... return fallback，但在 Mac 上直接 import triton 就会抛出异常，必须在代码层面 Mock 掉 Triton 的导入逻辑。
• Attention 后端回退 Bug：turbodiffusion/rcm/utils/attention.py 会根据 get_device_cc 分发后端，若无 flash_attn 会回退到 PyTorch SDPA。在 Mac MPS 上，SDPA 对 bfloat16 的支持可能存在缺陷，可能需要强制降级到 float32 或纯 CPU 运行。

四、维护者与社区

目前无正式 Release（releases 列表为空）。提交频率呈脉冲式，集中在特定功能开发期（如 2026年3月和8月有密集提交）。README 明确指出 'checkpoints and paper are not finalized'，属于典型的早期学术开源项目节奏。

流程与 Review 风格

缺乏正式的工程化贡献流程。仓库中没有 CONTRIBUTING.md、CLA/DCO 声明或 pre-commit 配置。贡献者通常直接提交 Issue 或发起 PR（如 PR #133 提交 Triton 算子融合，PR #66 提交 AMD 适配）。

偏向学术团队的实用主义风格。对技术探讨（如 Issue #99 loss 异常、Issue #126 rCM 策略）响应较快（有多次回复）；但对非核心路线图的工程需求（如 Issue #128 多卡并行、Issue #127 NPU 支持）往往暂不响应（0 comments）。（需验证具体 PR Review 细节）

渠道

• https://github.com/thu-ml/TurboDiffusion
• https://huggingface.co/TurboDiffusion

这里的规矩

• 使用中文或英文交流均可（Issue 列表中存在大量中文提问与讨论）。
• 提问或报 Bug 前需明确硬件环境与模型版本的匹配关系（README 强调 H100 需用非量化版，RTX 5090 需用量化版及 --quant_linear 参数）。

维护者现在最想要的帮助

• 多卡并行推理与分布式支持（社区在 Issue #5 和 #128 中多次求助，这正是你擅长的分布式训练性能方向，且可在 Mac 上重构纯 Python 调度逻辑）。
• Triton 算子融合与优化（社区已提交 PR #133 融合 LayerNorm，你可以利用 Colab T4 验证更多 Triton 级别的算子融合，避开无法编译的 CUTLASS 扩展）。
• LoRA 微调与训练管线完善（Issue #93 和 #92 提及，契合你的训练性能背景）。
• 跨硬件平台适配（社区在尝试 AMD PR #66 和 NPU Issue #127 适配，虽然你受限于 Mac，但可以关注纯 Python 层的解耦）。

五、切入方案

它现在缺什么（你无 GPU 也能补）

第一批 PR

怎么知道自己站住了

• Maintainers 接受并合并了你关于 sm_75 编译支持的 PR
• 你在 Mac 上成功跑通了 Mock 掉 CUDA ops 后的纯 Python 推理管线 (即使生成结果是噪声)
• 核心作者在涉及 FSDP 或 Triton 算子的 Issue 中主动 @ 你寻求建议