[sglang] sglang, DeepSeek V4.1 모델을 위한 Paged KV Cache 최적화: V41 및 FP4 포맷 도입

PR 링크: sgl-project/sglang#39123 상태: Merged | 변경: +2313 / -247

들어가며

최근 대규모 언어 모델(LLM)의 발전은 놀라운 속도로 이루어지고 있으며, 이러한 모델들을 효율적으로 서빙하는 것은 중요한 과제가 되었습니다. 특히, LLM의 추론 과정에서 발생하는 메모리 사용량과 지연 시간은 모델의 확장성과 사용자 경험에 직접적인 영향을 미칩니다.

sglang 레포지토리의 이번 Pull Request(PR)는 DeepSeek V4.1 모델의 Key-Value (KV) 캐시 레이아웃을 최적화하여 이러한 문제를 해결하고자 합니다. 기존의 V4 레이아웃을 개선한 V41 (528 B/token) 및 V41_FP4 (288 B/token) 포맷을 도입함으로써, 메모리 사용량을 줄이고 처리 속도를 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 기술적인 관점에서 어떤 의미를 가지는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 kv_layout.cuh, c1.cuh, c2.cuh 등 여러 CUDA 커널 파일과 파이썬 스크립트에서 이루어졌습니다. 주요 변경 사항을 파일별로 나누어 살펴보겠습니다.

1. kv_layout.cuh 및 관련 파일: 새로운 KV Cache 레이아웃 정의

가장 근본적인 변화는 새로운 KV 캐시 레이아웃인 V41과 V41_FP4의 정의입니다. 이 레이아웃들은 KVLayoutTraits와 PagedKV 클래스를 통해 구현되었습니다.

• V41 레이아웃: 기존 V4 레이아웃(584 B/token)보다 효율적인 528 B/token을 목표로 합니다. 이 포맷은 fp8 압축을 사용하며, 512개의 RoPE 차원, 16개의 ue8m0 스케일(32개 값당 하나)을 포함합니다.
• V41_FP4 레이아웃: 추가적인 압축을 통해 288 B/token을 달성합니다. 이는 512개의 e2m1 코드와 32개 값당 하나의 e4m3 스케일을 저장하며, 주로 extra-cache 시나리오에 사용됩니다. fp4 압축은 한 번만 수행되어 효율성을 높입니다.

kv_layout.cuh 파일에는 이러한 레이아웃을 관리하기 위한 KVLayout enum과 PagedKV 템플릿 클래스가 추가되었습니다. PagedKV<layout, page_bits>는 페이지, 행, 스케일 행 주소 지정 방식을 컴파일 타임 상수로 처리하여 런타임 오버헤드를 줄입니다.

또한, store_row<layout, vec> 함수는 벡터화된 저장(vectorized stores)을 지원하며, ue8m0 지수(exponent)를 amax 비트에서 직접 가져오는 방식으로 최적화되었습니다. 이는 기존의 나눗셈 연산을 제거하여 성능을 향상시킵니다.

2. c1.cuh 및 c2.cuh: KV Cache 쓰기 커널 업데이트

c1.cuh (Ratio-1 compressor)와 c2.cuh (Ratio-2 compressor) 파일에서는 기존의 KV 캐시 쓰기 커널들이 새로운 레이아웃을 지원하도록 수정되었습니다.

• 템플릿 매개변수 추가: flash_c1_decode_kernel 및 flash_c2_decode_kernel 함수에 deepseek_v4::KVLayout kLayout 템플릿 매개변수가 추가되었습니다. 이를 통해 커널은 컴파일 타임에 어떤 KV 캐시 레이아웃을 사용할지 결정하고, 해당 레이아웃에 맞는 최적화된 코드를 생성합니다.

• 레이아웃별 처리 분기: 커널 내부에서 if constexpr (kLayout == KVLayout::V41_FP4)와 같은 조건문을 사용하여 각 레이아웃에 따른 특화된 로직을 수행합니다. 예를 들어, V41_FP4 레이아웃의 경우, fp4 압축이 이미 완료된 데이터를 직접 저장하는 방식으로 처리됩니다.

• store_row 함수 호출: 새로운 store_row 함수는 kLayout을 인자로 받아, 해당 레이아웃에 맞는 데이터 저장 및 스케일 처리를 수행합니다. 이는 코드의 재사용성을 높이고 유지보수를 용이하게 합니다.

• kPageBytes 계산 변경: kPageBytes 계산 방식이 deepseek_v4::kv_page_bytes<kLayout>(kPageSize)를 사용하도록 변경되어, 각 레이아웃의 특성에 맞는 페이지 크기를 동적으로 계산합니다.

3. torch_quant.py: 순수 PyTorch 참조 구현

torch_quant.py 파일은 CUDA 커널의 동작을 검증하기 위한 순수 PyTorch 기반의 참조 구현입니다. 이 스크립트는 CUDA 커널과 바이트 단위로 동일한 결과를 내도록 하여, 커널의 정확성을 보장하는 데 중요한 역할을 합니다. 새로운 KV 캐시 포맷에 대한 검증 로직도 포함됩니다.

4. dequant_k_cache.py: V4.1 페이지 역양자화

dequant_k_cache.py 스크립트는 V4.1 형식의 KV 캐시 페이지를 bf16 형식의 prefill 작업 공간으로 역양자화(dequantize)하는 기능을 담당합니다. 이는 압축된 KV 캐시 데이터를 모델이 이해할 수 있는 형태로 변환하는 과정입니다.

5. DeepseekV4MemoryPool: 메모리 관리 및 레이아웃 설정

DeepseekV4MemoryPool 클래스는 KV 캐시 관리를 담당하며, 새로운 레이아웃 옵션을 지원하도록 업데이트되었습니다.

• kv_layout 및 compressed_kv_layout 속성: 각 토큰당 바이트 수, 페이지 바이트 수, 레이어별 접근자 등 KV 캐시 레이아웃 정보를 저장합니다.
• 환경 변수 지원: SGLANG_DSV4_KV_LAYOUT 환경 변수를 통해 v4, v41, auto 옵션을 지원하여 사용자가 KV 캐시 레이아웃을 동적으로 선택할 수 있게 합니다. SGLANG_DSV4_COMPRESSED_KV_LAYOUT 환경 변수는 압축 시 사용할 포맷(fp8, fp4, auto)을 지정합니다.
• 백엔드 로직 변경: 백엔드는 메모리 풀에서 레이아웃과 토큰당 바이트 수를 읽어오도록 변경되어, 하드코딩된 584 바이트 대신 동적으로 설정된 값을 사용합니다.

왜 이게 좋은가?

이번 PR은 다음과 같은 이유로 성능 및 효율성 측면에서 큰 이점을 제공합니다.

1. 메모리 사용량 감소

• V41 (528 B/token): 기존 V4 (584 B/token) 대비 약 9.6%의 메모리 절감 효과를 제공합니다.
• V41_FP4 (288 B/token): 추가적인 fp4 압축을 통해 약 50.7%의 메모리 절감 효과를 제공합니다. 이는 특히 배치 크기가 크거나 시퀀스 길이가 길어질 때 상당한 메모리 부담을 줄여줍니다.

2. 처리 속도 향상

PR 설명에 포함된 스토어 커널 지연 시간(latency) 표를 보면, 새로운 V41 레이아웃을 사용할 때 전반적으로 지연 시간이 감소하는 것을 확인할 수 있습니다. 예를 들어, store 커널의 경우 배치 크기 512에서 1.49us에서 1.39us로 약 6.7%의 성능 향상이 있었습니다. 다른 커널들에서도 유사한 성능 개선이 관찰됩니다.

kernel	layout	bs=1	bs=8	bs=64	bs=512
store	V41	1.02 -> 0.96	1.04 -> 0.98	1.09 -> 1.02	1.49 -> 1.39
store + RoPE	V41	1.02 -> 0.98	1.04 -> 1.00	1.09 -> 1.04	1.54 -> 1.41
fused K norm + RoPE	V41	1.21 -> 1.17	1.19 -> 1.17	1.23 -> 1.20	1.62 -> 1.62
c1	V41	1.47 -> 1.43	1.54 -> 1.41	1.52 -> 1.50	2.05 -> 1.95
c2	V41	1.82 -> 1.78	1.88 -> 1.84	1.97 -> 1.92	2.62 -> 2.52

참고: 위 표는 Baseline 대비 해당 PR 적용 후의 중앙값(median) 지연 시간(us) 변화를 보여줍니다. V4 레이아웃은 변경되지 않았습니다.

3. 유연성 및 확장성 증대

• 동적 레이아웃 선택: 환경 변수를 통해 사용자가 KV 캐시 레이아웃을 선택할 수 있게 됨으로써, 특정 하드웨어 환경이나 모델 요구사항에 맞춰 최적의 설정을 적용할 수 있습니다.
• 모듈식 설계: KVLayoutTraits와 PagedKV와 같은 추상화 계층을 도입하여, 새로운 KV 캐시 포맷이나 압축 기법을 추가하더라도 기존 코드에 미치는 영향을 최소화할 수 있습니다. 이는 향후 모델 및 기술 발전에 대한 유연한 대응을 가능하게 합니다.

4. 일반적인 교훈

• 양자화(Quantization)의 중요성: LLM 추론에서 양자화는 메모리 대역폭 병목 현상을 완화하고 처리량을 높이는 데 핵심적인 역할을 합니다. 특히 fp8, fp4와 같은 저정밀도 포맷을 효과적으로 활용하는 것이 중요합니다.
• 컴파일 타임 최적화: 가능한 경우, 템플릿 메타프로그래밍과 같은 기법을 사용하여 컴파일 타임에 레이아웃 관련 상수를 결정하고 분기를 처리하면 런타임 오버헤드를 크게 줄일 수 있습니다.
• 검증을 위한 참조 구현: 복잡한 CUDA 커널의 정확성을 보장하기 위해, 순수 PyTorch와 같은 참조 구현을 유지하는 것은 필수적입니다.

리뷰 피드백 반영

PR 설명에는 리뷰 댓글이 직접적으로 포함되어 있지 않지만, 코드 자체에 반영된 내용들을 통해 개발자의 깊은 고민을 엿볼 수 있습니다. 예를 들어, static_assert를 사용하여 다양한 제약 조건(헤드 차원, 로프 차원, 페이지 비트 등)을 컴파일 타임에 검증하는 것은 잠재적인 오류를 조기에 발견하고 코드의 견고성을 높이는 데 기여합니다. 또한, PDLTriggerSecondary와 같은 메커니즘은 병렬 처리 환경에서의 동기화 및 데이터 일관성을 보장하기 위한 설계로 보입니다.

결론

이번 sglang PR은 DeepSeek V4.1 모델을 위한 KV 캐시 레이아웃을 최적화함으로써, LLM 서빙의 효율성을 한 단계 끌어올렸습니다. V41 및 V41_FP4 포맷의 도입은 메모리 사용량을 획기적으로 줄이고 처리 속도를 향상시키는 동시에, 유연하고 확장 가능한 아키텍처를 제공합니다. 이러한 최적화는 LLM을 더 많은 사용자에게 더 빠르고 효율적으로 제공하는 데 중요한 기여를 할 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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