MUSA에서 MAGI-2 Preview, vLLM-Omni로 단계당 565ms 달성

TL;DR

Sand AI의 114B 오디오-비디오 생성 모델인 MAGI-2 Preview가 이제 torchada 호환성 레이어를 통해 Moore Threads MUSA GPU에서 디노이징 단계당 565ms(8×MTT S5000, 272p, 정적 컴파일)로 실행되며, 이는 첫 번째 작동 버전(단계당 15.6초) 대비 27.6배 개선된 수치입니다. 이러한 속도 향상은 호스트 동기화 제거, MoE 커널 융합, mHC 잔차 재작성, 전문가 병렬 샤딩 최적화를 수행한 일련의 PR 덕분입니다.

1. MAGI-2 Preview를 독특한 추론 워크로드로 만드는 요소

단일 스트림 DiT 백본은 텍스트, 비디오 잠재, 오디오 잠재 토큰을 하나의 자기 주의 시퀀스로 패킹합니다. 처음과 마지막 4개 레이어만 모달리티별 가중치를 가집니다.
adaLN 없음 – 타임스텝은 각 토큰의 처음 64개 채널에 직접 인코딩되어 비디오와 오디오 타임스텝을 독립적으로 처리할 수 있습니다.
멀티헤드 MoE(레이어 2-37)는 각 3072차원 토큰을 12개의 256차원 헤드로 분할합니다. 각 헤드는 256개 전문가(SwiGLU7) 중 상위 6개로 라우팅되며, 공유 전문가 1개와 모달리티별 공유 전문가 3개가 추가됩니다. 라우팅된 전문가는 매개변수의 약 95%를 차지합니다.
mHC 하이퍼 연결은 4-스트림 잔차(총 너비 12288)를 제공하며, 20회의 Sinkhorn 반복으로 계산된 이중 확률 4×4 혼합 행렬을 사용합니다.
샘플링은 100개의 비증류 단계를 실행하며, 각 단계는 두 개의 CFG 분기를 사용하여 전체 10초 비디오를 한 번에 디노이징합니다.

지표	값
총 매개변수	114B (토큰당 약 6B 활성)
백본	40레이어 DiT, hidden 3072, 24헤드×128
MoE	12헤드×256전문가, 헤드당 top-6
잔차	mHC, 4스트림, 20 Sinkhorn 반복
출력	10초, 125프레임, 12.5fps, 44.1kHz 스테레오
지원 해상도	272p (448×256) 및 540p (896×512)

다른 오픈 비디오 모델과 비교하여 MAGI-2 Preview는 토큰별 MoE 라우팅, mHC 잔차, 공동 오디오-비디오 생성을 단일 자기 주의 스트림에서 결합한 유일한 모델입니다.

2. torchada로 MUSA에서 CUDA 우선 모델 실행

torchada는 CUDA 스타일 임포트(torch.cuda, 장치 문자열, NCCL, CUDA Graph)를 가로채어 MUSA 등가물(torch.musa, MCCL, MUSAGraph)로 매핑합니다. 매핑은 의도적으로 정직하게 수행되므로(torch.cuda.is_available()은 False 유지) 플랫폼 검사는 torch.version.musa를 사용해야 합니다.
코드 재사용: eager PyTorch, torch.compile/Inductor, 업스트림 Triton 커널(융합 BF16 MoE, SwiGLU7, mHC)은 MUSA에서 변경 없이 실행됩니다. 런치 구성과 일부 플랫폼별 경로만 다릅니다.
주요 적응 사항:
- torch.isfinite는 MUSA에서 abs()<inf 표현식으로 대체되었습니다(torchada #120).
- Int32 라우팅 카운트(scatter_add_)는 MUSA에서 int64를 대체하여 호출당 약 39ms 절감(PR #8496).
- GEMM 피연산자 레이아웃이 스트림 메이저로 재정렬되어 단계당 약 80회 복사 제거(PR #8510).
- Triton-MUSA에는 BF16 원자 연산이 없으므로 융합 BF16 MoE 경로가 항상 선택됩니다.

"torchada는 MUSA 텐서에 대해 is_cuda를 true로 만들어 CUDA 스타일 코드가 평소처럼 실행되도록 합니다." – torchada 문서.

3. 최적화 타임라인 (단계당 지연 시간)

단계	단계당 시간 (ms)	주요 변경 사항
초기 부팅 (9월)	15,600	8×MTT S5000 기준
호스트 동기화 제거	8,800	`.item()` → `.tolist()`, all-gather 메타데이터 감소
MoE 융합 GEMM (MATE)	2,390	그룹화된 BF16 GEMM이 전문가별 Python 루프 대체
융합 W13 + SwiGLU7 Triton	954	게이트/업 + 활성화를 위한 단일 커널
Int32 라우팅 카운트	756	PR #8496
K-차원 타일링 (BLOCK_SIZE_K=64)	734	PR #8506
mHC 요소별 재작성	734 → 734 (기준) → 733.6	PR #8510 (FP32 수축)
Sinkhorn 단일 커널	–6.6	PR #8594
W13 누산기 분할	590–591	PR #8855 (커널 38–40% 빠름)
전체 PR 스택 (10월 19일)	565	위 모든 PR 적용, 정적 컴파일

27.6배 전체 속도 향상(15.6초 → 565ms)은 동일한 하드웨어와 워크로드(272p, 100단계, T2VA 프롬프트, SP8×CFG1+EP4)에서 측정되었습니다. 각 항목은 자체 기준선에 대해 평가되었으며, 숫자는 행 간 누적이 아닙니다.

4. 세부 기술적 개선 사항

4.1 호스트 측 동기화

전문가별 .item(): 일괄 .tolist()로 대체되어 단계당 약 3.4초 절감.
All-gather 토큰 수 메타데이터: 단계당 39회에서 3회로 줄여 약 61ms 절감.
torch.isfinite: MUSA에서 abs()<inf로 재구현하여 단계당 약 1.8초 절감.

4.2 MoE 커널 융합

MATE 그룹화 GEMM은 12×256차원 전문가 GEMM을 단일 BF16 GEMM으로 축소(21.9ms → 0.69ms).
융합 W13 + SwiGLU7 Triton 커널은 게이트, 업 프로젝션, 활성화를 결합하여 MoE 병목을 약 1.5초에서 1초 미만으로 줄임.
가중치 레이아웃 사전 인터리빙은 호출당 전치를 제거하여 단계 시간의 약 7.5% 절감.
Int32 라우팅 히스토그램은 라우팅 단계에서 약 1.2초 절감.
K-차원 타일링(BLOCK_SIZE_K=64)은 약 2%의 작은 개선.
결정적 라우트 구성(PR #8592)은 약 12ms 추가.

4.3 mHC 잔차 재작성

작은 K bmm을 FP32 요소별 수축 및 스트림 메이저 레이아웃으로 대체하여 단계당 약 250ms 절감(≈25%).
20회 Sinkhorn 반복을 단일 Triton 커널로 통합(-6.6ms).

4.4 어텐션

Torch 구현에서 MATE FlashAttention-3(FA3)로 전환하고 학습된 LSE 싱크를 사용하여 SP4×CFG2 구성에서 16.8% 감소.
RoPE, 모달리티 그룹화, 싱크-LSE 메타데이터 캐시(-6ms).

4.5 전문가 병렬 샤딩

헤드 샤딩 EP(EP4)는 시퀀스 병렬 처리와 분리되어 가장 빠른 구성(565ms)을 제공. EP8은 샤드가 많아 통신 오버헤드가 증가하여 더 느림(633ms).
Ulysses 송수신 주변 버퍼 복사 제거로 추가 10.6ms 절감(PR #8595).

4.6 컴파일

지역 torch.compile은 eager 실행 대비 단계당 시간을 약 16% 줄임.
동적 컴파일에서 비토큰 차원 고정은 1.4ms 커널을 55µs로 축소(PR #8627).
MUSA에서 동적 컴파일은 올바른 출력을 위해 Sinkhorn 수정(#8594)이 필요.

5. 측정된 성능 결과

5.1 PR 트리 종단 간 타이밍 (정적 vs 동적)

병렬 구성	최종 PR 라운드 전 (정적)	정적 컴파일 (ms/단계)	동적 컴파일 (ms/단계)
SP8×CFG1+EP4	714/715	665.7/669.9	678.3/678.9/683.7
SP4×CFG2+EP4	732.6	690.8	–
SP4×CFG2	783.1	729.9	–

565ms 수치(섹션 4)는 2026년 10월 19일에 측정된 최신 트리에서 나온 것으로, 다른 코드 베이스와 세션을 사용하므로 위 행과 직접 비교할 수 없습니다.

5.2 W13 누산기 분할 A/B 테스트 (PR #8855)

트리	정적 컴파일 (ms/단계)	출력 해시 (오디오 / 비디오)
대조군 (분할 없음)	690.8/691.1	`32f1b4fca4` / `44b9747e9e`
+ #8855 (분할)	591.3/590.4	동일

14.5% 감소는 비트 단위로 동일합니다. 동일한 SHA-256 해시가 수치적 동등성을 확인합니다.

5.3 최신 트리에서 전문가 병렬 정도 스윕

구성	평균 단계당 시간 (ms)	EP4 대비
SP8×CFG1, EP 꺼짐 (레거시)	697.8	+23.5%
SP8×CFG1+EP4	565.0	기준
SP8×CFG1+EP8	633.4	+12.1%

세 구성 모두 동일한 오디오 및 비디오 해시를 생성하여 속도 차이가 순전히 통신/계산 트레이드오프임을 확인합니다.

6. 샘플 출력 및 품질 노트

벤치마크 프롬프트(골든 리트리버, 272p, SP8×CFG1+EP4)는 단계당 565ms로 10초 비디오를 생성했습니다. 클립은 블로그 자산에서 확인할 수 있습니다.
공식 Sand AI 프롬프트(커피숍, 애니메이션, 클레이 애니메이션)는 T2VA 및 I2VA 모드로 실행되었습니다. 모든 샘플은 텍스트 설명과 일치하고, 0이 아닌 오디오를 포함하며, 커피숍 클립에는 들리는 영어 대화가 포함됩니다.
540p 미리보기는 정적 컴파일에서 약 3.6초/단계, 동적 컴파일에서 약 3.9초/단계로 실행됩니다. 이 수치는 다른 코드 트리와 컴파일 모드를 사용하므로 272p 타이밍과 직접 비교할 수 없습니다.
아직 정량적 다중 시드 품질 평가는 수행되지 않았습니다. 시각적 및 청각적 충실도는 수동으로 검증되었습니다.

7. 업스트림 상태 및 호환성

7.1 PR 현황 (2026년 10월 28일 기준)

PR	기능	상태	측정된 단계당 영향
#7206	융합 BF16 MoE (W13)	병합됨	–
#7243	구성 가능한 Dynamo 재컴파일 제한	병합됨	–
#7252	융합 SwiGLU7	병합됨	–
#7261	융합 mHC 사후 잔차	병합됨	–
#8496	int32 라우팅 카운트	병합됨	1,952→756ms
#8497	W13 단일 복사	병합됨	–
#8498	MUSA 지원 (FA3 varlen + sink)	병합됨	–
#8506	MUSA에서 MoE GEMM용 BLOCK_SIZE_K=64	병합됨	754.5→733.9ms
#8510	mHC FP32 수축 및 스트림 메이저 레이아웃	병합됨	984.7→733.6ms
#8511	SP와 독립적인 헤드 샤딩 EP	열림	864.6→717.5ms
#8592	비트 정확한 라우트 빌드	병합됨	–12ms
#8593	RoPE 및 sink-LSE 메타데이터 캐시	열림	–6ms
#8594	단일 커널 Sinkhorn	열림	–6.6ms
#8595	Ulysses 버퍼 복사 제거	열림	–10.6ms
#8627	어텐션 헤드 레이아웃 고정 (동적)	병합됨	1,429µs→55µs
#8855	W13 누산기 분할 (MUSA)	열림	690.8→590.4ms (정적)
#8915	1080p 리파이너 (블록 윈도우 어텐션)	초안	–

나열된 모든 PR은 병합되었거나 검토 중입니다. 이 게시물의 성능 수치는 각 PR이 측정된 시점의 코드 상태를 반영합니다.

7.2 호환성 경계

완전히 벤치마크됨: MUSA + SP8×CFG1+EP4 at 272p (섹션 5.1-5.3).
부분 지원: 540p 미리보기(샘플만), I2VA 모드(샘플만), 1080p 리파이너(초안 구현, 타이밍 없음).
검증되지 않음: 레이어별 오프로드, 양자화, Cache-DiT 및 기타 메모리 절약 기술.

8. 결론 및 향후 작업

속도 향상 요약: 호스트 동기화를 체계적으로 제거하고, MoE 커널을 융합하고, mHC 잔차를 재작성하고, 전문가 병렬 샤딩을 개선하여 15.6초 기준에서 단계당 565ms(27.6배)로 개선.
남은 병목 현상:
1. MoE W13 GEMM은 여전히 메모리 대역폭에 제한됨. 추가 타일링(128폭 K) 또는 사용자 지정 그룹 커널로 약 2–3ms 더 절감 가능.
2. 통신-계산 중첩: EP4는 단계당 60–90ms의 집단 통신이 발생하며 계산과 중첩되지 않음.
3. 유휴 호스트 시간: FlowUniPC 호스트 로직(약 19ms) 및 커널 내 실행 간격(약 14ms)이 남아 있음.
4. 동적 컴파일 오버헤드: 여러 커널에서 int64 크기 산술; int32로 변환하면 손실된 성능의 약 85% 회복.
다음 단계:
- MUSA에서 그래프 재생 활성화(현재 지역 컴파일보다 느림).
- 1080p 리파이너 통합 완료(블록 윈도우 어텐션, 메모리 계획).
- 병합된 PR만 사용하여 540p 벤치마킹 및 메인 브랜치를 포함하도록 CI 확장.
- 추가 지연 시간 감소를 위한 양자화 및 cache-DiT 기술 탐구.

9. 감사의 말

저자들은 MAGI-2 Preview를 오픈 소스로 공개하고 공식 프롬프트를 제공한 Sand AI, GPU 스택을 제공한 Moore Threads MUSA, torchada, MATE 팀, 그리고 원본 블로그 게시물에 나열된 PR을 구현하고 검토한 많은 기여자들에게 감사드립니다.

10. 재현성

컨테이너 이미지: registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview (다이제스트 sha256:2785b4626d1da1e0765a4cb515b9afc5349a22c6417e14c81091add0c1a9b902).

명령어 (8×MTT S5000에서 실행):

docker pull registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview
docker run -d --name magi2 -v /ipfs:/ipfs -v "$PWD/magi2-out":/work \
  --entrypoint sleep registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview infinity
docker exec magi2 bash -lc '
  cd /opt/vllm-omni-src &&
  python /opt/magi2/instrument-steps.py &&
  VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONUNBUFFERED=1 \
  python /opt/magi2/magi2_e2e_gate.py \
    --ulysses-degree 8 --cfg-parallel-size 1 \
    --enable-expert-parallel --expert-parallel-size 4 \
    --resolution 272p --height 256 --width 448 \
    --seconds 10 --fps 12.5 --num-frames 125 --num-inference-steps 100 \
    --no-pin-cpu-memory --metrics-dir /work/steps --save-outputs /work/outputs'

결정적 Inductor 동작을 위한 환경 변수는 소스 게시물의 부록 A.5에 설명된 대로 컨테이너 내부에 설정됩니다.

모든 숫자, 표, 인용은 2026년 10월 12일자 vLLM 블로그 게시물에서 직접 가져온 것입니다. 추가 데이터는 발명되지 않았습니다.