LLM Preemption 구현을 위한 Async KV Cache Dump/Load 및 Overlap 가능 여부

안녕하세요. 현재 Mobilint NPU에서 LLM serving scheduler를 구현하면서 KV cache preemption을 실험하고 있습니다.

현재 qbruntime 1.3.2 기준으로 dump_cache_memory / load_cache_memory는 동기식 API만 확인되는데, 몇 가지 확인하고 싶습니다.

  1. LLM KV cache의 dump/load를 asynchronous 또는 non-blocking하게 수행할 수 있는 API가 따로 존재하나요?
  2. 공개 API에는 없더라도 KV cache의 host/device transfer와 NPU inference를 동시에 진행할 수 있는 방법이 있나요?
  3. 예를 들어 한 request를 NPU에서 inference하는 동안 다른 request의 KV cache를 dump/load하여 transfer latency를 overlap시키는 것이 runtime 구조상 가능한가요?
  4. 이를 지원한다면 stream, event, future, DMA queue 등 어떤 synchronization/async mechanism을 사용해야 하나요?
  5. 현재 infer_async는 LLM에서는 지원되지 않는 것으로 확인했는데, LLM inference와 KV transfer overlap도 같은 제약 때문에 불가능한 것인지 궁금합니다.
  6. 현재 SDK에서는 지원되지 않지만 향후 runtime에서 async KV dump/load 또는 compute-transfer overlap 지원 계획이 있는지도 궁금합니다.

현재 저희 환경에서는 KV snapshot 약 64 MiB 기준으로 dump가 약 30~60 ms, load가 약 16~18 ms 정도 소요되어, preemption이 반복되면 synchronous KV transfer overhead가 전체 serving 성능에 크게 영향을 주고 있습니다.

따라서 현재 Mobilint NPU에서 이러한 KV transfer latency를 inference 뒤에 숨길 수 있는 공식적인 방법이 있는지 확인 부탁드립니다.

안녕하세요, 모빌린트 김건입니다.

1. Async/Non-blocking API 미지원

현재 qbruntime은 LLM KV-Cache의 load/dump 및 infer 함수에 대한 Async/Non-blocking API를 지원하지 않습니다. 일반 LLM 모델은 하나의 모델당 하나의 KV-Cache 세션만 사용하는 구조라, 동시 연산을 위해 고안된 Async API의 목적과 맞지 않기 때문입니다.

2. Batch Inference로 유사 기능 지원

다만, qbruntime은 모델 하나에 여러 개의 KV-Cache 세션을 두고, 여러 Request를 하나의 Batch로 묶어 동시에 처리하는 LLM batch inference이 가능한 모델들을 지원하고 있습니다. Batch inference 기능을 통해 동시에 여러 Request를 처리할 때의 dump/load로 인한 KV transfer를 뒤에 숨길 수 있을 것으로 예상됩니다.

해당 모델의 특징들은 다음과 같습니다.

  • 컴파일 타임에 정해진 최대 batch_size만큼 KV-Cache 세션을 가지고 있습니다.

  • 각 KV-Cache 세션은 독립적으로 dump/load 가능합니다.

  • 여러 Request를 하나의 Batch로 묶어 한번에 inference를 요청 가능합니다.

따라서 애플리케이션 레벨에서 여러 Request를 동시 처리하면서, 사용 중이지 않은 세션의 캐시를 갱신하는 것이 가능합니다.

Batch Inference 지원 모델의 경우 qbcompiler를 통한 컴파일을 지원하며, Hugging Face에서 사전에 컴파일된 모델도 제공하고 있습니다.

  • qbcompiler의 LlmConfig에서 batch_size를 1보다 큰 값으로 설정

  • 또는 Mobilint Hugging Face에 모델명 뒤 Batch16이 붙은 배포 모델 사용

참조 링크

  1. qbcompiler batch_size 설정

  2. Batch inference 지원 모델 검색 (Hugging Face)

  3. qbruntime batch inference 실행 API (Release Notes)

1 Like

답변 감사드립니다! 제안해 주신 Batch 방식도 확인했습니다.

저희가 Preemption 기반 스케줄링을 구현하고 있어 추가로 여쭤보고 싶습니다. 질문 3-6번에 나온 것처럼, Batch 세션을 쓸 때 Session A 연산과 Session B의 KV Transfer가 실제로 H/W 레벨에서 Overlap되어 실행되는지, 그리고 향후 Async KV API 지원 로드맵이 있으신지 궁금합니다!

1 Like

Batch inference를 지원하는 모델의 경우, 세션마다 독립된 KV-Cache 메모리를 가지고 있으며, 이에 따라 dump/load API 역시 세션 단위로 독립적으로 동작하도록 설계되어 있습니다.

따라서 서로 다른 세션에 대한 dump/load 요청을 멀티 스레드에서 동시에 호출하시면, H/W 레벨에서 overlap이 발생하는 구조입니다.

Python을 사용하신다면 asyncio 모듈을 활용하시거나, 별도의 멀티스레드 기반 비동기 파이프라인을 구성하시는 방향을 추천드립니다.

공식적인 Async KV API 지원은 검토해야 할 사항들이 있어 아직 구체적인 일정을 말씀드리기는 어려우나, 필요성은 계속 인지하고 있으며 내부적으로 논의해 나갈 예정입니다.

1 Like