vLLM 숨겨진 상태 추출 시스템

vLLM은 추측 디코딩 모델의 훈련 및 개발을 용이하게 하기 위해 네이티브 숨겨진 상태 추출 시스템을 도입했습니다(vllm>=v0.18.0에 포함, PR #33736 통해). 이 시스템을 사용하면 vLLM의 핵심 성능 최적화를 손상시키지 않으면서 토큰 시퀀스의 내부 중간 표현을 추출할 수 있습니다.

숨겨진 상태 추출 병목 해결

숨겨진 상태를 추출하는 것은 추측 디코딩에서 초안 모델을 훈련하는 데 핵심적입니다(작은 모델이 토큰을 예측하고 더 큰 “verifier” 모델이 이를 확인하는 방식). 초안 모델에 verifier의 내부 상태를 제공하면 정렬과 품질을 향상시킬 수 있기 때문입니다. 이전에는 연구자들이 두 가지 최적이 아닌 방법에 의존했습니다:

  • transformers 라이브러리 사용: 이 접근 방식은 vLLM의 분산 지원 및 성능 최적화를 잃게 하고, transformers와 vLLM 숨겨진 상태 간 불일치로 인해 버그가 발생할 수 있습니다.
  • vLLM 내부 패치: 내부 API를 수동으로 호출하면 유지 보수 부담이 크게 증가하고, prefix caching, auto-batching, async 서버와 같은 핵심 기능을 비활성화해야 합니다.

기술 설계 및 구현

vLLM 숨겨진 상태 추출 시스템은 표준 추론의 “핫 경로”에서 오버헤드를 피하면서 대형 숨겨진 상태 텐서의 막대한 메모리 요구 사항을 관리하도록 설계되었습니다. 숨김 크기가 4096인 Qwen3-8B 모델의 경우, 8k 토큰 시퀀스에 대해 4개의 레이어를 추출하면 약 268 MB의 데이터가 필요합니다.

“더미 모델” 아키텍처

새로운 런타임 오버헤드를 도입하지 않고 이를 구현하기 위해 vLLM은 기존 인프라를 활용합니다:

  1. Eagle-3 Plumbing: vLLM은 이미 Eagle-3 추측 디코딩을 위해 verifier 모델에서 초안 모델로 숨겨진 상태를 이동시키는 파이프라인을 보유하고 있습니다.
  2. 더미 초안 모델: 시스템은 이러한 숨겨진 상태를 받는 더미 초안 모델을 생성합니다. 어텐션 연산을 수행하는 대신, 이 모델은 더미 어텐션 레이어를 사용해 숨겨진 상태를 직접 자신의 KV 캐시에 삽입합니다.
  3. KV Connector API: vLLM은 확장 가능한 KV Connector API를 활용합니다—원래 Prefill/Decode Disaggregation에서 KV 캐시 추출을 위해 설계된 이 API를 사용해 더미 초안 모델의 KV 캐시(이제 숨겨진 상태를 포함)를 디스크에 저장하거나 다른 프로세스로 전송합니다.

숨겨진 상태를 KV 캐시 데이터로 취급함으로써 vLLM은 동일한 페이지드 메모리 시스템을 사용해 이를 관리할 수 있으며, 이는 prefix caching 및 청크형 prefill과의 호환성을 보장합니다.

사용법 및 제한 사항

사용자는 Python API를 통해 또는 특정 설정으로 vLLM 서버를 실행하여 숨겨진 상태를 추출할 수 있습니다. 서버는 더미 초안 모델을 설정하기 위한 --speculative_config와 커넥터를 정의하기 위한 --kv_transfer_config 두 가지가 필요합니다.

vllm serve Qwen/Qwen3-8B --speculative_config '{
	"method": "extract_hidden_states", 
	"num_speculative_tokens": 1, 
	"draft_model_config": {
		"hf_config": {
			"eagle_aux_hidden_state_layer_ids": [3, 18, 33, 36]
		}
	}
}' --kv_transfer_config '{
	"kv_connector": "ExampleHiddenStatesConnector", 
	"kv_role": "kv_producer", 
	"kv_connector_extra_config": {
		"shared_storage_path": "/tmp/hidden_states"
	}
}'

핵심 제약 조건:

  • 출력 형식: 서버는 kv_transfer_params 사전을 반환하며, 여기에는 token_ids와 hidden_states를 포함하는 .safetensors 파일을 가리키는 hidden_states_path가 포함됩니다.
  • 범위: 프롬프트 토큰과 해당 숨겨진 상태만 저장됩니다. 팀은 max_tokens=1인 v1/completions 엔드포인트 사용을 권장합니다.
  • 병렬성: 시스템은 단일 노드 다중 GPU 배포를 위해 --tensor-parallel-size와 --data-parallel-size를 지원합니다.

향후 로드맵

  • Speculators 통합: speculators 라이브러리(v0.5.0)는 PR #353을 통해 이 네이티브 vLLM 시스템을 사용하도록 업데이트되었으며, 초안 모델의 온라인 훈련을 가능하게 합니다.
  • 비동기 쓰기: 현재 ExampleHiddenStatesConnector는 블로킹 쓰기를 사용합니다; 향후 업데이트에서는 성능 향상을 위해 비동기 쓰기를 구현할 예정입니다.
  • 디바이스 간 전송: 디스크 기반 저장을 넘어선 방식으로, vLLM은 숨겨진 상태를 디바이스 간에 직접 전송하는 커넥터를 개발 중이며, 멀티 노드 환경 지원도 포함됩니다.

Sources