vipshop/cache-dit

A PyTorch-native inference engine with cache, parallelism, quantization and cpu offload for DiTs.

📚 Cache‑DiT란 무엇인가

Cache‑DiT는 이미지, 비디오, 오디오, 3D 생성에 사용되는 확산 기반 트랜스포머 모델(DiT)을 가속화하는 PyTorch 네이티브 추론 엔진입니다. Hugging Face의 Diffusers 라이브러리 위에 구축되었으며, 네 가지 주요 성능 기법을 추가합니다.

기능 설명 중요성
하이브리드 캐시 중간 키/값 텐서(DBCache, TaylorSeer, SCM 등)를 저장하여 이후 확산 단계에서 재계산 없이 재사용할 수 있게 합니다. 최대 약 9배 빠른 생성.
병렬 처리 컨텍스트, 텐서, 2D/3D, 텍스트 인코더, VAE, ControlNet을 위한 전용 병렬 처리. 대규모 모델을 여러 GPU 또는 GPU+CPU 구성에서 실행 가능하게 합니다.
양자화 8비트 및 4비트 가중치/활성화(W8A8, W4A4)를 지원합니다. 품질을 유지하면서 메모리와 계산을 줄입니다.
CPU 오프로드 모델의 일부를 거의 오버헤드 없이 CPU에 버킷별 레이어 단위로 오프로드합니다. 제한된 GPU 메모리에서 더 큰 DiT를 실행할 수 있습니다.

이 엔진은 컴파일 도구(예: torch.compile)와 호환되며, NVIDIA GPU, AMD GPU, Ascend NPU에서 작동하고 SGLang Diffusion, vLLM‑Omni, TensorRT‑LLM, ComfyUI, Jetson‑Containers 등의 기존 에코시스템에 플러그인할 수 있습니다.


🚀 할 수 있는 일

  • 40개 이상의 DiT 파이프라인 제품군(120개 이상의 구체적 변형)을 실행하며, FLUX, StableDiffusion‑3, PixArt‑Sigma, CogVideoX, StableAudio 등의 인기 모델을 포함합니다.
  • 캐시+컨텍스트 병렬 처리+torch 컴파일을 사용하면 9배 속도 향상을 제공합니다(저장소의 벤치마크 이미지에 표시됨).
  • GitHub Copilot/Claude-Code 에이전트를 위한 **"모델 통합 SKILL"**을 제공하여 단일 명령으로 새 DiT 파이프라인을 쉽게 추가할 수 있습니다.
  • 모든 nn.Module에 대해 양자화 및 레이어 단위 오프로드를 지원하므로 새 DiT가 자동으로 혜택을 받을 수 있습니다.

🛠️ 시작하기 (빠른 시작)

# PyPI에서 패키지 설치
pip install cache-dit

# 예: 캐시 및 병렬 처리를 사용하여 FLUX 모델 실행
python -m cache_dit.run \
    --model flux \
    --cache on \
    --parallelism context \
    --quantize w8a8

README는 cache‑dit.io의 자세한 문서를 가리킵니다:

  • 캐시 API – 다양한 캐시 백엔드를 활성화/비활성화하는 방법.
  • 컨텍스트 병렬 처리 – 확산 시간 단계를 GPU 전체에 분할.
  • 텐서 병렬 처리 – 트랜스포머 가중치를 분할.
  • 양자화 및 오프로드 – 저정밀 추론 플래그.

🤝 커뮤니티 및 에코시스템

Cache‑DiT는 성장하는 프로젝트 목록과 이미 통합되어 있습니다:

  • Diffusers (공식 Hugging Face 최적화 가이드)
  • ComfyUI (확산용 비주얼 노드 기반 UI)
  • SGLang, vLLM‑Omni, TensorRT‑LLM (대규모 서빙 프레임워크)
  • Jetson‑Containers (엣지 디바이스 배포)
  • xDiT, MindIE‑SD, SD.Next, stable‑diffusion.cpp 등의 다양한 오픈소스 파이프라인.

이러한 통합은 얇은 어댑터 레이어를 추가하는 풀 리퀘스트가 대부분이며, Cache‑DiT를 최소한의 변경으로 기존 코드베이스에 통합할 수 있음을 보여줍니다.


📄 라이선스 및 인용

  • Apache‑2.0 라이선스(배지에 표시됨).
  • 학술 사용을 위한 BibTeX 형식의 인용이 제공됩니다.

결론: Cache‑DiT는 확산 트랜스포머 모델을 위한 진지한 프로덕션 등급 가속 레이어입니다. 이미 Diffusers를 사용하고 있으며 모델을 다시 작성하지 않고 속도를 높이거나 메모리를 줄이거나 이기종 하드웨어에서 실행하려는 개발자를 대상으로 합니다.

관련

  • 프로젝트
  • 프로젝트
  • Dispatch
  • 프로젝트
  • 프로젝트