[sglang] SGLang, FP4 KV 캐시 도입으로 LLM 추론 성능 극대화: NVFP4 최적화 분석

2026년 4월 29일수정: 2026년 4월 29일

PR 링크: sgl-project/sglang#21954 상태: Merged | 변경: +None / -None

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부시지만, 그만큼 모델의 크기와 추론 시 요구되는 컴퓨팅 자원 또한 기하급수적으로 증가하고 있습니다. 특히, LLM 추론 성능의 병목 지점 중 하나는 바로 KV 캐시(Key-Value Cache)입니다. KV 캐시는 이전 토큰들의 계산 결과를 저장하여 다음 토큰 생성 시 재활용함으로써 추론 속도를 높이는 핵심적인 역할을 합니다. 하지만 이 KV 캐시가 차지하는 메모리 용량과 대역폭은 상당하여, 모델의 크기가 커질수록 심각한 성능 저하를 유발합니다.

이러한 문제를 해결하기 위해 SGLang 프로젝트에서는 "[1/4] NVFP4 KV cache: quantization strategy abstraction and kernel" PR을 통해 FP4 (4-bit Floating Point) 양자화를 활용한 새로운 KV 캐시 저장 방식을 도입했습니다. 이 PR은 특히 SM120 GPU 아키텍처를 타겟으로 하며, 기존 FP8 KV 캐시에 비해 정확도를 유지하면서도 추론 성능을 크게 향상시키는 것을 목표로 합니다. 본 글에서는 이 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 최적화가 LLM 추론 성능 향상에 기여하는지 기술적인 관점에서 깊이 있게 탐구하고자 합니다.

코드 분석

이번 PR은 주로 KV 캐시 양자화 전략을 추상화하고, 이를 위한 새로운 커널 및 유틸리티를 도입하는 데 초점을 맞추고 있습니다. 주요 변경 사항은 다음과 같이 파일별로 나누어 살펴볼 수 있습니다.

1. `python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py` (신규 파일)

이 파일은 FP4 KV 캐시 양자화 전략을 위한 추상 기본 클래스(FP4KVCacheQuantMethod)와 구체적인 구현 클래스들을 정의합니다. 전략 패턴(Strategy Pattern)을 사용하여 다양한 양자화 방식을 유연하게 적용할 수 있도록 설계되었습니다.

FP4KVCacheQuantMethod (ABC): 양자화 및 역양자화 로직을 캡슐화하는 추상 클래스입니다. create_buffers, quantize_and_store, dequantize_prev_kv 등의 추상 메서드를 정의하여, 각 양자화 전략이 구현해야 할 인터페이스를 명확히 합니다. 이 클래스는 CUDA 그래프 호환성을 염두에 두고 설계되었습니다.
NVFP4KVMethod: SM100 및 SM120 GPU를 위한 구체적인 FP4 양자화 전략을 구현합니다. 이 방식은 두 단계의 스케일링을 사용합니다: 첫 번째는 전역(per-tensor) FP32 스케일, 두 번째는 블록(per-block) FP8 E4M3 스케일입니다. 이는 기존 FP8 방식보다 더 높은 압축률을 제공하면서도 정확도를 유지하기 위한 설계입니다.
BlockFP4KVMethod: 블록 단위의 단일 레벨 스케일링을 사용하는 FP4 양자화 방식입니다. block_size=16으로 설정되어 MXFP4와 유사하지만, 블록 크기에서 차이가 있습니다.
FP4_KV_CACHE_QUANT_REGISTRY: 레시피 이름과 해당 양자화 메서드 클래스를 매핑하는 레지스트리입니다. 이를 통해 설정 파일 등에서 원하는 양자화 방식을 쉽게 선택하고 로드할 수 있습니다.
load_scales_from_model 메서드: 모델의 가중치로부터 FP32 전역 스케일을 로드하는 기능을 수행합니다. 이는 NVFP4KVMethod에서 사용되며, 모델 아키텍처에 따라 동적으로 스케일 값을 조정합니다. 특히 SM100과 SM120 간의 스케일링 계수 차이(E2M1_MAX 곱셈 여부)를 처리하는 로직이 포함되어 있습니다.
create_buffers 메서드: 각 양자화 전략에 필요한 KV 캐시 버퍼, 스케일 버퍼 등을 할당합니다. NVFP4KVMethod의 경우, FP4 데이터(k_buffer, v_buffer), 블록 스케일(k_scale_buffer, v_scale_buffer), 그리고 역양자화 시 사용될 FP8 E4M3 형식의 작업 공간(dq_k_buffer, dq_v_buffer)을 할당합니다.
quantize_and_store 메서드: 입력받은 FP16/BF16 KV 데이터를 FP4로 양자화하고, 해당 스케일과 함께 지정된 버퍼에 저장합니다. 이때 NVFP4KVQuantizeUtil의 quantize 함수를 사용합니다.
dequantize_prev_kv 메서드: 저장된 FP4 KV 데이터와 스케일을 FP8 E4M3 형식으로 역양자화합니다. 이는 주로 디코딩 단계에서 사용되며, FlashInfer의 FP8 prefill 커널에 입력으로 사용될 수 있도록 준비하는 과정입니다.

2. `python/sglang/srt/layers/quantization/kvfp4_tensor.py` (확장)

이 파일은 FP4 양자화 및 역양자화 유틸리티를 포함합니다. 기존의 KVFP4QuantizeUtil은 BlockFP4KVQuantizeUtil로 이름이 변경되었으며, 새로운 NVFP4KVQuantizeUtil이 추가되었습니다.

NVFP4KVQuantizeUtil: FlashInfer의 nvfp4_kv_quantize 및 nvfp4_kv_dequantize 커널을 사용하거나, fallback으로 PyTorch 기반의 구현을 사용합니다. 이 클래스는 NVFP4KVMethod에서 FP4 양자화 및 역양자화 연산을 수행하는 데 사용됩니다.
BlockFP4KVQuantizeUtil: 블록 단위 FP4 양자화를 처리합니다. block_size=16을 사용하며, 기존 KVFP4QuantizeUtil의 기능을 계승합니다. 리뷰어의 피드백에 따라 이름이 변경되었지만, 이전 버전과의 호환성을 위해 별칭(alias)이 유지되었습니다.

3. Unit Tests (신규)

새롭게 추가된 양자화 메서드 및 유틸리티에 대한 단위 테스트가 포함되었습니다. 레지스트리 동작 검증, 버퍼 형태 검증, 그리고 양자화 후 역양자화 시 원본 데이터와의 roundtrip 정확도 검증을 수행합니다. 특히 NVFP4는 CUDA 환경에서, BlockFP4는 CPU 환경에서 테스트됩니다.

왜 이게 좋은가?

이번 PR에서 도입된 FP4 KV 캐시 양자화 전략은 LLM 추론 성능 향상에 여러 가지 중요한 기여를 합니다.

1. 메모리 사용량 감소 및 대역폭 효율 증대

FP4 양자화는 기존 FP16/BF16 대비 KV 캐시 메모리 사용량을 1/4로 줄입니다. 이는 모델이 더 긴 시퀀스를 처리할 수 있도록 하거나(long context), 동일 하드웨어에서 더 많은 배치 크기를 처리할 수 있게 하여 처리량을 높입니다. 또한, 메모리 대역폭 요구량도 크게 줄어들어, 메모리 I/O가 병목인 환경에서 특히 효과적입니다.

2. 추론 속도 향상 (특히 디코딩 단계)

PR에서 제공된 벤치마크 결과는 NVFP4 KV 캐시가 디코딩 단계에서 FP8 KV 캐시 대비 1.18배의 속도 향상을 보여줍니다. 이는 FP4 데이터 자체의 크기가 작아 메모리 접근이 빨라지고, 최적화된 커널(XQA kernel)을 통해 효율적으로 처리되기 때문입니다. Prefill 단계에서는 약간의 성능 저하(0.99x)가 관찰되지만, 이는 디코딩 단계에서의 큰 이득으로 상쇄됩니다. 이는 LLM 추론의 주요 병목 구간인 디코딩 성능을 직접적으로 개선한다는 점에서 매우 중요합니다.

메모리 압축은 LLM 추론 성능의 핵심 병목을 해결하는 강력한 수단입니다. 특히 KV 캐시와 같이 대용량 데이터를 효율적으로 관리하는 것이 중요합니다.
정교한 양자화 기법(예: 다단계 스케일링)은 압축률을 높이면서도 정확도를 유지하는 데 필수적입니다.
하드웨어 아키텍처(SM100 vs SM120)의 특성을 고려한 최적화는 성능 극대화에 중요합니다.
전략 패턴과 같은 디자인 패턴은 복잡한 시스템의 유연성과 확장성을 높이는 데 기여합니다.

리뷰 피드백 반영

이번 PR은 여러 리뷰어들의 꼼꼼한 검토를 거쳤으며, 몇 가지 중요한 피드백이 반영되었습니다.

명명 규칙 개선: samuellees 리뷰어는 파일명과 클래스명에 "FP4"를 명시적으로 포함하도록 제안했습니다 (fp4_kv_cache_quant_method.py, NVFP4KVMethod 등). 이는 코드의 가독성과 명확성을 높이는 데 기여했습니다.
GPU 이름 제거: 특정 GPU 아키텍처 이름(예: SM120)을 클래스명에서 제거하여 일반성을 확보하라는 피드백이 있었습니다. 이는 코드의 재사용성을 높입니다.
상수 활용: 스케일링 계수 6.0에 대해 E2M1_MAX 상수를 재사용하도록 권장되었습니다. 이는 코드의 일관성과 유지보수성을 향상시킵니다.
FlashInfer 커널 우선 사용: 양자화/역양자화 유틸리티에서 FlashInfer의 최적화된 커널(nvfp4_kv_quantize, nvfp4_kv_dequantize)을 우선적으로 사용하도록 명확히 했습니다. 이는 성능 향상의 핵심 요소입니다.
KVFP4QuantizeUtil 이름 변경: samuellees와 b8zhong의 논의를 통해, 기존 KVFP4QuantizeUtil의 용도(MXFP4와 유사하나 block_size=16)를 더 명확히 하기 위해 BlockFP4KVQuantizeUtil로 이름 변경이 이루어졌습니다. 다만, 이전 API와의 호환성을 위해 별칭이 유지되었습니다.
CUDA 그래프 호환성: b8zhong의 질문에 따라, dequantize_prev_kv 메서드가 FP8 형식으로 결과를 반환하는 이유가 FP4 prefill 커널의 부재 때문이며, 이는 XQA 커널과의 호환성을 위한 설계임을 명확히 했습니다. 또한, 모든 연산이 FlashInfer 커널 또는 순수 텐서 연산을 사용하므로 CUDA 그래프 호환성을 유지한다는 점이 확인되었습니다.

이러한 리뷰 피드백의 반영은 코드의 품질, 명확성, 그리고 성능을 더욱 향상시키는 데 결정적인 역할을 했습니다.

References

참고 자료

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

PR Analysis 의 다른글

이전글 [cpython] Python subprocess.communicate() 타임아웃 성능 개선: 느린 자식 프로세스 응답 방식 변경
현재글 : [sglang] SGLang, FP4 KV 캐시 도입으로 LLM 추론 성능 극대화: NVFP4 최적화 분석
다음글 [sglang] FlashInfer TRTLLM-Gen MoE 커널 최적화: NemotronH 모델 지원 및 성능 향상