[sglang] AMD MI350x에서 Qwen3.5의 Long-Context Prefill 성능을 극대화하는 FP8 FMHA 최적화 분석

PR 링크: sgl-project/sglang#32046 상태: Merged | 변경: +38 / -0

들어가며

대규모 언어 모델(LLM)의 성능은 모델 자체의 능력뿐만 아니라, 이를 구동하는 하드웨어 및 소프트웨어 스택의 최적화에 크게 좌우됩니다. 특히 긴 컨텍스트(Long-Context)를 처리하는 능력은 최신 LLM의 핵심 요구사항 중 하나이며, 이를 위한 Prefill 단계의 효율성은 전체 응답 시간(TTFT, Time To First Token)에 결정적인 영향을 미칩니다.

오늘 분석할 GitHub PR은 sgl-project/sglang 레포지토리에서 AMD MI350x (gfx950) GPU 환경에서 Qwen3.5 모델의 긴 컨텍스트 Prefill 성능을 획기적으로 개선하기 위한 최적화를 다룹니다. 기존에는 Qwen3.5 모델이 head_dim 256을 사용할 때, 고속 aiter 어텐션 경로를 타지 못하고 느린 mha_batch_prefill 경로로 폴백하는 문제가 있었습니다. 이 PR은 ROCm/aiter에 새로 추가된 fp8 head_dim-256 FMHA (Flash Attention) 어셈블리 커널을 sglang에 통합하여, 이 특정 조건에서 컴퓨트 바운드 Prefill 어텐션의 속도를 높이는 것을 목표로 합니다.

문제 해결을 위한 접근

이 PR의 핵심은 aiter 라이브러리에 새로 추가된 flash_attn_varlen_fp8_pertensor_func 커널을 sglang의 어텐션 백엔드에 옵트인(opt-in) 방식의 고속 경로로 추가하는 것입니다. 이 커널은 특히 캐시된 프리픽스(cached prefix)가 없는 순수한 Prefill (no cached prefix) 상황에서 긴 컨텍스트의 처리량과 TTFT를 개선하는 데 초점을 맞춥니다. 이 최적화는 SGLANG_AITER_FMHA_FP8_HD256 환경 변수를 통해 활성화되며, gfx950 (MI350x) GPU에서만 작동하도록 설계되었습니다.

코드 분석: aiter_backend.py 변경사항

이 PR은 python/sglang/srt/layers/attention/aiter_backend.py 파일 단 하나에 변경사항을 적용합니다. 주요 변경점은 다음과 같습니다.

1. flash_attn_varlen_fp8_pertensor_func 임포트

aiter 라이브러리에서 새로운 FP8 FMHA 커널을 임포트합니다. 이 커널은 ROCm/aiter PR #3732를 통해 aiter 메인 브랜치에 병합된 기능입니다.

Before:

 try:
     from aiter import (
         flash_attn_varlen_func,
         get_mla_metadata_info_v1,
         get_mla_metadata_v1,
         mha_batch_decode_func,
         mha_batch_prefill_func,
         mha_decode_func,
         mha_prefill_func,
         paged_attention_v2_func,
         reshape_and_cache_v2,
     )

After:

 try:
     from aiter import (
+        flash_attn_varlen_fp8_pertensor_func,
         flash_attn_varlen_func,
         get_mla_metadata_info_v1,
         get_mla_metadata_v1,
         mha_batch_decode_func,
         mha_batch_prefill_func,
         mha_decode_func,
         mha_prefill_func,
         paged_attention_v2_func,
         reshape_and_cache_v2,
     )

2. 새로운 고속 경로 추가

AiterAttnBackend.forward_extend 함수 내에 특정 조건을 만족할 때 새로운 flash_attn_varlen_fp8_pertensor_func 커널을 호출하는 고속 경로를 추가합니다. 이 경로는 다음 조건들이 모두 충족될 때 활성화됩니다:

• SGLANG_AITER_FMHA_FP8_ASM=1 환경 변수가 설정되어 있고, gfx950 (MI350x) GPU에서 실행 중일 때.
• forward_batch.forward_mode.is_extend(): extend 모드일 때.
• not any(forward_batch.extend_prefix_lens_cpu): 캐시된 프리픽스가 전혀 없는 순수한 Prefill일 때.
• window_size == (-1, -1): 슬라이딩 윈도우가 사용되지 않을 때.
• sinks is None: 어텐션 싱크(sinks)가 없을 때.
• self.logits_soft_cap == 0.0: 로짓 소프트 캡이 0.0일 때 (커널이 소프트 캡을 지원하지 않음).
• layer.qk_head_dim == 256 및 layer.v_head_dim == 256: 헤드 차원이 256일 때.
• self.kv_cache_dtype == fp8_dtype: KV 캐시 데이터 타입이 FP8일 때.

이 조건들 중 하나라도 만족하지 않으면, 기존의 mha_batch_prefill_func 경로로 폴백합니다.

Before:

             if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
                 window_size = (layer.sliding_window_size, -1)
 
             if self.kv_cache_is_vectorized_5d:
                 return forward_extend_vectorized_5d(
                     self,

After:

             if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
                 window_size = (layer.sliding_window_size, -1)
 
+            if (
+                get_bool_env_var(

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