vLLM에서 Reflection Beam 1.0: Day 0 서빙, 에이전트 최적화 커널, DFlash2 추측 디코딩

vLLM은 Reflection의 Beam(501B 파라미터, 토큰당 23B 활성 파라미터를 가진 mixture-of-experts 모델)에 대한 day 0 서빙 지원을 추가하고, NVIDIA Blackwell 하드웨어에 특화하여 최적화했습니다. 이 결과는 완전한 프로덕션 서빙 기능(추론, 도구 호출, 구조화된 출력, 텐서/데이터/전문가 병렬 처리)과 코딩 에이전트가 자율적으로 최적화한 커널 스택을 결합하여, 배치-1 디코딩을 vLLM의 첫 번째 작동 구현보다 2.71배, 프리필을 1.87배 더 빠르게 만들었기 때문에 중요합니다.

단일 요청은 TP8에서 Beam의 DFlash2 드래프트 헤드를 사용하여 초당 708토큰에 도달하며, 수용 길이는 3.7입니다. 적응형 검증은 높은 동시성에서 처리량을 10% 더 추가합니다.

모델은 Docker 이미지 vllm/vllm-openai:reflection-beam으로 서빙할 수 있습니다:

docker run --gpus all --ipc=host -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  vllm/vllm-openai:reflection-beam \
  reflection/Beam-1.0-501B-A23B-NVFP4 \
  -tp 4 \
  --max-num-seqs 128 \
  --reasoning-parser beam \
  --tool-call-parser beam --enable-auto-tool-choice \
  --kv-cache-dtype fp8 \
  --attention-config '{"backend":"FLASHINFER"}' \
  --kernel-config '{"enable_flashinfer_autotune":false}' \
  --speculative-config '{
    "method": "dflash",
    "model": "reflection/Beam-1.0-501B-A23B-DFlash2",
    "num_speculative_tokens": 7,
    "attention_backend": "FLASH_ATTN"
  }'

Beam 모델 아키텍처: 52개 레이어, 하이브리드 어텐션, DeepSeek-V3 스타일 MoE

Beam은 hidden size 5120을 가진 52개 레이어 MoE 디코더입니다. MoE 라우팅은 DeepSeek-V3를 따릅니다. 아키텍처의 독특한 선택은 어텐션 레이어와 잔차 스트림 정규화 방식에 있습니다.

어텐션: 슬라이딩 윈도우 레이어 + NoPE 전체 어텐션 레이어

Beam은 어텐션을 두 가지 방식으로 나눕니다. 43개 레이어는 RoPE를 사용하여 마지막 512개 토큰에 대해 슬라이딩 윈도우 어텐션을 사용합니다. 나머지 9개 레이어(여섯 번째 레이어마다, 그리고 마지막 레이어)는 위치 인코딩 없이(NoPE) 전체 컨텍스트에 대해 어텐션을 수행합니다.

RoPE가 없을 때 키 수가 증가함에 따라 어텐션이 평평해지기 때문에, Beam은 NoPE 레이어에 SSMax를 적용합니다. SSMax는 각 쿼리를 (1 + p/2048)^0.1로 스케일링합니다. 여기서 p는 1부터 시작하는 위치입니다. 따라서 컨텍스트가 길어질수록 소프트맥스가 더 선명해집니다.

Beam은 48개의 쿼리 헤드, 8개의 KV 헤드, 헤드 차원 128을 가진 grouped-query attention을 사용합니다. Q와 K는 RoPE 또는 SSMax 전에 헤드별로 RMS 정규화됩니다. 출력 프로젝션(o_proj) 전에 각 헤드의 출력은 sigma(gate + bias)로 곱해집니다. 여기서 gate는 자체 전체 너비 프로젝션에서 오며 학습된 헤드별 바이어스가 있습니다. 따라서 헤드는 토큰별로 줄이거나 끌 수 있습니다.

어텐션 설계는 KV 캐시를 작게 유지합니다. 전체 시퀀스에 대해 KV를 유지하는 것은 9개의 전체 어텐션 레이어뿐입니다: 토큰당 9 * 2 * 8 * 128 값, 이는 BF16에서 36KiB, FP8에서 18KiB입니다. 전체 1M 토큰 컨텍스트는 샤딩 전에 약 18GiB의 FP8 KV 캐시가 필요합니다.

MoE 및 잔차 스트림: FP32 라우터, 샌드위치 norm, 깊이 스케일링 잔차

Beam은 NVFP4에서 d_ff 2,560을 가진 256개의 라우팅 전문가를 가지고 있습니다. 라우팅은 DeepSeek-V3와 일치합니다: 시그모이드 점수와 보정 바이어스가 8개 전문가 그룹 중 최상위 4개를 선택한 다음, 최상위 8개 전문가를 선택하고, 재정규화된 가중치는 2.5로 스케일링됩니다. DeepSeek-V3와 달리 Beam에는 공유 전문가가 없습니다. 레이어 0~2는 MoE 대신 d_ff 20,480을 가진 밀집 SwiGLU MLP를 사용합니다.

라우터는 완전히 FP32입니다: 5120 x 256 라우터 가중치, 누적, 로짓 모두 FP32로 실행됩니다.

두 가지 정규화 선택이 잔차 스트림을 형성합니다. 모든 하위 레이어에는 앞뒤로 RMSNorm이 있습니다("샌드위치 norm"). n번째 잔차 추가는 h + x 대신 h <- sqrt((n-1)/n) * h + x/sqrt(n)을 계산합니다. 따라서 단위 분산 하위 레이어 출력으로 잔차 스트림은 깊이에 따라 증가하지 않고 모든 깊이에서 단위 분산을 유지합니다.

에이전트 기반 최적화로 배치-1 디코딩 2.71배 향상

vLLM의 첫 번째 Beam 구현은 최적화되지 않았고 거의 순수 PyTorch였습니다. 그런 다음 코딩 에이전트가 자율적으로 최적화했습니다: 모델을 프로파일링하고, 커널을 작성하고, 각 변경 사항을 벤치마킹했습니다. TP8에서 7가지 변경으로 배치-1 디코딩이 단계당 8.95ms에서 3.30ms로, 8 x 2K 프리필이 328ms에서 176ms로 빨라졌습니다.

7가지 변경 사항은 다음과 같습니다:

  1. norm, 잔차, norm 융합(-40% 디코딩). 104개 하위 레이어 경계 각각에서 사후 norm, 스케일링된 잔차 추가, 다음 사전 norm이 하나의 Triton 커널로 실행되며, FP32 수학과 참조 구현의 BF16 반올림 지점을 사용합니다.
  2. Q/K norm, RoPE/SSMax, 출력 게이트 융합. 헤드별 Q/K RMSNorm 및 RoPE(또는 SSMax)가 하나의 커널로 실행되고, 시그모이드 게이트는 다른 커널로 실행됩니다.
  3. 두 텐서 병렬 all-reduce를 norm 커널에 융합, FlashInfer의 멀티캐스트 all-reduce 기반. MoE 후에 동일한 커널이 상위 8개 전문가 출력도 결합합니다.
  4. FP32 라우터용 특수 GEMM 커널. 라우터는 이전에 일반 경로로 폴백되었습니다. vLLM의 소형 배치 FP32 라우터 커널은 이제 FP32 가중치, 누적, 로짓을 유지하면서 Beam의 5120 x 256 형태를 지원합니다.
  5. QKV 및 게이트 프로젝션을 저지연 GEMM으로 병합. Q, K, V, 게이트 가중치는 로드 시 하나의 행렬로 패킹되어 네 개의 프로젝션이 하나의 GEMM으로 실행됩니다. 소형 배치의 경우 저지연 CuTe DSL GEMM 또는 FlashInfer의 TinyGEMM을 사용합니다.
  6. 융합 KV 캐시 업데이트. Q/K norm 및 RoPE 커널은 K와 V를 KV 캐시에 기록하여 별도의 캐시 쓰기 커널을 제거합니다.
  7. 세 개의 밀집 레이어에 대한 단일 커널 밀집 MLP 활성화.

레이턴시 측정은 GB200에서 TP8 모델 포워드 레이턴시로, 배치-1 디코딩과 약 100K 컨텍스트에서 8 x 2,048 토큰 프리필에 대해 수행되었습니다.

융합 디코딩 그래프는 하나의 MoE 레이어를 13개 커널로 실행

이러한 변경 후, 하나의 MoE 레이어는 TP8에서 배치-1 디코딩 시 13개 커널로 실행됩니다. 그 중 6개 커널은 Beam용으로 작성, 융합 또는 확장되었습니다. 모든 작은 Beam 특정 연산은 이웃에 병합되고, 두 all-reduce는 norm 커널 내부에 있습니다.

나머지 7개 커널은 라이브러리에서 가져옵니다: TRTLLM-gen 어텐션, cuBLAS 출력 프로젝션, TRTLLM NVFP4 MoE 경로(양자화, 두 라우팅 커널, 두 그룹 GEMM). 남은 시간의 대부분은 이러한 라이브러리 커널에 있으며, vLLM은 이를 다음 최적화 대상으로 식별합니다.

어텐션 DP + MegaMoE로 초당 168.8K 토큰 달성

높은 동시성 영역에서 Beam 성능은 FlashInfer의 CuTe DSL NVFP4 MegaMoE 커널로 구동되는 어텐션 데이터 병렬 처리와 MegaMoE를 결합하여 더욱 향상됩니다. 이 커널은 통신을 전문가 계산과 융합할 수 있습니다. Finalize + norm 경계 융합은 이 구성에서도 적용 가능합니다.

구성 피크 입력 토큰 처리량
어텐션 DP + MegaMoE 168.8K 토큰 / 초
TP (flashinfer_trtllm, flashinfer_nvlink_one_sided) 98.7K 토큰 / 초

어텐션 DP + MegaMoE는 TP4보다 피크 프리필 처리량이 1.71배 높았습니다. 벤치마크는 4 x GB300에서 입력 시퀀스 길이 32K 토큰, 출력 시퀀스 길이 1 토큰으로 프리필 성능을 분리하여 사용했습니다.

활성화하려면 MoE 및 병렬 설정을 다음으로 교체하십시오:

--moe-backend flashinfer_moe_ep_cutedsl \
  --data-parallel-size 4 \
  --enable-expert-parallel

DFlash2 추측 디코딩: 동시성 1에서 초당 708토큰

Beam에는 DFlash2 드래프트 헤드가 포함되어 있으며, vLLM은 day 0부터 지원합니다. DFlash2는 한 번에 하나의 토큰을 드래프트하는 대신 대상 모델의 hidden states에서 단일 포워드 패스로 전체 토큰 블록을 예측합니다.

드래프터는 5개의 슬라이딩 윈도우 레이어로 구성되며, 5개의 대상 레이어에서 hidden states를 읽고 대상의 임베딩과 LM 헤드를 공유합니다. 그룹 컨볼루션과 작은 후보 선택기가 병렬 드래프터가 잃는 왼쪽에서 오른쪽 의존성의 일부를 복원합니다.

"enable_adaptive_verification": true로 활성화되는 적응형 검증은 각 드래프트 토큰이 수용될 가능성을 예측하고 단계당 고정된 수 대신 비용이 들만한 토큰만 검증합니다. DFlash2에는 신뢰도 헤드가 없으므로 수용 추정은 드래프터 자체 로짓에서 나옵니다. 드래프터의 컨볼루션, norm, KV 쓰기는 융합되어 있으며 전체 드래프트 단계는 CUDA 그래프에서 실행됩니다. FP8 KV 캐시는 수용 변경 없이 드래프터에서도 작동합니다.

TP8에서의 서빙 결과는 SPEED-Bench의 throughput_1k 분할의 낮은 엔트로피 범주에서 나옵니다: 약 1K 토큰 프롬프트, 주로 코드 완성 및 텍스트 정렬, 탐욕 디코딩, 최대 4K 출력 토큰. 낮은 엔트로피 텍스트는 예측하기 쉽기 때문에 개방형 채팅에서는 더 낮은 수용을 기대해야 합니다.

동시성 고정 드래프트 (tok/s) 적응형 (tok/s) 변경 수용 길이
1 708 707 -0.1% 3.71 -> 3.61
8 3,232 3,221 -0.3% 3.67 -> 3.57
32 7,744 8,275 +6.9% 3.68 -> 3.29
128 15,012 16,632 +10.8% 3.68 -> 3.23

적응형 검증은 부하가 걸릴 때 효과적입니다. 동시성 128에서 출력 처리량을 10.8% 높이고 p99 TPOT를 9% 줄이면서 단계당 검증되는 드래프트 토큰 수는 줄입니다. 수용 길이는 약간 감소하지만 절약된 계산이 이를 상쇄합니다. 낮은 동시성에서는 추가 드래프트 토큰이 거의 무료이므로 효과는 중립적입니다.

배포 지침: 병렬 설정, 추론, 도구 호출

vLLM 게시물 상단의 명령은 권장 기본값입니다: 단일 GB300 노드에서 TP4, DFlash2 및 FP8 KV 캐시 사용. 적응형 검증을 활성화하려면 --speculative-config에 "enable_adaptive_verification": true를 추가하십시오. 위에서 설명한 다른 최적화는 기본적으로 켜져 있습니다.

가장 낮은 레이턴시를 위해 TP4 또는 TP8을 사용하십시오. 더 높은 처리량을 위해 데이터 병렬 어텐션과 전문가 병렬 처리를 사용하십시오(예: -dp 4 --enable-expert-parallel).

beam 추론 및 도구 호출 파서는 체크포인트에 포함된 응답 템플릿을 읽습니다. 기본적으로 vLLM은 모든 응답을 해당 템플릿으로 제한하므로 도구 호출은 항상 파싱되고 각 도구의 JSON 스키마와 일치합니다. response_format 및 thinking_token_budget는 평소처럼 작동합니다. 기본 제약을 끄려면 VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0을 설정하십시오.

SkyRL을 통한 Day 0 강화 학습 지원

vLLM은 SkyRL 팀과 협력하여 Beam에 대한 day 0 강화 학습 지원을 가능하게 했습니다. LoRA, BF16 전체 미세 조정, MXFP8 RL이 기본적으로 지원되며, 500B 규모를 위한 업그레이드된 RL 스택이 제공됩니다.

SkyRL 팀은 vLLM을 추론에, Megatron을 훈련에 사용하여 DAPO 레시피로 Beam을 사후 훈련했습니다. LoRA, 전체 미세 조정, MXFP8 RL은 50개의 훈련 단계 또는 100개의 옵티마이저 업데이트(추론 minimal)에 걸쳐 유사한 성능 향상을 달성하며, AIME25에서 평균 정확도(avg@8)가 약 61%에서 약 79%로 점프합니다. Megatron과 vLLM 간의 평균 절대 로그 확률 차이는 LoRA 및 BF16 전체 미세 조정에 대해 Rollout Routing Replay(R3)를 활성화한 상태에서 0.01 미만, MXFP8 RL의 경우 0.03 미만으로 유지됩니다.

Beam 서빙 스택을 구축한 사람들

vLLM은 이 작업을 Inferact, Reflection, Nvidia 팀 간의 긴밀한 협력으로 설명합니다. Inferact 기여자는 Woosuk Kwon, Ziqi Zhao, Ao Shen, Tiezhen Wang입니다. Reflection 기여자는 Kaixing Wu, Yoni Gozlan, Neil Gollapudi, Thibaut Barroyer, Aymen Kallala, Matthias Reso, Florent Gbelidji입니다. Nvidia 기여자는 Siyuan Fu입니다.

Sources