Disaggregated Prefill / Decode¶

개요¶

Prefill과 decode는 NPU에 요구하는 자원의 성격이 다릅니다. Prefill은 프롬프트 전체를 한 번에 처리하는 compute-bound 작업이고, decode는 스텝마다 토큰 하나를 만들며 KV 캐시 읽기가 병목입니다. 두 단계를 같은 인스턴스에서 실행하면 긴 prefill이 그 뒤에 대기 중인 decode 스텝을 지연시킵니다.

Disaggregated prefill / decode(PD disaggregation)는 두 단계를 별도의 vLLM 인스턴스로 분리합니다. Prefill 인스턴스가 프롬프트의 KV 캐시를 계산하면, decode 인스턴스가 그 KV 캐시를 네트워크로 가져와 토큰을 생성합니다. 이렇게 하면 양쪽의 규모와 설정을 독립적으로 조정할 수 있습니다.

RBLN vLLM은 이를 RblnNixlConnector로 구현합니다. nixl-rbln을 사용해 KV 캐시를 RDMA로 전송하는 KV connector입니다.

NIXL(NVIDIA Inference Xfer Library)은 vLLM의 KV connector들이 기반으로 삼는 전송 라이브러리입니다. 엔진이 옮기려는 메모리를 등록하는 방법과 옮기는 방법을 하나씩 제공하고, transport마다 백엔드 플러그인을 통해 연결됩니다. RBLN NPU 메모리는 nixl-rbln을 거쳐야 NIXL에 등록할 수 있습니다. nixl-rbln은 아래 절에서 다룹니다.

아키텍처¶

RblnNixlConnector는 vLLM의 pull 방식 NIXL connector를 확장한 것이므로, 전송을 주도하는 쪽은 decode입니다. Decode가 prefill 인스턴스의 블록을 읽어 자신의 KV 캐시로 가져옵니다. 반대 방향의 connector도 사용할 수 있습니다.

요청은 두 인스턴스를 다음 순서로 통과합니다.

1. 두 인스턴스 앞의 프록시가 요청을 먼저 prefill 인스턴스로, 이어서 decode 인스턴스로 전달하면서 KV 전송 파라미터를 함께 전달합니다.
2. Prefill 인스턴스가 프롬프트를 실행하고, KV 캐시를 담은 블록들을 등록합니다.
3. Decode 인스턴스가 prefill 인스턴스와 handshake한 뒤 해당 블록들을 자신의 KV 캐시로 직접 pull하고 decode를 시작합니다.

Push 방향에서는 전송 자체가 반대가 됩니다. Decode 인스턴스가 KV 캐시를 받을 블록을 등록하고, prefill 인스턴스가 프롬프트 처리를 마친 뒤 그 블록에 직접 씁니다.

두 인스턴스는 모델, dtype, attention 백엔드가 일치해야 합니다. Connector는 handshake 과정에서 compatibility hash로 세 가지가 같은지 확인하고, 다르면 진행을 거부합니다. KV 캐시 레이아웃도 일치해야 하지만, 레이아웃은 hash에 들어가지 않으므로 handshake가 따로 검사합니다.

Connector 자신의 버전, 전송 방향, 그리고 speculative decoding을 쓰는 경우 draft 모델도 hash에 함께 들어갑니다. 따라서 pull 인스턴스와 push 인스턴스는 짝이 되지 못합니다.

사전 준비 사항¶

• vllm model implementation: --model-impl vllm.
• 디바이스 텐서. 두 전송 방식 모두 디바이스 텐서를 써야 합니다. vllm model implementation이 KV 캐시를 기본적으로 디바이스 텐서로 두므로 끄지만 않으면 됩니다.
• 인스턴스를 실행하는 모든 호스트에 서로 연결된 RDMA(RoCE) NIC. peer와 연결되는 NIC 하나가 rail 하나를 담당하므로, NIC이 여럿인 호스트는 rail도 여럿이고, 두 인스턴스는 양쪽이 모두 가진 rail로 전송할 수 있습니다.
• Device-to-device 방식을 사용하려면 nixl-rbln.

Connector는 다음을 모두 만족하는 RoCE NIC을 고릅니다. NPU와 같은 NUMA 노드에 있어야 하고, 포트가 ACTIVE여야 하고, IPv4 주소가 있어야 하며, 호스트의 기본 경로가 걸린 인터페이스가 아니어야 합니다. 호스트에 그런 NIC이 있는지는 아래 명령으로 확인합니다.

1
2
3
4
5

rbln-smi -j                                      # 각 NPU의 `pci.numa_node`
rdma link show                                   # 포트 상태, 그리고 각 RDMA 장치가 쓰는 netdev
ip route show default                            # 제외되는 관리 인터페이스
ip -4 -brief addr show                           # IPv4 주소를 가진 netdev
grep . /sys/class/infiniband/*/device/numa_node  # 각 RDMA 장치가 속한 NUMA 노드

rbln-smi는 RBLN 드라이버에 포함되어 있고 나머지는 iproute2와 sysfs에 들어 있으므로, RDMA를 실행할 수 있는 호스트에는 이미 설치되어 있습니다. 엔진은 이런 NIC이 있는 NUMA 노드의 NPU에 배치해야 합니다.

설정¶

두 인스턴스 모두 --kv-transfer-config와 함께 실행합니다. Connector는 여기에서 다음 키를 읽습니다.

kv_connector

Pull 방향은 RblnNixlPullConnector, 실험적인 push 방향은 RblnNixlPushConnector로 설정합니다. RblnNixlConnector는 pull connector의 별칭으로 남아 있습니다. 기존 배포가 쓰고 있는 이름이기 때문입니다.

kv_role

kv_producer, kv_consumer, kv_both 중 하나입니다. Connector를 설정하면 이 값을 반드시 지정해야 하고, kv_both는 두 역할을 다 합니다. 아래 예제에서는 두 인스턴스 모두 kv_both를 씁니다.

kv_buffer_device

Device-to-device는 "rbln", host-bounce는 "cpu"입니다. 그 외의 값은 거부됩니다.

engine_id

인스턴스를 식별합니다. Decode 쪽은 전달받은 식별자를 자신이 기대한 값과 대조합니다.

kv_connector_extra_config

Connector 자신의 설정입니다. 아래에 정리했습니다.

kv_connector_extra_config는 upstream NIXL connector와 공유하는 설정이고, upstream도 자신의 키를 여기에서 읽습니다. 이 connector가 읽는 키는 다음과 같습니다.

swa_window_mode

기본값: false. Sliding-window 그룹의 블록을 window가 움직이는 단위로 자르고, window가 걸친 곳만 전송합니다. Device-to-device에서만 쓸 수 있습니다.

chunk_mode

기본값: false. 블록을 chunk로 자릅니다. 이 경우 요청의 마지막 블록에서 토큰이 든 chunk만 나갑니다. 마지막 토큰이 걸친 chunk는 여전히 통째로 나갑니다. Device-to-device에서만 쓸 수 있습니다.

chunk_tokens

기본값: 설정 없음. 이때는 prefill 스텝 하나(--max-num-batched-tokens)를 씁니다. 이 값은 chunk 크기의 하한입니다. 실제 chunk 크기는 다음을 만족하는 가장 작은 값입니다. 그 하한 이상이어야 하고, descriptor 하나가 덮는 토큰 범위를 나누어떨어지게 해야 합니다. 스텝 하나가 닫는 chunk가 최대 하나이므로, prefill 스텝 하나 이상이어야 합니다.

push_stream

기본값: false. Prefill이 블록을 닫는 대로 하나씩 전송합니다. 이 설정이 꺼져 있으면 요청의 KV를 프롬프트가 끝날 때까지 보관합니다. Push 방향과 device-to-device 전용입니다.

stripe_width

기본값: 설정 없음. 이때 RBLN NIXL 백엔드는 1로 간주하여 전송 하나가 rail 하나를 씁니다. 0을 주면 두 인스턴스가 공유하는 모든 rail에 연결이 열립니다. 그 외의 값은 거부됩니다.

link_down_exit_s

기본값: 0. 이때는 종료하지 않습니다. 로컬 링크를 전부 잃은 인스턴스는 그 시간이 지나면 종료하므로, 배포 쪽에서 새 인스턴스를 띄울 수 있습니다. kv_role: kv_producer일 때만 설정할 수 있습니다. Decode 인스턴스는 처리 중인 요청을 유지한 채 바깥에서 교체되기 때문입니다.

enforce_handshake_compat은 upstream connector의 키이고 기본값은 true입니다. false로 설정하면 handshake가 수행하는 compatibility hash 검사를 건너뜁니다.

chunk_mode와 swa_window_mode는 둘 다 fabric을 건너는 양을 줄이며, 서로 독립입니다. 자르는 자리가 서로 다릅니다. chunk_mode를 켜면 요청의 마지막 블록에서 토큰이 없는 부분은 fabric으로 나가지 않으므로, 처리하는 프롬프트에 비해 block_size가 클수록 이득이 커집니다. swa_window_mode는 sliding-window attention을 쓰는 모델에만 씁니다. 그런 모델의 window 그룹은 프롬프트 길이와 무관하게 블록 하나를 통째로 차지하는데 그 블록에서 유효한 토큰은 window 길이만큼이므로, window에 비해 block_size가 클수록 이득이 커집니다. Push 방향에서 push_stream은 전송의 tail을 줄입니다. 프롬프트가 길수록, 프롬프트가 끝난 시점에는 그 요청의 KV가 그만큼 더 많이 건너가 있습니다.

예제¶

한 호스트에 pull 방향 인스턴스 둘을 띄웁니다. 첫 번째 NPU에 prefill 인스턴스, 두 번째 NPU에 decode 인스턴스, 그리고 그 앞에 프록시를 둡니다.

Prefill 인스턴스입니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

RBLN_VISIBLE_DEVICES=0 \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve Qwen/Qwen3-0.6B \
  --model-impl vllm \
  --port 8100 \
  --block-size 1024 \
  --max-model-len 4096 \
  --max-num-batched-tokens 512 \
  --kv-transfer-config \
  '{"kv_connector":"RblnNixlPullConnector","kv_role":"kv_both","kv_buffer_device":"rbln"}'

Decode 인스턴스는 같은 설정을 쓰고, NPU와 포트만 다릅니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

RBLN_VISIBLE_DEVICES=1 \
VLLM_NIXL_SIDE_CHANNEL_PORT=5659 \
vllm serve Qwen/Qwen3-0.6B \
  --model-impl vllm \
  --port 8200 \
  --block-size 1024 \
  --max-model-len 4096 \
  --max-num-batched-tokens 512 \
  --kv-transfer-config \
  '{"kv_connector":"RblnNixlPullConnector","kv_role":"kv_both","kv_buffer_device":"rbln"}'

프록시는 upstream vLLM의 toy_proxy_server.py입니다.

1
2
3
4

python3 toy_proxy_server.py \
  --port 8192 \
  --prefiller-port 8100 \
  --decoder-port 8200

요청은 프록시로 보냅니다. 두 인스턴스로 직접 보내지 않습니다.

1
2
3

curl http://localhost:8192/v1/completions \
  -H 'Content-Type: application/json' \
  -d '{"model": "Qwen/Qwen3-0.6B", "prompt": "The capital of France is", "max_tokens": 16}'

toy_proxy_server.py는 upstream의 예제입니다. 라우팅 정책도, 헬스 체크도, 확장도 없습니다. 운영 환경에서는 두 인스턴스 앞에 라우터를 둡니다.

위의 두 인스턴스는 한 호스트에 있으므로 side-channel 포트를 서로 다르게 지정했습니다. 호스트를 둘로 나누면 포트는 더 이상 겹치지 않고, 바뀌는 것은 VLLM_NIXL_SIDE_CHANNEL_HOST입니다. 각 인스턴스는 이 변수에 지정된 주소를 알리고, peer는 그 주소로 접속합니다. 기본값이 localhost이므로 설정하지 않으면 peer가 받는 주소는 peer 자신의 주소입니다. 이 경우 handshake가 성립하지 않습니다. 이 side channel은 ZMQ over TCP이며 handshake만 실어 나릅니다. KV 캐시 자체는 RoCE로 갑니다.

검증된 구성은 레시피에 있습니다. Prefill과 decode를 서로 다른 노드에서 서비스하고, 그 앞에 llm-d 라우터를 둡니다.

동작 확인¶

로그를 보면 이 인스턴스가 어느 경로로 올라왔는지, peer와 연결되었는지 알 수 있습니다. 다음 접두사로 로그를 검색할 수 있습니다.

RBLN NIXL (D2D): registered

이 인스턴스의 KV 캐시가 device-to-device 방식으로 NIXL에 등록되었습니다. 등록은 워밍업 뒤에 이루어지며, 등록을 미룬다는 RBLN NIXL (D2D): deferring registration이 그보다 먼저 나옵니다.

RBLN NIXL: nixl-rbln not available

이 인스턴스는 upstream NIXL의 UCX transport로 동작하고 있습니다. RBLN 백엔드를 쓰고 있지 않습니다.

remote descriptor(s) for engine

peer와의 handshake가 끝났고, 이 인스턴스가 전송에 쓸 descriptor를 만들었습니다. 이 줄은 그 peer가 처음 필요해지는 요청에서, 그리고 전송을 거는 쪽에서 나옵니다. Pull 방향에서는 decode 인스턴스, push 방향에서는 prefill 인스턴스입니다. KV 캐시 전송은 handshake가 완료된 뒤 시작됩니다.

문제 해결¶

no usable data-plane RoCE NIC on numa=N 오류로 시작하지 못합니다

그 NUMA 노드에 사전 준비 사항의 조건을 만족하는 NIC이 없다는 뜻입니다. 어느 노드를 확인했는지는 메시지에 나옵니다. 그 노드에서 사전 준비 사항의 확인 명령을 실행해 보면 됩니다.

두 인스턴스가 등록은 하는데 handshake까지 가지 못합니다

각자 등록 줄은 찍는데 descriptor 줄이 뒤따르지 않습니다. 대개 side-channel 주소가 원인입니다. 호스트가 둘일 때 이 주소를 지정하는 방법은 예제에 나와 있습니다.

전송 단위¶

Descriptor의 길이는 메모리를 등록하는 시점에 고정됩니다. 전송할 때는 어떤 descriptor를 보낼지만 고릅니다. Descriptor 하나의 일부만 보낼 수는 없습니다. 따라서 부분만 채워진 마지막 블록도 통째로 전송됩니다. chunk_mode를 켜면 달라집니다. 마지막 블록을 chunk로 잘라 토큰이 없는 chunk는 빼고 보냅니다. Device-to-device에서는 블록 하나가 shard마다 하나씩 여러 descriptor에 걸칩니다.

push_stream은 descriptor가 언제 나가는지를 바꿉니다. 보내는 집합은 같습니다. 블록은 prefill이 닫는 대로 나가고, 남은 것은 요청이 끝날 때 나갑니다.

따라서 block_size는 PD disaggregation만의 설정이 아닙니다. 컴파일러의 KV 캐시 partition 길이를 정하고, sub-block prefix cache 단위를 제약하며, sliding window를 쓰는 모델에서는 그 window 길이의 배수여야 합니다. 전송까지 더해 네 가지를 함께 보고 바꿔야 합니다.

제약 사항¶

• Prefill 인스턴스와 decode 인스턴스는 같은 block_size를 써야 합니다. 값이 다르면 거부됩니다.
• Sub-block prefix caching과 KV connector는 둘 다 full-block prefix cache hit을 늘리는 쪽으로 동작하므로, 스케줄러는 요청마다 둘 중 하나만 씁니다. 스케줄러는 더 많은 토큰을 확보하는 쪽을 쓰고, 같으면 sub-block이 우선합니다. 자세한 내용은 Prefix Caching에서 확인할 수 있습니다.