作为一款可扩展的高性能推理 DiT 引擎,TACO DiT 能够兼容 GPU 以及紫霄硬件,其核心能力概括如下:
·支持主干网络混合并行
TACO DiT 提供六种基础并行方法,其中 Data、Sequence、PipeFusion 和 CFG 可以灵活组合使用,以适配底层硬件并扩展至大规模集群。此外,DistriFusion 和 TensorParallel 两种并行策略则作为独立方案运行。
·支持解码模块 VAE 并行
针对扩散模型后处理中的解码 VAE 模块在生成高分辨率图像时可能出现的 OOM 问题,TACO DiT 开发了 Patch Parallel 版本的 VAE,有效提升了处理效率。
·支持多种模型
目前,TACO DiT 已兼容多种主流文生图和视频生成模型,包括 Stable Diffusion、Flux 以及 CogVideoX 等。
·兼容其他加速模块
TACO DiT 能够与多种计算加速模块协同工作,例如 FlashAttention、TensorRT 和 torch.compile 等算子及编译加速方案,同时支持在公有云和私有化环境中部署和使用。