본문으로 건너뛰기

[sglang] [AMD MI355X] DeepSeek-V4 최적화: FP8 Quantization의 불필요한 Copy 제거하기

PR 링크: sgl-project/sglang#33165 상태: Merged | 변경: +273 / -3

들어가며

대규모 언어 모델(LLM)의 추론 성능을 극대화하기 위해 FP8 양자화(Quantization)는 이제 선택이 아닌 필수입니다. 특히 DeepSeek-V4와 같은 최신 모델을 AMD의 차세대 가속기인 MI355X (gfx950)에서 구동할 때, 하드웨어 가속 커널인 CK(Composable Kernel)의 성능을 온전히 활용하는 것이 중요합니다.

기존 SGLang의 구현에서는 AMD 환경에서 w8a8 (Weight 8-bit, Activation 8-bit) GEMM을 수행할 때, 활성화 함수(Activation)의 Scale 데이터를 하드웨어가 요구하는 레이아웃으로 맞추기 위해 매번 메모리 복사(Relayout Copy)를 수행하고 있었습니다. 이는 특히 MLA(Multi-head Latent Attention)나 MoE(Mixture of Experts)와 같이 Linear 레이어가 빈번하게 호출되는 구조에서 무시할 수 없는 오버헤드를 발생시킵니다.

본 포스트에서는 이 불필요한 복사 과정을 제거하여 Zero-copy view로 전환하고, 전체적인 추론 성능(TPOT)을 개선한 PR을 분석해 보겠습니다.

문제의 핵심: Column-Major vs Row-Major

AMD MI355X의 gemm_a8w8_blockscale_bpreshuffle 커널은 효율적인 연산을 위해 Activation Scale 데이터가 [num_groups, tokens] 형태의 Column-major 레이아웃으로 배치되어 있기를 기대합니다.

하지만 기존의 aiter_per1x128_quant 함수는 기본적으로 Row-major로 Scale을 생성했고, 이를 하드웨어 요구사항에 맞추기 위해 .t().contiguous().t()와 같은 연산을 수행했습니다. 여기서 .contiguous()는 실제 메모리 상의 데이터를 재배치하며 새로운 텐서를 할당하는 Element-wise Copy를 유발합니다.

코드 분석: Zero-copy로의 전환

1. Quantization 단계에서의 레이아웃 최적화

가장 큰 변화는 양자화 커널을 호출할 때부터 하드웨어가 원하는 레이아웃으로 데이터를 쓰도록(emit) 요청하는 것입니다.

Before:

# 기존: Row-major로 생성 후 수동으로 재배치(Materialize)
q_input, x_scale = aiter_per1x128_quant(
    input_2d,
    quant_dtype=aiter.dtypes.fp8,
    transpose_scale=False, # Row-major
)
if materialize_bpreshuffle_scale:
    x_scale = materialize_bpreshuffle_fp8_scale(x_scale) # 실제 메모리 복사 발생

After:

# 개선: 처음부터 Transposed 레이아웃으로 생성 요청
emit_bpreshuffle_scale = (
    materialize_bpreshuffle_scale and input_2d.shape[0] >= 2
)
q_input, x_scale = aiter_per1x128_quant(
    input_2d,
    quant_dtype=aiter.dtypes.fp8,
    transpose_scale=emit_bpreshuffle_scale, # 하드웨어 친화적 레이아웃
)
if emit_bpreshuffle_scale:
    # 메모리 복사 없이 Stride 정보만 수정하는 View 생성
    x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale)

여기서 view_aiter_fused_rms_transposed_fp8_scale 함수는 내부적으로 torch.as_strided를 사용하여 데이터의 실제 위치는 바꾸지 않고, PyTorch가 이 데이터를 해석하는 방식(Stride)만 변경합니다. 결과적으로 커널 런칭 오버헤드와 메모리 대역폭 소모를 동시에 줄였습니다.

2. M=1 (Single Token) 케이스의 정교한 처리

리뷰어인 kkHuang-amd는 M=1인 경우의 레이아웃 특성을 정확히 짚어냈습니다. 토큰이 하나인 경우 [1, G] 레이아웃과 [G, 1] 레이아웃은 메모리 상의 바이트 순서가 동일합니다. 따라서 이 경우에는 굳이 복잡한 로직을 타지 않고 기존의 materialize 경로를 유지하더라도 PyTorch의 최적화 덕분에 실제 복사가 일어나지 않는 No-op view가 됩니다.

이를 위해 코드에 input_2d.shape[0] >= 2라는 가드(Guard) 조건이 추가되었습니다.

왜 이게 좋은 최적화인가?

1. 성능 수치 (Performance Gains)

DeepSeek-V4 모델을 MI355X TP8 환경에서 벤치마킹한 결과, 다음과 같은 유의미한 성능 향상이 관찰되었습니다.

  • TPOT (Time Per Output Token): 동시성(Concurrency) 수준에 따라 1.2% ~ 2.0% 개선되었습니다.
  • Total Throughput: 전체 토큰 처리량이 약 1.4% ~ 1.7% 향상되었습니다.

수치상으로는 작아 보일 수 있지만, 모델의 모든 Linear 레이어에서 발생하는 고정적인 오버헤드를 제거했다는 점에서 매우 효율적인 최적화입니다.

2. 비트 단위 동일성 (Bit-identical Correctness)

이 최적화는 결과값을 근사하는 것이 아니라, 데이터의 배치 방식만 바꾸는 것이기 때문에 기존 방식과 비트 단위로 완전히 동일한(Bit-identical) 결과를 보장합니다. 이는 모델의 정확도(Accuracy) 손실 없이 순수하게 성능만 올릴 수 있음을 의미합니다.

3. 하드웨어 추상화 계층의 통합

기존에는 as_strided 연산이 코드 곳곳에 산재해 있었으나, 이번 PR을 통해 view_aiter_fused_rms_transposed_fp8_scale이라는 공통 헬퍼 함수로 로직이 통합되었습니다. 이는 유지보수성을 높이고 향후 다른 커널 최적화 시에도 재사용할 수 있는 기반을 마련했습니다.

결론

이번 최적화 사례는 하드웨어 가속기(AMD MI355X)의 특성을 깊이 이해하고, 상위 프레임워크(PyTorch/SGLang) 수준에서 데이터 레이아웃을 선제적으로 맞춤으로써 불필요한 연산을 제거한 훌륭한 사례입니다. 시니어 엔지니어로서 우리는 단순히 기능을 구현하는 것을 넘어, 데이터의 흐름(Data Flow) 속에서 발생하는 숨겨진 복사 비용을 찾아내고 이를 제거하는 데 집중해야 합니다.

이러한 마이크로 최적화들이 모여 DeepSeek-V4와 같은 거대 모델의 실시간 서비스 가능성을 결정짓게 됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글