vLLM v0.20.0 하이브리드 SSM 모델을 위한 분산 서빙 추가

vLLM v0.20.0 하이브리드 SSM 모델을 위한 분산 서빙 추가

소개

vLLM은 기존 표준 트랜스포머 워크플로를 변경하지 않고 하이브리드 SSM‑FA 모델을 위한 분산 프리필/디코드 지원을 위해 NIXL 기반 KV 커넥터를 확장했습니다.

하이브리드 아키텍처인 NVIDIA Nemotron‑H는 Mamba 스타일 SSM 레이어와 전체 어텐션 레이어를 교차 배치하여 선형 시간 SSM 효율성과 어텐션 표현력을 결합합니다. 기존 NIXL 분산 P/D 설계는 일관된 KV 캐시 형식을 가정했지만, FA와 SSM 레이어가 서로 다른 레이아웃과 크기로 상태를 저장하기 때문에 하이브리드 모델에는 적용되지 않습니다.

배경: NIXL KV 전송 워크플로

표준 트랜스포머 모델의 경우, NIXL 분산 P/D는 메모리 영역을 등록하고, 블록별 디스크립터를 생성하며, 핸드셰이크를 수행하고, RDMA를 통해 블록을 전송함으로써 동작합니다. 각 워커는 자신의 KV 캐시 텐서를 NIXL에 등록한 뒤, 각 블록에 대해 (address, length, device_id)를 지정하는 디스크립터를 생성합니다. 핸드셰이크는 프리필‑디코드 쌍당 한 번 메타데이터를 교환합니다. 스케줄러는 디코드 워커에게 어떤 블록 ID를 가져올지 알려주고, 디코드 워커는 블록 ID를 디스크립터 ID로 매핑한 뒤 RDMA READ를 발행하고 완료를 폴링합니다. M개의 등록된 영역과 각 영역당 N개의 블록이 있을 때, 디스크립터 인덱스는 r * N + b 로 계산됩니다.

도전 과제: FA와 SSM 상태는 근본적으로 다름

하이브리드 모델은 FA 레이어가 토큰당 K/V 쌍을 저장하고, SSM 레이어가 고정 크기의 conv 상태와 SSM 상태를 저장하기 때문에 일관된 디스크립터 가정이 깨집니다.

FA 레이어의 KV 캐시 형태는 [num_blocks, 2, block_size, num_kv_heads, head_dim](또는 변형)이며, SSM 레이어는 conv 상태 (conv_dim, state_len)와 SSM 상태 (num_heads, head_dim, state_size)를 유지합니다. 이러한 상태는 토큰 기반이 아니므로 블록 크기의 개념이 다릅니다: 각 SSM 블록은 전체 상태 스냅샷을 의미합니다. HMA는 FA와 SSM 그룹 간 메모리를 풀링하고, SSM 행을 패딩하여 두 뷰가 동일한 물리 텐서와 바이트 단위 페이지 크기를 공유하도록 합니다. 따라서 일관된 (address, length)를 가진 단일 디스크립터 리스트로는 FA와 SSM 뷰를 동시에 올바르게 인덱싱할 수 없습니다.

이중 디스크립터 뷰

해결책은 동일한 물리 메모리 위에 두 개의 별도 디스크립터 리스트를 등록하는 것입니다: 하나는 FA 디스크립터용, 다른 하나는 SSM 디스크립터용이며, 단일 NIXL 전송 핸들 아래에서 연결됩니다. FA 디스크립터는 첫 num_descs = M * N_phys 슬롯을 차지하며, 각 영역에 대해 K와 V를 별도로 인덱싱합니다. SSM 디스크립터는 그 뒤를 이어, 각 Mamba 레이어가 블록당 네 개의 서브‑디스크립터 (x, B, C, SSM) 로 표현됩니다. 블록 ID → 디스크립터 ID 매핑은 FA 그룹의 경우 region * N_phys + block_id, SSM 그룹의 경우 mamba_region_id * N_log + block_id + num_descs 로 수행됩니다.

물리 블록 크기 vs. 논리 블록 크기

FA 레이어는 어텐션 커널을 위해 별도의 물리 블록 크기가 필요할 수 있지만, SSM 레이어는 논리 블록을 그대로 사용하므로 두 디스크립터 섹션의 블록 수가 달라집니다. 비율 logical_block_size / kernel_block_size 를 사용해 물리 블록 수를 physical_blocks = logical_blocks * ratio 로 계산하는데, 이는 FA 레이어에만 적용됩니다. SSM 레이어는 항상 논리 블록 수를 사용합니다. 이 정보는 _physical_blocks_per_logical 에서 추적되며, 프리필과 디코드 인스턴스가 텐서‑패럴렐 크기가 다를 때 서로 다를 수 있습니다. 블록‑ID‑디스크립터‑ID 매핑은 그룹이 FA인지 SSM인지에 따라 적절한 스트라이드를 사용합니다.

3‑디스크립터 Conv 전송

이기종 텐서‑패럴렐 설정에서는 Conv 상태가 DS 형식으로 배치되어 각 디코드 랭크가 연속적인 x, B, C 슬라이스를 세 개의 디스크립터 영역을 통해 읽을 수 있게 합니다. 이를 통해 추가 버퍼나 재배열 없이 제로‑카피 RDMA가 가능합니다. DS 레이아웃 (dim, state_len) 은 각 서브‑프로젝션을 메모리 상에서 연속적으로 배치합니다: [x][B][C][SSM]. 디코드 워커는 단일 NIXL READ 작업 내에서 세 개의 연속적인 읽기를 발행해 자신의 Conv 상태 조각을 획득합니다. 이렇게 하면 디코드 측에 스테이징 버퍼를 할당할 필요가 없고, 전송 후 재배열을 없애며, 소유한 1/TP 비율의 Conv 상태만 전송하고, HMA 패딩 바이트도 건너뛸 수 있습니다. 이 접근법은 동질 및 이기종 TP 모두에 적용되며, 후자의 경우 Mamba 레이어당 네 개의 디스크립터 영역 (x, B, C, SSM) 을 사용합니다.

전체 흐름: Nemotron‑H 예시

구체적인 예시를 통해 nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 모델을 TP=2 환경에서 분산 P/D 로 서빙하는 방법을 보여줍니다. 모델은 Mamba와 FA 레이어가 교차되는 52개의 레이어를 가지고, HMA에 의해 5개의 그룹(4 Mamba, 1 FA)으로 묶여 6개의 공유 KV 캐시 텐서를 형성합니다. FA 레이어는 [num_blocks, 2, block_size=400, 4, 128] 형태를 사용하고, SSM 레이어는 [num_blocks, 3, 3072](conv)와 [num_blocks, 48, 64, 128](ssm) 형태를 사용합니다. HMA 패딩 후 두 뷰는 동일한 페이지 크기를 공유합니다. 프리필 인스턴스는 6개의 텐서를 등록하고, 모든 영역 × N_phys 블록에 대해 FA 디스크립터(K와 V를 별도) 를 생성한 뒤, 블록당 네 개의 서브‑리전으로 구성된 Mamba 디스크립터를 추가합니다. 디코드 인스턴스는 적절한 스트라이드를 사용해 블록 ID를 디스크립터 ID로 매핑하고, FA와 Mamba 디스크립터를 모두 포함한 단일 READ 를 사전 준비하여 발행하고 완료를 폴링합니다. 중간 버퍼나 데이터 재배열이 전혀 필요하지 않습니다.

성능

하이브리드 SSM 모델에 대한 분산 P/D 는 프리필 간섭을 격리함으로써 높은 동시성에서도 공동 배치 서빙 처리량과 동등하거나 이를 초과합니다. NVLink가 연결된 8× H200 GPU 환경에서 nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 모델을 대상으로, 공동 배치 기준(단일 인스턴스, TP=8)과 분산 구성(프리필 TP=4 + 디코드 TP=4, GPU 수 동일)을 비교했습니다. 동시 사용자 수를 8명에서 256명까지 늘리면서, 분산 파레토 곡선은 배치 크기가 커질수록 공동 배치 곡선을 압도했으며, 디코드가 프리필과 분리될 때 GPU당 초당 출력 토큰 수가 더 높아졌습니다.

시작하기

하이브리드 SSM 모델을 분산 P/D 로 실행하려면 VLLM_SSM_CONV_STATE_LAYOUT=DS 를 설정하고, 프리필 및 디코드 인스턴스를 적절한 인수와 함께 시작하십시오. 프리필 인스턴스 예시 명령:

VLLM_SSM_CONV_STATE_LAYOUT=DS vllm serve nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.85 \
    --trust-remote-code \
    --max-model-len 8192 \
    --block-size 128 \
    --no-disable-hybrid-kv-cache-manager \
    --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

디코드 인스턴스는 kv_role 을 적절히 설정한 유사 명령을 사용합니다(소스에 표시되지 않음). DS 레이아웃은 이기종 TP 환경에서 필요하지만, 그렇지 않은 경우 선택 사항입니다.

제한 사항 및 향후 작업

현재 지원은 Mamba2에만 한정되며, Mamba1 모델과 GDN 레이어는 아직 지원되지 않습니다. 추론 디코딩과의 상호 작용은 충분히 검증되지 않았으며, HMA가 활성화된 경우 혼합 블록 크기는 아직 지원되지 않습니다. Mamba1의 SSM 시간 형태는 Conv 분해에 필요한 intermediate_size 재구성을 방해합니다. GDN 지원은 분산 로드맵에 명시되어 있습니다. 추론 디코딩 및 HMA와 함께하는 블록‑크기 비율 처리 역시 향후 과제로 남아 있습니다.

감사의 글

Thomas Parnell (IBM Research)와 Roi Koren (NVIDIA)에게 기여에 대한 감사를 전합니다.

Sources