[vllm] vLLM의 NIXL KV 전송을 활용한 GDN(Gated Delta Net) 모델 지원 최적화

PR 링크: vllm-project/vllm#41869 상태: Merged | 변경: +0 / -0

들어가며

최근 vLLM은 Mamba와 같은 SSM(State Space Model) 계열 모델을 효율적으로 서빙하기 위해 NIXL 커넥터를 통한 KV 전송 기능을 확장하고 있습니다. 이번 PR은 Qwen3.5와 같은 GDN(Gated Delta Net) 모델을 지원하기 위해 컨볼루션 상태 레이아웃(conv-state layout)을 개선하고, 이기종 텐서 병렬(Heterogeneous TP) 환경에서 발생하는 블록 매핑 문제를 해결하는 데 초점을 맞추고 있습니다.

코드 분석

1. vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py

기존 Mamba2 모델은 컨볼루션 상태가 [x, B, C] 구조였으나, GDN 모델은 [Q, K, V] 구조를 가집니다. 이를 일반화하기 위해 MambaConvSplitInfo 클래스를 수정하여 하드코딩된 필드 대신 local_proj_dims 튜플을 사용하도록 변경했습니다.

Before:

@property
def x_bytes(self) -> int:
    return self.x_local * self.conv_rows * self.conv_dtype_size

After:

@property
def proj_bytes(self) -> tuple[int, int, int]:
    row_bytes = self.conv_rows * self.conv_dtype_size
    return tuple(d * row_bytes for d in self.local_proj_dims)

이 변경으로 인해 다양한 SSM 아키텍처의 프로젝션 차원을 유연하게 처리할 수 있게 되었습니다.

2. vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

원격 디스크립터 영역을 빌드하는 로직에서 가독성을 높이고, 이기종 TP 환경에서의 오프셋 계산 로직을 간소화했습니다.

Before:

# Build 4 remote desc regions (x, B, C, ssm)

After:

# Build 4 remote desc regions (proj0, proj1, proj2, ssm)
conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio)

_conv_decomp.remote_conv_offsets를 활용하여 P와 D 사이의 TP 비율에 따른 복잡한 분기 처리를 제거하고, 모델 아키텍처에 독립적인 오프셋 계산이 가능해졌습니다.

왜 이게 좋은가

1. 범용성 확보: 특정 모델(Mamba2)에 종속적이었던 로직을 GDN(Qwen3.5)까지 확장함으로써, 향후 등장할 다양한 SSM 기반 모델을 별도의 커널 수정 없이 지원할 수 있는 구조를 갖추었습니다.
2. 정확성: 이기종 TP 환경(예: 1P2D, 2P4D)에서 physical_blocks_per_logical 불일치를 해결하여, GSM8K 벤치마크 기준 독립 실행 모델과 유사한 정확도(0.323 ± 0.03)를 유지함을 확인했습니다.
3. 유지보수성: local_proj_dims를 통한 추상화는 코드의 가독성을 높이고, 리뷰어들의 피드백대로 'x, B, C'와 같은 특정 아키텍처 용어 대신 'proj0, proj1, proj2'와 같은 일반적인 명칭을 사용하여 향후 아키텍처 변경 시 코드 수정 범위를 최소화했습니다.

리뷰 피드백 반영

리뷰 과정에서 async scheduling이 TP>1 환경에서 문제를 일으킬 수 있다는 점이 지적되었고, 이를 반영하여 테스트 스크립트에서 특정 모델 설정 시 --no-async-scheduling을 명시하도록 수정되었습니다. 또한, 코드의 명명 규칙을 더 일반화하여 특정 모델에 종속되지 않는 API 설계를 지향했습니다.

참고 자료

• https://github.com/vllm-project/vllm

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.