[sglang] NVIDIA SM90 GPU를 위한 SGLang SubBlock Sparse Attention 최적화: Sage FP8 Compute 도입

2026년 9월 9일수정: 2026년 9월 9일

PR 링크: sgl-project/sglang#37982 상태: Merged | 변경: +900 / -68

들어가며

최근 대규모 언어 모델(LLM) 및 멀티모달 모델의 발전은 GPU 하드웨어의 한계를 끊임없이 시험하고 있습니다. 특히, 트랜스포머 아키텍처의 핵심인 어텐션 메커니즘은 계산량과 메모리 사용량 측면에서 병목 현상을 일으키는 주요 원인 중 하나입니다. 이러한 문제를 해결하기 위해 SGLang 프로젝트에서는 다양한 최적화 기법을 도입하고 있으며, 이번 PR은 NVIDIA의 최신 GPU 아키텍처인 SM90 (Hopper)에서 SubBlock Sparse Attention의 성능을 극대화하기 위한 중요한 개선 사항을 담고 있습니다.

이 PR은 기존의 BF16 기반 SubBlock Sparse Attention에 더해, SM90 아키텍처에 특화된 Sage FP8 Compute 모드를 도입합니다. 이는 FlashInfer의 성공적인 사례에서 영감을 받아, INT8 양자화된 Query/Key와 FP8으로 계산된 Value/Softmax 확률을 활용하여 연산 효율성을 높이는 것을 목표로 합니다. 모델 가중치는 그대로 유지하면서도, 특정 하드웨어에서 더 빠른 연산을 가능하게 하는 것이 핵심입니다. 이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 최적화가 왜 효과적인지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR은 주로 두 개의 파일을 수정 및 추가했습니다:

python/sglang/kernels/ops/attention/subblock_sage_fp8_sm90.py: SM90 아키텍처를 위한 Sage FP8 Compute를 구현하는 새로운 Triton 커널 및 관련 로직을 정의합니다.
python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py: 서버 시작 시 GPU 기능 및 SpargeAttention 의존성을 확인하고, 새로운 sage_fp8 compute 모드를 지원하도록 설정을 업데이트합니다.

1. `subblock_sage_fp8_sm90.py`: SM90 Sage FP8 커널 구현

이 파일은 SM90 GPU에서 SubBlock Sparse Attention을 위한 Sage FP8 Compute를 담당하는 핵심 로직을 포함합니다. 주요 변경 사항은 다음과 같습니다:

SageAttention2 커널 통합: SM90 아키텍처의 특성을 활용하는 spas_sage_attn 라이브러리의 SageAttention2 커널을 사용합니다. 이 커널은 64-토큰 Query 블록과 128-토큰 Key 블록을 사용하며, Query와 Key를 INT8로, Value와 Softmax 확률을 FP8 (E4M3)으로 온라인 양자화합니다. 모델 가중치는 BF16으로 유지됩니다.
_load_sparge_attention_sm90_ops 함수: SpargeAttention 라이브러리에서 필요한 함수들 (e.g., get_vanilla_qk_quant, block_map_lut_triton, transpose_pad_permute_cuda, scale_fuse_quant_cuda, qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90)을 동적으로 로드합니다. 만약 SpargeAttention이 설치되지 않았다면, 사용자에게 설치를 안내하는 에러 메시지를 발생시킵니다.
_routing_plan_to_block_map 함수: SubBlock 라우팅 계획을 SpargeAttention 커널이 요구하는 Dense Boolean Map 형태로 변환하는 로직입니다. Triton 커널을 사용하여 효율적으로 이 변환을 수행합니다. Query 블록 크기(64)와 Key 블록 크기(128)에 맞춰 최적화되었습니다.
subblock_sage_fp8_sm90_attention 함수: 실제 어텐션 연산을 수행하는 메인 함수입니다. 입력 텐서(Q, K, V)의 유효성을 검증하고, 필요한 경우 데이터를 전처리한 후, SpargeAttention의 네이티브 SM90 커널을 호출합니다. 최종 결과는 입력과 동일한 BF16 형식으로 반환됩니다.

코드 인용 (예시 - subblock_sage_fp8_sm90.py의 일부):

# Before (기존 BF16 SubBlock Attention 로직 - 개념적 예시, 실제 diff에는 없음)
# ... 기존 BF16 연산 ...

# After (새로운 Sage FP8 SM90 어텐션 로직)
# ... (위의 _load_sparge_attention_sm90_ops, _routing_plan_to_block_map 호출) ...

        q_hnd = q.transpose(1, 2).contiguous()
        k_hnd = k.transpose(1, 2).contiguous()
        v_hnd = v.transpose(1, 2).contiguous()
        k_mean = k_hnd.mean(dim=-2, keepdim=True)
        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(
            q_hnd,
            k_hnd,
            k_mean,
            SAGE_FP8_SM90_QUERY_BLOCK_SIZE, # 64
            SAGE_FP8_SM90_KEY_BLOCK_SIZE,   # 128
        )
        # ... (Value FP8 양자화 및 커널 호출) ...
        kernel(
            q_int8,
            k_int8,
            v_fp8,
            output_hnd,
            lut,
            valid_block_counts,
            q_scale,
            k_scale,
            v_scale,
            1,  # HND tensor layout
            False,
            1,  # per-block Q/K scales
            softmax_scale,
        )

2. `minimax_h3.py`: 설정 및 유효성 검사 업데이트

이 파일은 SGLang 서버의 설정 및 파이프라인 구성을 담당합니다. 새로운 Sage FP8 Compute 모드를 지원하기 위해 다음과 같은 변경이 이루어졌습니다:

validate_server_args 함수 수정: AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN이 선택되었을 때, attention_backend_config에서 compute_mode를 확인합니다. 지원되는 모드는 bf16 (기본값)과 sage_fp8입니다.
SM90 호환성 검사: compute_mode가 sage_fp8로 설정된 경우, 현재 GPU의 컴퓨트 캡 능력이 9.0 (SM90)인지 확인합니다. 그렇지 않으면 오류를 발생시켜 잘못된 하드웨어에서의 사용을 방지합니다.
SpargeAttention 의존성 확인: sage_fp8 모드를 사용하기 위해 필요한 spas_sage_attn 라이브러리의 로딩을 시도하고, 실패 시 설치 방법을 안내합니다.

코드 인용 (예시 - minimax_h3.py의 일부):

# Before (기존 설정 로직 - 개념적 예시)
# ...
# if selected_backend is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN:
#     # 기존 BF16 관련 로직
#     ...

# After (새로운 Sage FP8 모드 지원 로직 추가)
        if selected_backend is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN:
            attention_config = server_args.attention_backend_config or {}
            compute_mode = str(attention_config.get("compute_mode", "bf16"))
            if compute_mode not in ("bf16", "sage_fp8"):
                raise ValueError(
                    "SubBlock compute_mode must be 'bf16' or 'sage_fp8', got "
                    f"{compute_mode!r}."
                )
            if compute_mode == "sage_fp8":
                capability = current_platform.get_device_capability()
                if capability is None or capability.to_int() != 90:
                    # ... SM90 검증 로직 ...
                from sglang.kernels.ops.attention.subblock_sage_fp8_sm90 import (
                    _load_sparge_attention_sm90_ops,
                )

                _load_sparge_attention_sm90_ops() # SpargeAttention 로딩 시도

왜 이게 좋은가?

이 PR은 여러 측면에서 중요한 성능 개선 및 기술적 발전을 이루었습니다.

1. 성능 향상 (Speed Tests)

PR 설명에 포함된 벤치마크 결과는 Sage FP8 Compute 모드의 효과를 명확하게 보여줍니다. 8개의 H200 GPU (SM90 기반)에서 테스트된 결과에 따르면:

Dense FA 대비 속도 향상: SubBlock+Sage 모드는 Dense FA 대비 최대 51.09%까지 생성 시간 단축을 보였습니다. 이는 특히 REF2VA-V 케이스에서 두드러집니다.
SB (BF16 SubBlock) 대비 속도 향상: Sage FP8 모드는 기존의 BF16 SubBlock 모드보다도 평균적으로 6.32% ~ 20.90% 더 빠른 성능을 제공합니다. 이는 FP8 연산의 효율성을 입증합니다.

코드 인용 (벤치마크 결과 테이블):

| Case       | Duration | Dense FA (s) | SB (s) | SB ↓   | SB+Sage (s) | SB+Sage ↓ |
|------------|----------|--------------|--------|--------|-------------|-----------|
| T2VA       | 5s       | 39.21        | 30.38  | 22.52% | 28.46       | 27.42%    |
| T2VA       | 10s      | 111.84       | 78.16  | 30.11% | 67.90       | 39.29%    |
| T2VA       | 15s      | 218.72       | 144.44 | 33.96% | 121.35      | 44.52%    |
| REF2VA-V   | 15s      | 842.06       | 520.66 | 38.17% | 411.83      | 51.09%    |

2. 메모리 사용량 유지

성능 향상에도 불구하고, Peak Reserved Memory는 Dense FA 및 SB 모드와 거의 동일하게 유지되었습니다. 이는 FP8/INT8 양자화가 주로 연산 단계에서 이루어지며, 모델 가중치나 중간 활성화 값의 메모리 footprint에 큰 영향을 미치지 않음을 시사합니다. 이는 메모리 제약이 있는 환경에서도 이점을 제공합니다.

3. 정확도 유지

성능 향상이 정확도 저하로 이어지는 경우는 흔하지만, 이 PR은 이를 최소화했습니다. Video SSIM 측정 결과, SB+Sage 모드는 SB (BF16 SubBlock) 대비 SSIM 값이 평균 0.00118 ~ 0.00918 포인트 하락하는 데 그쳤습니다. 이는 FP8/INT8 양자화가 모델의 최종 출력 품질에 미치는 영향을 최소화하면서 성능을 개선했음을 의미합니다.

코드 인용 (SSIM 결과 테이블):

| Case       | Duration | SB      | SB+Sage |
|------------|----------|---------|---------|
| T2VA       | 5s       | 0.821579| 0.818747|
| T2VA       | 10s      | 0.870746| 0.865564|
| REF2VA-V   | 15s      | 0.891421| 0.884516|

4. 하드웨어 특화 최적화

이 최적화는 SM90 (Hopper) 아키텍처의 Tensor Core 기능을 최대한 활용합니다. FP8 연산은 Hopper 아키텍처에서 BF16보다 훨씬 높은 처리량을 제공하므로, 이를 활용하는 것은 최신 하드웨어의 성능을 극대화하는 데 필수적입니다. 또한, SubBlock 라우팅을 SM90 네이티브 커널의 Q64xK128 블록 레이아웃에 맞게 조정함으로써, 커널의 효율성을 높였습니다.

5. 모듈성 및 확장성

새로운 sage_fp8 모드는 옵트인(opt-in) 방식으로 제공되며, 기존 BF16 모드는 기본값으로 유지됩니다. 이는 사용자가 필요에 따라 성능과 정확도 사이의 균형을 선택할 수 있게 합니다. 또한, SpargeAttention과 같은 외부 라이브러리를 활용하여 최신 GPU 기능을 통합하는 접근 방식은 향후 새로운 하드웨어 아키텍처나 최적화 기법이 등장했을 때 SGLang의 확장성을 높여줍니다.

일반적 교훈 및 시사점

이 PR은 다음과 같은 중요한 기술적 교훈을 제공합니다:

최신 하드웨어 기능 활용의 중요성: NVIDIA의 최신 GPU 아키텍처(SM90)는 FP8과 같은 새로운 데이터 타입을 지원하며, 이를 적극적으로 활용하는 것이 성능 향상의 핵심입니다. 하드웨어의 발전 속도에 맞춰 소프트웨어도 최적화되어야 합니다.
양자화 기술의 발전: INT8 및 FP8 양자화는 연산 효율성을 크게 높이면서도 정확도 손실을 최소화하는 방향으로 발전하고 있습니다. 특히 온라인 양자화는 모델 가중치를 변경하지 않고도 성능을 개선할 수 있는 유연한 방법입니다.
커널 최적화의 중요성: Triton과 같은 DSL을 사용하여 GPU 아키텍처에 특화된 고성능 커널을 작성하는 것은 LLM 추론 성능에 결정적인 영향을 미칩니다. SubBlock 라우팅과 네이티브 커널의 블록 레이아웃을 일치시키는 것은 이러한 최적화의 좋은 예입니다.
외부 라이브러리 활용 전략: 고도로 최적화된 외부 라이브러리(예: SpargeAttention)를 통합하는 것은 개발 시간을 단축하고 최신 기술을 빠르게 도입하는 효과적인 방법입니다. 다만, 의존성 관리 및 호환성 검증이 중요합니다.
벤치마킹과 정확도 검증의 필수성: 성능 개선은 반드시 정확도 손실과 함께 평가되어야 합니다. 다양한 시나리오에서의 속도 테스트와 함께 SSIM, SNR 등의 정확도 지표를 종합적으로 검토하는 것이 중요합니다.

리뷰 피드백 반영

리뷰 과정에서 mickqian님이 지적한 테스트 파일 경로 문제는 중요한 부분이었습니다. 새로운 테스트 파일 test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py가 SGLang의 테스트 등록 규칙(test/registered/<kind>/<subsystem>/)에 맞지 않아 린트 실패를 유발했습니다. 이 피드백을 통해 테스트 파일이 올바른 경로로 이동되었고, README의 참조 경로도 업데이트되었습니다. 이는 코드 품질 유지 및 CI/CD 파이프라인의 원활한 작동을 위해 필수적인 절차입니다.

결론

이번 PR은 SGLang의 SubBlock Sparse Attention에 SM90 GPU를 위한 Sage FP8 Compute 모드를 성공적으로 도입했습니다. 이는 최신 하드웨어 기능을 활용하고 양자화 기술을 적용하여 상당한 성능 향상을 달성하면서도 정확도를 거의 유지하는 뛰어난 성과입니다. 이러한 최적화는 LLM 및 멀티모달 모델의 추론 속도를 높여 사용자 경험을 개선하고, 더 효율적인 AI 서비스 제공을 가능하게 할 것입니다. 앞으로도 SGLang은 지속적인 최적화를 통해 AI 모델의 성능 한계를 넓혀나갈 것으로 기대됩니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] NVFP4 MoE All-to-All 성능 최적화: Phased Dispatch 기법 분석
현재글 : [sglang] NVIDIA SM90 GPU를 위한 SGLang SubBlock Sparse Attention 최적화: Sage FP8 Compute 도입
다음글 [Liger-Kernel] Ascend NPU 성능 극대화: Liger-Kernel의 커널 최적화 분석