[sglang] SGLang, Sol-Attn 도입으로 비디오 생성 속도 1.23배 향상
PR 링크: sgl-project/sglang#33702 상태: Merged | 변경: +465 / -0
들어가며
최근 AI 기반 비디오 생성 기술은 놀라운 발전을 거듭하고 있지만, 여전히 생성 속도는 중요한 병목 지점으로 남아있습니다. 특히 MiniMax-H3와 같은 모델은 수많은 어텐션 레이어를 거치며 상당한 연산량을 요구합니다. 이러한 문제를 해결하기 위해, NVLabs에서 개발한 학습 없이 실시간으로 어텐션을 희소화하는 기술인 Sol-Attn이 주목받고 있습니다. 본 PR은 SGLang에 Sol-Attn을 새로운 어텐션 백엔드로 통합하여, 비디오 생성 과정의 속도를 크게 향상시키는 것을 목표로 합니다.
기존 SGLang은 MiniMax-H3 모델에서 FlashAttention (FA3)을 사용하여 forward_varlen 방식으로 어텐션을 처리했습니다. 이는 약 50개의 레이어와 각 스텝마다 약 38,000개의 토큰을 처리해야 하는 MiniMax-H3의 특성상 상당한 연산 부담을 야기했습니다. Sol-Attn은 이러한 어텐션 연산을 동적으로 라우팅하여 불필요한 계산을 줄임으로써 속도 향상을 가능하게 합니다.
Sol-Attn이란 무엇인가?
Sol-Attn은 비디오 생성 추론 속도를 높이기 위한 학습 없는(training-free), 온더플라이(on-the-fly) 희소 어텐션 방법론입니다. 전체 프록시 스코어 맵을 구체화하지 않고, 온라인 소프트맥스(online-softmax) 패스를 통해 임계값(threshold) 기반으로 어텐션 블록을 라우팅합니다. 이를 통해 비디오 생성 및 편집 작업에서 최대 2.1배 ~ 2.3배의 속도 향상을 보고했으며, 전체 Sol-Engine 스택 적용 시에는 더 큰 폭의 가속 효과를 기대할 수 있습니다.
본 PR은 Sol-Attn 커널 자체를 SGLang에 통합하는 데 초점을 맞추었으며, H3 모델에서 기본 설정으로 약 1.15배 ~ 1.23배의 디노이징 속도 향상과 함께 우수한 품질(PSNR 31 dB 이상)을 달성했습니다.
코드 변경 분석
이번 PR은 크게 두 부분으로 나누어 볼 수 있습니다. 첫째, SGLang의 문서에 Sol-Attn 백엔드를 추가하고 사용법 및 설정 옵션을 설명합니다. 둘째, Sol-Attn을 SGLang의 어텐션 백엔드로 등록하고, MiniMax-H3 모델에서 이를 활용할 수 있도록 관련 코드를 구현합니다.
1. 문서 업데이트 (docs/docs/sglang-diffusion/attention_backends.mdx)
새로운 어텐션 백엔드인 sol_attn이 추가되었음을 알리고, 설치 방법과 설정 파라미터에 대한 설명을 포함합니다. 특히, Sol-Attn 사용을 위해서는 별도의 sol-attn 패키지 설치가 필요함을 명시하고 있습니다.
Before:
--- a/docs/docs/sglang-diffusion/attention_backends.mdx
+++ b/docs/docs/sglang-diffusion/attention_backends.mdx
@@ -64,6 +64,11 @@
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`SAGE_ATTN_3`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Requires SageAttention3 installed per upstream instructions.</td>
</tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sol_attn`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`SOL_ATTN`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Requires the upstream <code>sol-attn</code> package. Install with <code>pip install git+https://github.com/NVlabs/Sana.git@sol-engine#subdirectory=techniques/sparse_backends</code>. BF16, head dim 128. Configure via <code>--attention-backend-config</code>.</td>
+ </tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`video_sparse_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`VIDEO_SPARSE_ATTN`</td>
After:
--- a/docs/docs/sglang-diffusion/attention_backends.mdx
+++ b/docs/docs/sglang-diffusion/attention_backends.mdx
@@ -339,6 +344,69 @@
</tbody>
</table>
+**Sol-Attn (`sol_attn`)**
+
+<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
+ <colgroup>
+ <col style={{width: "20%"}} />
+ <col style={{width: "16%"}} />
+ <col style={{width: "46%"}} />
+ <col style={{width: "18%"}} />
+ </colgroup>
+ <thead>
+ <tr style={{borderBottom: "2px solid #d55816"}}>
+ <th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Parameter</th>
+ <th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>Type</th>
+ <th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Description</th>
+ <th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>Default</th>
+ </tr>
+ </thead>
+ <tbody>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`tau`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`float`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Routing threshold scale. Higher values select fewer exact KV blocks.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`1.0`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`thresh_type`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`str`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Threshold mode: `diag` or `exact`.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`diag`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sink_tokens`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Exact KV sink length for prefix tokens such as text/audio rows.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sink_start`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Start index of the exact KV sink range.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`dense_steps`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Use dense attention for the first N denoising steps.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`10`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`dense_layers`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`str`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Layer indices kept dense, e.g. `0,1` or `0-2`.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0,1`</td>
+ </tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`kv_splits`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int | str`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>KV split factor passed to the Sol-Attn kernel. Use `auto` on long sequences.</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`auto`</td>
+ </tr>
+ </tbody>
+</table>
+
## Platform support matrix
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
@@ -413,6 +481,16 @@ Some backends require additional configuration. You can pass these parameters vi
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA-only (optional dependency).</td>
</tr>
+ <tr>
+ <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`sol_attn`</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>No</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
+ <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA-only. Requires <code>sol-attn</code>. Install with <code>pip install git+https://github.com/NVlabs/Sana.git@sol-engine#subdirectory=techniques/sparse_backends</code>. Configure via <code>--attention-backend-config</code>.</td>
+ </tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`video_sparse_attn`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
2. Sol-Attn 백엔드 구현 (python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py)
이 파일은 Sol-Attn을 SGLang의 어텐션 백엔드로 통합하는 핵심 로직을 담고 있습니다. SolAttnBackend 클래스는 백엔드의 메타데이터를 정의하고, SolAttnImpl 클래스는 실제 어텐션 연산을 수행합니다. 특히, _get_sol_attn_runtime_config 함수는 --attention-backend-config를 통해 전달된 파라미터들을 파싱하여 Sol-Attn 커널에 전달하는 역할을 합니다. _parse_layer_ranges와 _resolve_kv_splits 같은 헬퍼 함수들은 복잡한 설정 값들을 처리합니다.
Before: (해당 파일은 새로 생성되었으므로 Before diff는 없습니다.)
After:
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py
@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import re
+
+import torch
+
+from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+ AttentionBackend,
+ AttentionImpl,
+ AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_SOL_ATTN_HEAD_DIM = 128
+
+
def _parse_layer_ranges(spec: str | int | None) -> frozenset[int]:
+ if spec is None:
+ return frozenset()
+ if isinstance(spec, int):
+ return frozenset({spec})
+ layers: set[int] = set()
+ for item in str(spec).split(","):
+ item = item.strip()
+ if not item:
+ continue
+ if "-" in item:
+ start, end = item.split("-", 1)
+ layers.update(range(int(start), int(end) + 1))
+ else:
+ layers.add(int(item))
+ return frozenset(layers)
+
+
def _resolve_kv_splits(q: torch.Tensor, kv_splits: int | str | None) -> int:
+ if kv_splits not in (None, "auto"):
+ return int(kv_splits)
+ arch = tuple(torch.cuda.get_device_capability(q.device))
+ if arch == (9, 0) and q.shape[1] >= 65536:
+ try:
+ import cuda.bindings.driver # noqa: F401
+ import cutlass.cute # noqa: F401
+
+ return 4
+ except ImportError:
+ pass
+ return 1
+
+
def _get_sol_attn_runtime_config() -> dict:
+ server_args = get_global_server_args()
+ cfg = getattr(server_args, "attention_backend_config", None) or {}
+ dense_layers = cfg.get("dense_layers", "0,1")
+ sink_start = cfg.get("sink_start", 0)
+ return {
+ "tau": float(cfg.get("tau", 1.0)),
+ "thresh_type": str(cfg.get("thresh_type", "diag")),
+ "kv_splits": cfg.get("kv_splits", "auto") ,
+ "sink_tokens": int(cfg.get("sink_tokens", 0)),
+ "sink_start": None if sink_start is None else int(sink_start),
+ "dense_steps": int(cfg.get("dense_steps", 10)),
+ "dense_layers": _parse_layer_ranges(dense_layers),
+ }
+
+
+class SolAttnBackend(AttentionBackend):
+ accept_output_buffer: bool = True
+
+ @staticmethod
+ def get_supported_head_sizes() -> list[int]:
+ return [_SOL_ATTN_HEAD_DIM]
+
+ @staticmethod
+ def get_enum() -> AttentionBackendEnum:
+ return AttentionBackendEnum.SOL_ATTN
+
+ @staticmethod
+ def get_impl_cls() -> type[SolAttnImpl]:
+ return SolAttnImpl
+
+
+class SolAttnImpl(AttentionImpl):
+
+ def __init__(
+ self,
+ num_heads: int,
+ head_size: int,
+ causal: bool,
+ softmax_scale: float,
+ num_kv_heads: int | None = None,
+ prefix: str = "",
+ **extra_impl_args,
+ ) -> None:
+ del num_heads, num_kv_heads, extra_impl_args
+ if head_size != _SOL_ATTN_HEAD_DIM:
+ raise ValueError(
+ f"Sol-Attn requires head_size={_SOL_ATTN_HEAD_DIM}, got {head_size}"
+ )
+ self.causal = causal
+ self.softmax_scale = softmax_scale
+ self.prefix = prefix
+ self.layer_idx = self._parse_layer_idx(prefix)
+
+ @staticmethod
+ def _parse_layer_idx(prefix: str) -> int | None:
+ match = re.search(r"blocks\.(\d+)", prefix)
+ if match is None:
+ return None
+ return int(match.group(1))
+
+ def _should_use_dense(self) -> bool:
+ cfg = _get_sol_attn_runtime_config()
+ try:
+ from sglang.multimodal_gen.runtime.managers.forward_context import (
+ get_forward_context,
+ )
+
+ step = int(get_forward_context().current_timestep)
+ except AssertionError:
+ step = 0
+
3. MiniMax-H3 모델 통합 (python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py)
리뷰 댓글에서 언급된 것처럼, 이 PR은 MiniMax-H3 모델에 특화된 변경사항을 포함합니다. 기존 H3 모델의 forward_varlen 함수는 prefix 정보를 제대로 전달하지 않았는데, Sol-Attn은 dense_layers 설정을 위해 이 prefix 정보가 필요합니다. 따라서 해당 PR은 H3 모델의 forward_varlen 함수를 수정하여 prefix 정보를 Sol-Attn 백엔드로 전달하도록 합니다.
Before:
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -12,7 +12,7 @@
def forward_varlen(self, input_ids, attention_mask, attention_meta: AttentionMetadata, output_ids=None):
# TODO: This is a hacky way to get the prefix. We should refactor this.
# For now, we assume that the prefix is the same for all attention layers.
- prefix = ""
+ prefix = self.get_prefix(attention_meta)
# TODO: This is a hacky way to get the prefix. We should refactor this.
# For now, we assume that the prefix is the same for all attention layers.
After:
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -12,7 +12,7 @@
def forward_varlen(self, input_ids, attention_mask, attention_meta: AttentionMetadata, output_ids=None):
# TODO: This is a hacky way to get the prefix. We should refactor this.
# For now, we assume that the prefix is the same for all attention layers.
- prefix = ""
+ prefix = self.get_prefix(attention_meta)
# TODO: This is a hacky way to get the prefix. We should refactor this.
# For now, we assume that the prefix is the same for all attention layers.
왜 이게 좋은가?
성능 향상
PR 설명에 따르면, MiniMax-H3 모델에서 Sol-Attn 백엔드를 사용할 경우 다음과 같은 성능 향상을 얻을 수 있습니다.
-
H3 t2va (Text-to-Video Animation):
- 기존 FlashAttention (FA3): 27.0초
- Sol-Attn: 23.5초 (약 1.15배 속도 향상)
- 품질: PSNR 32.9 dB (FA3 대비 우수)
-
H3 ref2va (Reference-to-Video Animation):
- 기존 FlashAttention (FA3): 119.8초
- Sol-Attn: 97.2초 (약 1.23배 속도 향상)
- 품질: PSNR 31.1 dB (FA3 대비 우수)
이 수치들은 Sol-Attn이 단순히 속도만 높이는 것이 아니라, 품질 저하 없이 효율적인 연산을 가능하게 함을 보여줍니다. 특히, dense_steps와 tau 같은 파라미터를 조정하여 더 높은 속도 향상(최대 1.62배)을 얻을 수도 있지만, 품질 저하가 발생할 수 있으므로 기본 설정값(tau=1.0, dense_steps=10, dense_layers='0,1')이 권장됩니다.
일반적 교훈
- 특화된 커널의 중요성: 비디오 생성과 같이 특정 도메인에서 반복적으로 발생하는 연산 패턴(예: 긴 시퀀스의 어텐션)은 특화된 커널(Sol-Attn)을 통해 큰 성능 향상을 얻을 수 있습니다. 범용적인 솔루션보다 특정 문제에 최적화된 도구를 사용하는 것이 효과적입니다.
- 동적 희소성 활용: 모든 연산이 동일한 중요도를 갖는 것은 아닙니다. Sol-Attn처럼 실시간으로 연산의 중요도를 판단하고 불필요한 부분을 건너뛰는 동적 희소성(dynamic sparsity) 기법은 효율성을 극대화할 수 있습니다.
- 프레임워크 통합의 중요성: SGLang과 같이 모델 서빙 프레임워크는 다양한 최적화 기법(어텐션 백엔드, 양자화 등)을 쉽게 통합하고 적용할 수 있는 인터페이스를 제공해야 합니다. 이를 통해 사용자는 복잡한 최적화 과정을 직접 구현하지 않고도 이점을 누릴 수 있습니다.
- 품질-속도 트레이드오프 관리: 성능 향상을 추구할 때 품질 저하를 간과해서는 안 됩니다. Sol-Attn의 다양한 설정 옵션처럼, 사용자가 성능과 품질 사이의 균형점을 찾을 수 있도록 유연한 설정과 명확한 가이드라인을 제공하는 것이 중요합니다.
리뷰 피드백 반영
리뷰어 niehen6174의 지적처럼, 이 PR은 MiniMax-H3 모델에 특화된 변경사항을 포함합니다. H3 모델은 다른 모델들과 달리 forward_varlen 함수에서 어텐션 레이어의 prefix 정보를 전달하는 방식이 달랐습니다. Sol-Attn은 dense_layers 설정을 위해 이 prefix 정보가 필요한데, 이를 위해 H3 모델의 forward_varlen 함수가 수정되었습니다. 다른 모델들은 이미 prefix 정보를 잘 전달하고 있어 별도의 수정이 필요 없었으며, 이는 향후 다른 모델에 Sol-Attn을 적용할 때 고려해야 할 사항입니다. 또한, mickqian의 요청에 따라 CI가 재실행되었습니다.
결론
이번 PR은 SGLang에 Sol-Attn 희소 어텐션 백엔드를 성공적으로 통합하여, 특히 MiniMax-H3 모델에서의 비디오 생성 속도를 크게 향상시켰습니다. 이는 AI 모델의 효율성을 높이고 더 빠른 결과물을 얻는 데 기여할 것입니다. 앞으로 Sol-Attn의 적용 범위를 넓히고, 관련 기능(예: Morton 3D 토큰 재정렬 지원)을 추가하여 더욱 강력한 성능을 제공할 것으로 기대됩니다.
참고 자료
- https://nvlabs.github.io/Sana/Sol-Attn/
- https://arxiv.org/abs/2607.24027
- https://github.com/NVlabs/Sana/tree/sol-engine
- https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
- https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/server_args.py
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [ultralytics] RT-DETR FLOPs 프로파일링 성능 최적화 및 안정화
- 현재글 : [sglang] SGLang, Sol-Attn 도입으로 비디오 생성 속도 1.23배 향상
- 다음글 [hermes-agent] Hermes Agent: 10배 빠른 프로젝트 그룹화 최적화 분석
댓글