실시간 응답 중 모델 레이어를 GPU 간 이동시키기 (8ms 소요, 재-프리필 없음)

저는 jitllm이라는 추론 엔진을 개발해 왔습니다. 이 아이디어는 운영체제 및 JIT 컴파일러 분야의 제 배경에서 나왔습니다. 즉, 어떤 하드웨어가 있든 그에 맞는 커널을 생성하고, 레이어를 프로세스처럼 취급하여 엔진이 배치하고 이동시킬 수 있도록 하는 것입니다. 비디오에서는 T4가 장착된 클라우드 VM 환경에서 Llama-3.2-3B의 모든 레이어가 응답 중간에 한 GPU에서 다른 GPU로 그리고 다시 이동하는 것을 볼 수 있습니다. 또한, 제 노트북에서는 모델이 CPU, CUDA, Vulkan (iGPU)에 걸쳐 분할되어 생성되는 동안 재배치됩니다. KV 캐시는 페이지화되며 레이어와 함께 이동하므로 프롬프트를 재처리하지 않아 출력 결과가 이동이 없었던 실행과 일치합니다.

이 디자인이 열어주는 공간에 흥분됩니다. 생성된 커널은 모든 개선 사항이 모든 기기에 도달한다는 것을 의미하며, 이 엔진은 한 번도 본 적 없는 하드웨어에 대해서도 융합하고 조정할 수 있습니다. 페이지화와 재배치는 배치를 스케줄링으로 바꿉니다: 메모리보다 큰 모델, 서로 다른 GPU, 대화가 진행됨에 따라 변화하는 작업들입니다.

코드 -> https://github.com/jitllm/jitllm
제출자: /u/web3samy [link] [comments]