DiffusionGemma: vLLM에서 네이티브로 처음 지원되는 첫 번째 Diffusion LLM
DiffusionGemma: vLLM에서 네이티브로 처음 지원되는 첫 번째 Diffusion LLM
vLLM은 Gemma4 백본을 기반으로 한 26B-parameter 이산 디퓨전 언어 모델인 DiffusionGemma를 통합했습니다. 이 통합은 vLLM에서 디퓨전 LLM(dLLM)에 대한 첫 번째 네이티브 지원을 의미하며, 순차적 토큰 생성에서 반복적 블록 정제로의 전환을 가능하게 하여 낮은 배치 크기에서 지연 시간을 크게 줄입니다.
DiffusionGemma 아키텍처 및 디코딩
DiffusionGemma는 왼쪽에서 오른쪽으로 토큰을 하나씩 생성하는 표준 자동 회귀 디코딩을 대체하여, 256 토큰 길이의 고정 캔버스를 반복적으로 디노이즈하는 과정을 사용합니다. 이 접근 방식은 메모리 대역폭 압력을 추가 컴퓨트와 trade-off 하여, 여러 디노이징 단계에서 모델이 여러 토큰을 병렬로 정제할 수 있도록 합니다.
듀얼 모드 운영
DiffusionGemma는 두 가지 distinto 모드에서 작동하는 단일 가중치 세트를 사용합니다:
- Encoder Mode: 인과적 어텐션을 사용하여 프롬프트를 프리필하고 완료된 블록을 KV 캐시에 "커밋"합니다.
- Decoder Mode: 양방향 어텐션을 사용하여 캔버스의 모든 위치가 다른 모든 위치에 어텐션하여 동시에 정제할 수 있도록 합니다.
인코더 모드가 표준 인과적 어텐션을 사용하기 때문에, vLLM의 자동 프리픽스 캐싱은 수정 없이도 호환되며 작동합니다.
샘플링 루프와 엔트로피 경계 디노이징
생성은 256 토큰 블록 단위로 이루어집니다. 프롬프트가 프리필된 후, 캔버스는 랜덤 토큰으로 초기화됩니다. 모델은 각 위치에서 후보 토큰을 샘플링하는 디노이징 루프에 진입합니다.
어떤 토큰을 유지할지 결정하기 위해, DiffusionGemma는 엔트로피 경계 규칙을 사용합니다: 누적 엔트로피가 고정된 예산을 초과할 때까지 가장 확신도가 높은 토큰부터 낮은 토큰 순으로 받아들입니다. 캔버스가 수렴 상태가 되면—즉, argmax 예측이 여러 단계 동안 안정적이고 토큰당 평균 엔트로피가 임계값 이하로 떨어지거나 단계 제한에 도달하면—토큰은 인코더 패스를 통해 커밋되고, 다음 블록에 대해 프로세스가 반복됩니다.
안정성을 위한 셀프 컨디셔닝
수렴 속도를 높이고 안정성을 높이기 위해 모델은 셀프 컨디셔닝을 사용합니다. 디노이징 단계 사이에서 모델은 이전 예측의 전체 소프트맥스 분포에 조건을 부여합니다. 이 분포는 토큰 임베딩의 확률 가중 평균으로 변환되어 게이트된 MLP를 통해 캔버스 임베딩에 추가되며,これにより 모델은 이전 믿음에 대한 메모리를 제공합니다.
vLLM 구현 세부 사항
DiffusionGemma 통합을 위해 비자동 회귀 디코딩 패턴을 지원해야 했습니다. vLLM은 여러 아키텍처 추상화를 활용하여 이를 달성했습니다:
추측 디코딩 데이터 경로
vLLM은 기존의 추측 디코딩 경로를 재사용하여, 현재 캔버스를 완전히 거부되거나 완전히 받아들여지는 드래프트 토큰 세트로 취급했습니다.これにより 팀은 핵심 스케줄러와 모델 러너를 최소한의 변경으로 유지할 수 있었습니다.
ModelState 인터페이스
모델 러너를 포크하지 않기 위해 vLLM은 ModelState 추상화를 활용했습니다. 이는 일반 러너를 변경하지 않고 DiffusionGemma가 사용자 정의 동작을 정의할 수 있는 훅을 제공합니다:
| 훅 | DiffusionGemma에서의 목적 |
|---|---|
prepare_inputs() |
캔버스 임베딩과 셀프 컨디셔닝 적용을 처리합니다. |
prepare_attn() |
요청당 인코더(인과적)와 디노이즈(양방향) 어텐션 간을 전환합니다. |
custom_sampler() |
기본 샘플러를 대체하기 위해 DiffusionSampler를 설치합니다. |
add_request() / remove_request() |
요청별 디퓨전 상태(캔버스, 확률)의 생명 주기를 관리합니다. |
동적 퍼-시퀀스 인과적 어텐션
단일 배치에 프리필, 디노이즈, 커밋 단계의 요청이 혼합될 수 있기 때문에, vLLM은 동적 퍼-시퀀스 인과적 어텐션을 구현했습니다.これにより 어텐션 마스크가 각 요청의 특정 인과성 요구에 적응할 수 있습니다. 이 기능은 Triton Attention (TRITON_ATTN) 및 FlashAttention 4 (FLASH_ATTN) 백엔드 모두에서 지원됩니다.
대칭 슬라이딩 윈도우 어텐션
슬라이딩 윈도우 어텐션을 사용하는 레이어에 대해, DiffusionGemma는 디노이징 단계에서 대칭 윈도우를 요구합니다. 인과적 요청은 한쪽 방향 윈도우(앞쪽의 $W$ 토큰에 어텐션)를 사용하는 반면, 양방향 요청은 양쪽 각각 $W$ 토큰에 어텐션하여 총 윈도우 크기가 $2W + 1$가 됩니다.
성능 및 양자화
처리량 결과
DiffusionGemma의 아키텍처는 extrêmement low-latency 추론을 가능하게 합니다. 단일 GPU에서 배치 크기 1로 수행된 벤치마크는 자동 회귀 기준선에 비해 상당한 처리량 향상을 보여줍니다:
- H200 (FP8): 초당 1,288 생성 토큰 (~6$ times$ 표준 자동 회귀 기준선, ~3$ times$ 멀티-토큰 예측).
- H100 (FP8): 초당 1,008 생성 토큰 (~5$ times$ 표준 자동 회귀 기준선, ~2.6$ times$ 멀티-토큰 예측).
양자화된 체크포인트
LLM Compressor를 사용하여 모델의 양자화 버전이 생성되었으며, compressed-tensors 형식으로 제공됩니다. 지원되는 형식은 다음과 같습니다:
- FP8: 완전 동적 활성화를 갖는 양자화된 가중치.
- NVFP4: 가중치와 활성화가 모두 NVFP4 형식으로 양자화됨.