본문으로 건너뛰기

[sglang] [AMD gfx950] GLM-5.2 MLA 최적화: FP8 양자화와 Zero-Copy 레이아웃 전환

PR 링크: sgl-project/sglang#30519 상태: Merged | 변경: +29 / -11

들어가며

최근 대규모 언어 모델(LLM) 최적화의 핵심 중 하나는 MLA(Multi-head Latent Attention) 구조를 얼마나 효율적으로 처리하느냐에 달려 있습니다. 특히 AMD의 최신 가속기인 gfx950(MI355X) 환경에서 GLM-5.2 모델을 실행할 때, 기존에는 가중치 로딩 방식과 메모리 레이아웃의 불일치로 인해 하드웨어의 잠재력을 100% 활용하지 못하는 병목 지점이 있었습니다.

이번 PR(#30519)은 GLM-5.2의 MLA 연산 과정에서 발생하는 두 가지 주요 병목을 해결합니다.

  1. BF16 가중치로 인한 Slow Path: 가중치가 BF16으로 로드되어 고성능 FP8 커널 대신 느린 torch.bmm을 사용하던 문제.
  2. 불필요한 메모리 복사: 연산 결과의 레이아웃이 맞지 않아 flatten 과정에서 발생하는 at::native::direct_copy 오버헤드.

이 글에서는 실제 코드 변경 사항을 통해 시니어 엔지니어가 어떻게 하드웨어 특성에 맞춰 소프트웨어를 최적화하는지 살펴보겠습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. 가중치 로드 시 FP8 양자화 강제 (deepseek_weight_loader.py)

기존 GLM-5.2 모델은 kv_b_proj 가중치를 BF16으로 로드했습니다. 이로 인해 실행 시점에 forward_mla 함수는 고성능 aiter 커널을 사용하지 못하고 rocBLAS 기반의 일반적인 torch.bmm으로 폴백(fallback)되었습니다.

Before:

# 별도의 양자화 로직 없이 BF16 상태로 유지됨
w_kc, w_vc = w.unflatten(
    0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

After:

# GLM-5.2 모델이고 gfx950 환경일 경우, 로드 시점에 FP8(e4m3fn)로 즉시 양자화
if (
    _use_aiter_gfx95
    and self.config.architectures
    and self.config.architectures[0] == "GlmMoeDsaForCausalLM"
    and w.dtype == torch.bfloat16
):
    w, self_attn.w_scale = input_to_float8(w, dtype=torch.float8_e4m3fn)

w_kc, w_vc = w.unflatten(
    0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

왜 좋은가? 가중치를 로드할 때 미리 fp8_e4m3fn으로 변환해두면, 추론(Inference) 루프마다 발생하는 양자화 비용을 제거할 수 있습니다. 또한, 이를 통해 aiter 라이브러리의 batched_gemm_a8w8 커널을 사용할 수 있는 조건을 충족시켜 연산 속도를 비약적으로 높입니다.

2. Zero-Copy를 위한 Batch-Major 레이아웃 전환 (forward_mla_rocm.py)

가장 흥미로운 최적화는 메모리 레이아웃 변경입니다. 기존에는 GEMM 연산 결과가 (heads, batch, dim) 형태여서, 이후 o_proj로 넘어가기 위해 flatten을 호출할 때 메모리 재배치(copy)가 발생했습니다.

Before:

attn_bmm_output = (
    batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
        X=attn_output,
        WQ=attn.w_vc.transpose(-1, -2),
        w_scale=attn.w_scale,
        group_size=128,
        YQ=None,
        transpose_bm=False, # 기존: Batch-Major가 아님
        transpose_bm_in=True,
        dtype=torch.bfloat16,
    )
)

After:

# 1. 출력을 담을 버퍼를 (batch, heads, dim) 순서로 미리 할당
_bmm_buf = torch.empty(
    attn_output.shape[0],
    attn.num_local_heads,
    attn.w_vc.shape[-1],
    device=attn_output.device,
    dtype=torch.bfloat16,
)
# 2. transpose_bm=True를 설정하여 커널이 직접 Batch-Major로 쓰도록 유도
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
    X=attn_output,
    WQ=attn.w_vc.transpose(-1, -2),
    w_scale=attn.w_scale,
    group_size=128,
    YQ=_bmm_buf,
    transpose_bm=True, # 변경: Batch-Major 출력 활성화
    transpose_bm_in=True,
    dtype=torch.bfloat16,
)

왜 좋은가? transpose_bm=True를 설정하면 커널이 데이터를 쓸 때부터 최종적으로 필요한 모양인 (batch, heads, dim) 순서로 기록합니다. 이렇게 하면 이후 코드에서 수행되는 flatten() 연산이 실제 데이터를 복사하지 않고 메타데이터만 바꾸는 Free View가 됩니다. 레이어당 약 5us의 direct_copy 시간을 절약할 수 있는데, 수백 개의 레이어를 가진 LLM에서는 무시할 수 없는 차이를 만듭니다.

성능 향상 및 결과 분석

이번 최적화의 결과는 수치로 증명되었습니다. MI355X TP4 환경에서 GLM-5.2-MXFP4 모델을 테스트한 결과입니다.

  • Throughput (TPUT): 낮은 동시성(Concurrency 4)에서 최대 14% 향상.
  • Latency (TPOT): Median TPOT 기준 최대 12.5% 감소.
  • 정확도: GSM8K 벤치마크 결과, 오차 범위 내(0.928 -> 0.922)로 유지되어 기능적 무결성을 확인했습니다.

리뷰어인 Jacob0226은 추가 테스트를 통해 ITL(Inter-Token Latency)이 약 4.1% 개선되었음을 확인하며 이 PR의 유효성을 뒷받침했습니다.

일반적인 교훈: 시니어의 관점

  1. 하드웨어 특화 게이팅: _use_aiter_gfx95와 같은 조건을 통해 특정 하드웨어(gfx950)에서만 최적화가 동작하도록 설계하여, 다른 환경(NVIDIA 등)에서의 회귀(Regression)를 방지했습니다.
  2. 데이터 레이아웃의 중요성: 딥러닝 최적화는 단순히 연산 속도를 높이는 것을 넘어, 연산 사이의 '데이터 이동'과 '복사'를 줄이는 것이 핵심입니다. transpose_bm 옵션 하나가 전체 파이프라인의 Zero-copy를 가능하게 했습니다.
  3. 로드 타임 오버헤드 전가: 런타임에 반복되는 작업(양자화)을 모델 로드 시점으로 옮기는 것은 전형적이지만 강력한 최적화 기법입니다.

이 PR은 AMD 하드웨어 생태계에서 SGLang이 얼마나 정교하게 최적화되고 있는지를 보여주는 좋은 사례입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글