[GPT-SoVITS] GPT-SoVITS DiT 샘플러에 cache-dit(DBCache) 블록 캐싱 도입하기: forward를 세 조각으로 나눈 이유

PR 링크: RVC-Boss/GPT-SoVITS#2868 상태: Merged | 변경: +733 / -10

들어가며

GPT-SoVITS의 V3/V4/V5 계열은 mel을 생성할 때 F5-TTS 계열의 DiT(Diffusion Transformer) 를 flow-matching sampler로 여러 번 돌립니다. V5dev 기준으로 32 sampling step에 CFG(classifier-free guidance) positive/negative 두 branch가 돌아가므로, 요청 하나당 DiT forward가 64번 발생합니다. 추론 시간의 상당 부분이 이 스테이지에 몰려 있습니다.

그런데 diffusion sampling의 중후반부로 가면 스텝 간 hidden state의 변화가 매우 작아집니다. 이 관찰을 이용한 것이 vipshop/cache-dit의 DBCache(Dual Block Cache) 입니다. 앞쪽 몇 개 블록만 돌려보고, 그 잔차(residual)가 이전 스텝과 거의 같다면 나머지 중간 블록들의 결과를 이전 스텝의 잔차로 대체하는 방식입니다.

이 PR은 외부 의존성(pip install cache-dit) 없이 그 의미론을 직접 구현해 DiT sampler에 붙였고, 기본값은 off 입니다. 핵심 아이디어는 다음과 같습니다.

front : trunk(x) -> blocks[0:Fn] -> x_F        (every step)
        R  = x_F - trunk(x)                    Fn residual, the decision signal
        fd = sum|R - R_prev| / (sum|R_prev| + eps)
hit   : x_mid = x_F + MnResidual_prev, then norm_out -> proj_out only
miss  : x_F -> blocks[Fn:depth] -> x_M;  MnResidual = x_M - x_F
head  : x_mid -> norm_out -> proj_out -> vel   (every step)

코드 분석

1. dit.py — forward()를 trunk / run_blocks / head로 분할

가장 중요한 변경입니다. 기존 forward()는 embedding, 블록 루프, 출력 projection이 하나의 함수에 뭉쳐 있었습니다.

Before

        if self.long_skip_connection is not None:
            residual = x

        for block in self.transformer_blocks:
            if use_grad_ckpt:
                x = checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False)
            else:
                x = block(x, t, mask=mask, rope=rope)

        if self.long_skip_connection is not None:
            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))

        x = self.norm_out(x, t)
        output = self.proj_out(x)

After

    def run_blocks(self, x, t, mask, rope, start=0, end=None, use_grad_ckpt=False):
        """`front` is blocks[0:fn], `mid` is blocks[fn:depth]."""
        end = self.depth if end is None else end
        for block in self.transformer_blocks[start:end]:
            if use_grad_ckpt:
                x = checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False)
            else:
                x = block(x, t, mask=mask, rope=rope)
        return x

    def head(self, x, t, residual=None):
        """`norm_out -> proj_out`, computed on every step, hits included."""
        if self.long_skip_connection is not None:
            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
        x = self.norm_out(x, t)
        return self.proj_out(x)

self.transformer_blocks[start:end]는 nn.ModuleList 슬라이싱이라 새 ModuleList를 만들지만 내부 모듈은 동일 객체를 참조합니다. 즉 weight는 공유되고 복사되지 않습니다. 루프 본문은 한 글자도 바뀌지 않았기 때문에 off 경로의 수치 동일성을 보장하기가 쉽습니다.

왜 forward()를 black box로 감싸지 않았을까요? PR 설명에 따르면 잔차 R은 trunk(x)(블록 진입 직전 값) 기준으로 정의되어야 하고, head(norm_out은 timestep t에 조건화된 AdaLN)는 hit 스텝에서도 반드시 새 t로 다시 계산되어야 하기 때문입니다. 출력 전체를 캐싱하면 timestep 조건이 반영되지 않으므로, 블록 루프 자체를 잘라야 했습니다.

2. dit.py — 새 forward()의 분기

        if dbcache is None or not dbcache.enabled:
            # ---- original path: all blocks, then head ----
            x = self.run_blocks(x, t, mask, rope, 0, self.depth, use_grad_ckpt)
        else:
            fn = max(0, min(int(dbcache.fn), self.depth - 1))
            x_f = self.run_blocks(x, t, mask, rope, 0, fn, use_grad_ckpt)
            fd = dbcache.front_update(dbcache_branch, x_f, x)
            if dbcache.can_cache(dbcache_branch, fd):
                x = x_f + dbcache.mid_residual(dbcache_branch)
                dbcache.on_hit(dbcache_branch)
            else:
                x = self.run_blocks(x_f, t, mask, rope, fn, self.depth, use_grad_ckpt)
                dbcache.on_miss(dbcache_branch, x, x_f)

        output = self.head(x, t, residual)

몇 가지 포인트:

• fn을 [0, depth-1]로 clamp하여 mid 구간이 최소 1블록은 남도록 합니다.
• fd는 상대 L1 차이라 activation 스케일에 무관(scale free)합니다. 덕분에 threshold 하나로 여러 모델/입력에 대응할 수 있습니다.
• dbcache_branch(0=positive, 1=negative)로 CFG branch별 ledger를 분리합니다. 두 branch는 condition이 달라 잔차 분포도 다르므로, 섞으면 엉뚱한 잔차를 재사용하게 됩니다.
• hit일 때도 x_f + MnResidual_prev로 front의 최신 결과에 잔차를 더하므로, 단순히 이전 출력을 재사용하는 것보다 현재 스텝 정보를 더 많이 보존합니다.

3. models.py / cache_dit.py — 상태는 sampling run 단위

CFM(DBCacheMixin, nn.Module)으로 mixin을 붙이고, inference() 호출마다 DBCacheState를 새로 만들며 스텝마다 step_begin(j)을 부릅니다. 상태가 요청 단위로 생성되므로 off -> on -> off 순으로 호출해도 두 번째 off가 bit-identical합니다. 또한 hit 판정에 필요한 GPU→CPU 동기화(.item())를 스텝당 한 번으로 제한해 sync 비용이 이득을 잡아먹지 않도록 설계했습니다.

hit 조건에는 warmup 외에 세 개의 gate가 추가됩니다: max_cached(run당 캐시 스텝 수), max_cont(연속 hit 수), max_accum(누적 sum(fd) drift). 연속으로 캐시를 쓰면 오차가 누적되는 문제를 통제하는 장치입니다.

4. TTS.py — 안전한 스위치

        if cfg.enabled and str(getattr(self.configs, "version", "")) == "v5turbo" and not cfg.force:
            # 4-step turbo: nothing to cache, keep it off
            print(i18n("cache-dit: not supported for v5turbo (4 steps), staying off"))
            cfg.fn = 0
        cfm.set_dbcache(cfg)

auto는 16 step 이상에서만 F8을 켭니다. 4-step turbo는 warmup(기본 8)이 이미 전체 run을 덮어 캐시할 스텝이 없으므로, 켜봤자 품질 위험만 생긴다는 판단입니다. run()에서 sample_steps를 결정한 직후 self.apply_cache_dit(inputs, sample_steps)를 호출하고, API/WebUI/CLI 모두 같은 필드 이름(cache_dit, dbcache_*)과 GSV_DBCACHE* 환경변수 fallback을 공유합니다.

왜 이게 좋은가

측정 결과 (32-step v5dev, 동일 noise)

config	cached forwards	mel relL1	DiT stage
off	0/64	—	1.00x
f8, thr 0.02	12/64 (18.8%)	1.19e-3	~1.12x
f4, thr 0.02	14/64 (21.9%)	1.39e-3	~1.18x
f8, thr 0.05	29/64 (45.3%)	9.4e-3	~1.34x

보수적 기본값(0.02)에서도 DiT 스테이지가 약 12% 빨라지고 mel 오차는 1e-3 수준입니다. threshold를 0.05로 올리면 1.34x까지 가지만 오차는 약 8배 커집니다. max_cached=4로 묶으면 8/64 캐시, relL1 7.4e-4처럼 더 엄격한 품질 예산도 가능합니다.

일반적 교훈

1. 블랙박스 캐싱보다 구조적 분할. 모델 출력 전체를 캐시하지 않고 trunk / run_blocks / head로 잘라 timestep 조건 부분(head)은 매번 계산했습니다. 무엇을 캐시하고 무엇을 재계산할지 모델 구조에 맞춰 정해야 품질이 유지됩니다.
2. 상대 지표로 결정하기. fd를 정규화해 스케일과 무관한 threshold를 쓸 수 있게 했습니다.
3. off 경로의 bit-identical 검증. fp16/fp32 × long_skip_connection × static cache × inference/training 24개 조합과 threshold=0 케이스까지 비트 단위로 확인했습니다. 기본 off 최적화를 메인 경로에 넣을 때 리뷰어가 신뢰할 수 있는 가장 강력한 근거입니다.
4. 상태 수명은 요청 단위로. 전역 캐시 상태가 다음 요청에 새어나가지 않도록 sampling run마다 state를 생성했습니다.
5. GPU sync 최소화. hit 판정에 .item()이 필요하지만 스텝당 한 번으로 묶어 오버헤드를 억제했습니다.

한계도 명확히 적혀 있습니다. V3/V4는 시간 측정이 없고, skip_first_step 등 같은 스테이지를 건드리는 다른 스위치와의 조합은 아직 측정되지 않았습니다. hit rate가 참조 음성과 텍스트에 의존하므로 0.02는 "보편 상수"가 아니라 "보수적 기본값"으로 읽어야 합니다.

참고 자료

• vipshop/cache-dit — 이 PR이 참고한 DBCache 원본 구현
• torch.nn.ModuleList — transformer_blocks 슬라이싱에 사용된 컨테이너
• torch.utils.checkpoint — run_blocks에서 유지된 gradient checkpointing API
• torch.Tensor.item — 스텝당 한 번 호출되는 GPU→CPU 동기화 지점

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.