[sglang] LTX2.3 HQ Denoising 성능 최적화: Attention Skip을 활용한 효율적인 모델 호출

PR 링크: sgl-project/sglang#24298 상태: Merged | 변경: +None / -None

들어가며

최근 sglang 레포지토리의 Pull Request([codex] Optimize LTX2.3 HQ denoising split passes)는 LTX2.3 HQ 가이드 Denoising 과정의 성능을 크게 향상시키는 중요한 최적화를 포함하고 있습니다. 기존 코드에서는 Denoising 과정 중 각 스텝에서 모델을 호출할 때, 불필요한 Attention 계산이 수행되어 성능 저하의 원인이 되었습니다. 이 PR은 이러한 비효율성을 제거하고, 특히 HQ(High Quality) 모드에서 각 모델 호출이 단일 확장 배치(expanded-batch) 항목만을 처리할 때 Attention 관련 옵션을 직접 모델 인자로 전달함으로써 최적화를 달성했습니다. 본 글에서는 이 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 변경이 성능 향상에 기여하는지, 그리고 일반적인 딥러닝 모델 최적화에 어떤 교훈을 주는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py 파일에 집중되어 있습니다. 주요 변경 내용을 파일별로 나누어 살펴보겠습니다.

1. _ltx2_res2s_sde_step 함수의 개선

이 함수는 SDE(Stochastic Differential Equation) 기반의 Denoising 스텝을 처리합니다. 기존 코드에서는 sigma_up 또는 sigma_next가 0인 경우를 체크하여 조기 반환하는 로직이 있었지만, 이 로직이 CUDA bool 동기화(sync)를 유발할 수 있었습니다. PR에서는 이 부분을 개선하여 terminal이라는 새로운 인자를 추가했습니다.

Before:

-        if bool((sigma_up == 0).any()) or bool((sigma_next == 0).any()):
-            return denoised_sample.to(dtype=sample.dtype)

After:

+        if terminal:
+            return denoised_sample.to(dtype=sample.dtype)

terminal 인자가 True일 경우, 즉 SDE 스텝의 마지막 단계인 경우에만 조기 반환하도록 변경되었습니다. 이는 res2s의 SDE 터미널 경로에서 내부 CUDA bool 동기화를 피하기 위한 조치입니다. 터미널 단계가 아닌 일반적인 스텝에서는 이 동기화가 불필요하며, 오히려 성능에 부정적인 영향을 줄 수 있기 때문입니다.

2. _build_ltx2_model_kwargs 및 관련 함수 추가

PR은 HQ 모드에서의 최적화를 위해 새로운 헬퍼 함수들을 도입했습니다. 핵심 아이디어는 HQ 모드에서 각 모델 호출이 단일 expanded_batch_item만을 처리할 때, perturbation_configs 리스트 대신 직접 disable-attention 옵션을 모델 인자로 전달하는 것입니다.

• _ltx2_guidance_perturbation_config: 개별 pass_spec으로부터 모델 인자로 전달될 perturbation_config 딕셔너리를 생성합니다.
• _build_ltx2_guidance_perturbation_configs: 여러 pass_spec에 대해 perturbation_configs 튜플을 생성합니다. 이는 HQ 모드가 아닌 경우에 사용됩니다.
• _apply_ltx2_guidance_pass_kwargs: pass_spec의 disable-attention 옵션을 model_kwargs 딕셔너리에 직접 복사합니다. 이 함수는 HQ 모드에서 사용됩니다.

기존 _build_ltx2_model_kwargs의 일부 (개념적):

# ...
if perturbation_configs:
    kwargs["perturbation_configs"] = perturbation_configs
# ...

새로운 로직 (HQ 모드): _stage2_midpoint_model_call 함수 내에서 use_split_pass_kwargs 플래그를 통해 HQ 모드인지 아닌지를 판단합니다. HQ 모드(server_args.pipeline_class_name == "LTX2TwoStageHQPipeline")일 경우, _apply_ltx2_guidance_pass_kwargs 함수를 사용하여 skip_video_self_attn_blocks, skip_audio_self_attn_blocks, disable_a2v_cross_attn, disable_v2a_cross_attn 등의 옵션을 model_kwargs_chunk에 직접 설정합니다.

Before (개념적, evaluate_stage1_guided_x0 내부):

-                    perturbation_configs = tuple(
-                        {
-                            "skip_video_self_attn_blocks": pass_spec.skip_video_self_attn_blocks,
-                            "skip_audio_self_attn_blocks": pass_spec.skip_audio_self_attn_blocks,
-                            "skip_a2v_cross_attn": pass_spec.disable_a2v_cross_attn,
-                            "skip_v2a_cross_attn": pass_spec.disable_v2a_cross_attn,
-                        }
-                        for pass_spec in pass_specs
-                        for _ in range(batch_size_local)
-                    )
+# ...
+                            for model_kwargs_chunk, perturbation_config in zip(
+                                self._split_ltx2_model_kwargs(
+                                    batched_model_kwargs, split_sizes
+                                ),
+                                perturbation_configs,
+                                strict=True,
+                            ):
+                                model_kwargs_chunk["perturbation_configs"] = (
+                                    perturbation_config,
+                                )

After (개념적, evaluate_stage1_guided_x0 내부):

+        use_split_pass_kwargs = (
+            server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
+        )
+# ...
+                            for index, (model_kwargs_chunk, pass_spec) in enumerate(
+                                zip(
+                                    self._split_ltx2_model_kwargs(
+                                        batched_model_kwargs, split_sizes
+                                    ),
+                                    split_pass_specs,
+                                    strict=True,
+                                )
+                            ):
+                                if use_split_pass_kwargs:
+                                    self._apply_ltx2_guidance_pass_kwargs(
+                                        model_kwargs_chunk, pass_spec
+                                    )
+                                else:
+                                    model_kwargs_chunk["perturbation_configs"] = (
+                                        split_perturbation_configs[index],
+                                    )

이 변경은 HQ 모드에서 각 모델 호출이 단일 pass_spec만을 처리하므로, perturbation_configs라는 리스트 전체를 전달하는 대신 필요한 disable-attention 옵션만 직접 전달하여 오버헤드를 줄입니다. 이는 특히 batch_size_local * len(pass_specs)가 클 때 효율적입니다.

3. evaluate_stage1_guided_x0 함수의 분기 처리 개선

evaluate_stage1_guided_x0 함수는 Stage 1 Denoising을 평가하는 핵심 로직을 담고 있습니다. 이 함수 내에서 use_split_stage1_guided_passes 조건에 따라 모델 호출 방식이 달라지는데, 이번 PR에서는 HQ 모드에 대한 처리가 더욱 명확해졌습니다.

• HQ 모드(use_split_pass_kwargs가 True)에서는 _apply_ltx2_guidance_pass_kwargs를 통해 개별 model_kwargs_chunk에 직접 Attention Skip 옵션을 적용합니다.
• HQ 모드가 아닌 경우(use_split_pass_kwargs가 False)에는 기존과 같이 perturbation_configs를 사용하여 여러 pass_spec의 설정을 한 번에 전달합니다.

이러한 분기 처리는 각 모드의 특성에 맞게 최적화된 방식으로 모델을 호출하도록 보장하며, 특히 HQ 모드에서의 불필요한 perturbation_configs 생성을 방지합니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 벤치마크 결과는 이 최적화의 효과를 명확히 보여줍니다. H100 80GB GPU에서 ltx_2_3_hq_pipeline을 사용한 테스트 결과, 다음과 같은 성능 개선이 관찰되었습니다:

Metric	Parent	Optimized	Delta
LTX2AVDenoisingStage	10054.60 ms	9737.69 ms	+3.15%
Avg denoise step	601.32 ms	582.47 ms	+3.13%
Median denoise step	636.53 ms	617.63 ms	+2.97%
E2E baseline dump	15468.94 ms	15192.81 ms	+1.78%

전반적으로 Denoising 관련 메트릭에서 약 2~3%의 성능 향상이 있었으며, End-to-End(E2E) 처리 시간도 개선되었습니다. 이러한 성능 향상은 다음과 같은 이유로 달성되었습니다:

1. 불필요한 CUDA 동기화 제거: _ltx2_res2s_sde_step 함수에서 terminal 인자를 도입하여 불필요한 CUDA bool 동기화를 제거함으로써, 특히 SDE 스텝의 마지막 단계에서 발생하는 오버헤드를 줄였습니다.
2. Attention 계산 최적화: HQ 모드에서 각 모델 호출이 단일 expanded_batch_item만을 처리할 때, perturbation_configs 리스트 전체를 전달하는 대신 필요한 disable-attention 옵션만을 직접 모델 인자로 전달함으로써 모델 내부의 불필요한 Attention 계산을 건너뛰도록 했습니다. 이는 GPU 연산량을 줄이고 메모리 접근을 최적화하는 효과를 가져옵니다.

일반적인 교훈

이 PR은 딥러닝 모델 최적화에 있어 몇 가지 중요한 교훈을 제공합니다:

1. 컨텍스트에 맞는 최적화: 모든 상황에 동일한 최적화 기법을 적용하는 것은 비효율적일 수 있습니다. HQ 모드와 같이 특정 조건(단일 expanded_batch_item 처리)에서는 perturbation_configs 리스트 대신 직접 인자를 전달하는 것이 더 효율적입니다. 모델의 내부 동작 방식과 호출 컨텍스트를 이해하는 것이 중요합니다.
2. GPU 동기화 오버헤드 최소화: GPU 연산에서 동기화는 성능 병목의 주요 원인 중 하나입니다. _ltx2_res2s_sde_step 함수의 변경처럼, 불필요한 동기화를 식별하고 제거하는 것은 성능 향상에 직접적인 영향을 미칩니다. 특히 조건부 로직에서 동기화 발생 여부를 신중하게 고려해야 합니다.
3. 명확한 API 설계: 새로운 헬퍼 함수(_apply_ltx2_guidance_pass_kwargs, _ltx2_guidance_perturbation_config 등)의 도입은 코드의 가독성과 유지보수성을 높입니다. 각 함수의 역할이 명확해짐으로써, 향후 유사한 최적화나 디버깅이 용이해집니다.
4. 벤치마킹의 중요성: PR 설명에 포함된 상세한 성능 벤치마크 결과는 최적화의 효과를 객관적으로 입증합니다. 실제 하드웨어 및 워크로드에서의 성능 측정은 최적화의 성공 여부를 판단하는 데 필수적입니다.

결론

이번 PR은 LTX2.3 HQ Denoising 파이프라인의 성능을 실질적으로 개선하는 중요한 최적화를 성공적으로 적용했습니다. 불필요한 CUDA 동기화를 제거하고, HQ 모드에서의 모델 호출 방식을 최적화하여 Denoising 과정의 효율성을 높였습니다. 이러한 변경은 더 빠르고 효율적인 멀티모달 생성 모델의 실행을 가능하게 하며, 딥러닝 모델 최적화에 대한 귀중한 통찰력을 제공합니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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