[vllm] 한 프로그램이 4개 head를 처리하게 하라: DeepSeek-V4 inverse RoPE + FP8 quant 커널 3.5배 가속기

PR 링크: vllm-project/vllm#63540 상태: Merged | 변경: +102 / -61

들어가며

DeepSeek-V4 attention 출력은 wo_a projection으로 들어가기 전에 두 가지 처리를 거칩니다. 하나는 RoPE를 되돌리는 inverse RoPE이고, 다른 하나는 block FP8 quantization입니다. vLLM은 이 두 단계를 fused_inv_rope_fp8_quant라는 Triton 커널 하나로 합쳐 두었습니다.

문제는 이 커널이 (token, head) 하나당 single-warp 프로그램 하나를 띄우는 구조였다는 점입니다. 16k 토큰에 rank당 32 head라면 프로그램이 524,288개 생깁니다. 각 프로그램이 옮기는 데이터는 약 1.5 KiB(bf16 1 KiB 읽기, FP8 512B 쓰기)에 불과합니다. 그 결과 B200에서 실효 대역폭이 1.5 TB/s에 그쳤고, DeepSeek-V4-Flash TP2 기준 16k prefill step의 8.6%를 이 커널 하나가 차지했습니다.

PR 작성자는 흥미로운 증거를 하나 제시합니다. 기존 커널은 H100과 H200에서 똑같이 631 µs가 걸렸습니다. H200의 HBM 대역폭이 43% 더 높은데도 시간이 같다는 것은 이 커널이 bandwidth-bound가 아니라 latency-bound라는 뜻입니다. SM당 동시에 상주할 수 있는 CTA 수에는 상한(약 32)이 있습니다. 그런데 CTA 하나가 들고 있는 데이터가 너무 작으니, 메모리 시스템에 in-flight 상태인 바이트가 HBM을 채우기에 부족했던 것입니다.

코드 분석

변경은 vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py 한 파일에 집중되어 있습니다. 테스트에는 케이스가 하나 추가되었습니다.

1. 프로그램당 head 수: BLOCK_H

def _heads_per_program(heads_per_group: int) -> int:
    """Heads one program rotates and quantizes; they must share an output group.

    One 512-dim head is only 1 KiB of bf16, so a program per (token, head) keeps
    too few bytes in flight to saturate HBM on large prefills.
    """
    return math.gcd(heads_per_group, 4)

gcd를 쓰는 이유가 핵심입니다. 한 프로그램이 처리하는 head들은 반드시 같은 output group(wo_a의 group)에 속해야 주소 계산이 단순해집니다. BLOCK_H가 heads_per_group의 약수이면 블록이 group 경계를 넘지 않습니다. V4 설정은 모두 group당 8 head이므로 BLOCK_H=4가 됩니다. group당 6 head 같은 비정형 설정이라면 2로 자연스럽게 fallback합니다.

커널 내부의 인덱싱은 다음과 같이 바뀌었습니다.

Before

pid_gh = tl.program_id(1).to(tl.int64)
g = pid_gh // heads_per_group
head_in_group = pid_gh % heads_per_group
global_head = pid_gh
qb_start = head_in_group * CHUNKS_PER_HEAD

After

pid_hb = tl.program_id(1).to(tl.int64)
# This program handles BLOCK_H consecutive heads; BLOCK_H divides
# heads_per_group, so they all land in group g.
first_head = pid_hb * BLOCK_H
g = first_head // heads_per_group
head_in_group = first_head % heads_per_group + tl.arange(0, BLOCK_H)

head_in_group이 스칼라에서 길이 BLOCK_H의 벡터로 바뀌었습니다. 그래서 이후의 모든 주소 계산은 [:, None] broadcasting으로 leading dimension 하나만 추가하면 됩니다. 입력 로드도 1D에서 2D tile로 확장됩니다.

input_base = (
    o_ptr
    + pid_token * o_stride_token
    + (first_head + tl.arange(0, BLOCK_H))[:, None] * o_stride_head
)
x = tl.load(input_base + offsets[None, :]).to(tl.float32)

이제 프로그램 하나가 [4, 512] tile을 한 번에 다룹니다. torch.profiler에서 보면 grid가 [16384, 32, 1]에서 [16384, 8, 1]로 줄었습니다. SM당 상주 CTA 수는 32에서 28로 비슷하게 유지되지만, CTA당 데이터가 4배이므로 in-flight 데이터는 약 3.5배로 늘어납니다.

2. RoPE partner를 다시 로드하지 않고 레지스터에서 swap

GPT-J 스타일 RoPE는 (even, odd) 쌍을 회전시킵니다. 기존 구현은 각 원소의 짝을 구하려고 offsets ^ 1 주소로 global memory를 한 번 더 읽었습니다.

Before

x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(
    tl.float32
)
cs_idx = tl.maximum(rope_local >> 1, 0)
cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0)
sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0)
x_add = x * cos_v + x_partner * sin_v
x_sub = x * cos_v - x_partner * sin_v
is_even = (rope_local & 1) == 0
rotated = tl.where(is_even, x_add, x_sub)
x = tl.where(is_rope, rotated, x)

After

pair = tl.arange(0, HEAD_DIM // 2)
is_rope_pair = pair >= NOPE_DIM // 2
cs_idx = tl.maximum(pair - NOPE_DIM // 2, 0)
cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)[None, :]
sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)[None, :]
even, odd = tl.split(tl.reshape(x, (BLOCK_H, HEAD_DIM // 2, 2)))
rope = is_rope_pair[None, :]
rot_even = tl.where(rope, even * cos_v + odd * sin_v, even)
rot_odd = tl.where(rope, odd * cos_v - even * sin_v, odd)
x = tl.reshape(tl.join(rot_even, rot_odd), (BLOCK_H, HEAD_DIM))

동작 방식은 이렇습니다. x를 (BLOCK_H, 256, 2)로 reshape한 뒤 tl.split으로 마지막 축을 even과 odd로 나눕니다. 회전을 적용한 다음 tl.join과 reshape로 원래 레이아웃에 되돌립니다.

cos/sin 로드도 pair 단위(256개)로 줄어 중복 로드가 사라졌습니다. 또 cos/sin은 [None, :]로 BLOCK_H head 전체에 broadcast되므로, 4개 head가 같은 cos/sin 벡터를 공유합니다.

fp32 연산의 종류와 순서는 그대로이므로 결과는 bitwise identical입니다. PR 설명에 따르면 이 변경은 1번과 함께 있을 때 의미가 큽니다. 4-head 블록에서 re-load를 유지하면 252 µs, 레지스터 swap을 쓰면 152 µs로 1.66배 차이가 납니다.

3. Quantization도 3D로 확장

# Before
x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
# After
x_3d = tl.reshape(tl.abs(x), (BLOCK_H, CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
block_absmax = tl.maximum(tl.max(x_3d, axis=2), eps)

UE8M0 packing도 같은 방식으로 axis=2 reduction으로 확장되었습니다. 알고리즘은 바뀌지 않았고 차원만 하나 늘었습니다.

4. grid 크기에 따라 num_warps 선택

class CompileKey:
    ...
    quantize: bool
    num_warps: int
...
num_warps=WarmupChoices(1, 4),

선택 규칙은 프로그램 수가 2048 미만이면(decode) 4 warps, 그 이상이면(prefill) 1 warp입니다.

• decode에서는 프로그램이 적어서 SM이 남습니다. 따라서 프로그램 하나 안에서 병렬성을 키우는 편이 낫습니다.
• prefill에서는 프로그램이 충분히 많습니다. 1 warp로 CTA 수를 최대화하는 편이 유리합니다.

num_warps를 CompileKey에 넣고 WarmupChoices(1, 4)로 warmup하기 때문에 두 variant 모두 startup 시점에 컴파일됩니다. 런타임 JIT 컴파일은 발생하지 않습니다.

5. 테스트: gcd fallback 커버

# H24_G4: 6 heads per group, so each program takes 2 heads instead of 4.
[(64, 8), (32, 4), (128, 8), (24, 4)],
ids=["H64_G8", "H32_G4", "H128_G8", "H24_G4"],

왜 이게 좋은가

성능 수치

아래는 H=32, 16k 토큰 기준 커널 시간입니다.

GPU	Before	After	속도 향상	실효 대역폭
B200	536 µs	152 µs	3.53x	5.3 TB/s
H200	631 µs	208 µs	3.04x	3.9 TB/s
H100	631 µs	293 µs	2.15x	2.7 TB/s

H100과 H200이 이제 서로 다른 시간을 보입니다. 커널이 latency-bound에서 bandwidth-bound로 전환되었다는 증거입니다.

end-to-end 효과는 다음과 같습니다.

• vllm bench forward: prefill 16384가 237 ms에서 223 ms로 줄었습니다(−6%).
• vllm bench serve: TTFT p99가 5.1~5.6% 감소했습니다.
• decode: 호출당 최대 0.3 µs 느려집니다. 43 layer를 합쳐도 TPOT의 0.2% 수준입니다.

리뷰어 algebra-MCX는 다른 소프트웨어 스택(Triton 3.6, CUDA 12.9)의 H200에서 독립적으로 검증했습니다. 16k 토큰에서 688 µs가 257 µs로 줄어 2.68x를 얻었고, 24개 케이스 모두 bitwise identical이었습니다. 절대 수치는 측정 방식(CUDA_LAUNCH_BLOCKING=1, CUPTI 미사용) 때문에 다르지만, 속도 향상 추세는 재현되었습니다. 벤치마크 스택이 달라도 결론이 유지된다는 점은 좋은 신호입니다.

설계상 돋보이는 점

• Batch invariance 유지: num_warps는 배치 크기에 따라 바뀝니다. 하지만 모든 config가 bitwise 동일한 결과를 내고 프로그램 하나는 한 토큰만 처리합니다. 따라서 어떤 토큰의 출력도 배치 크기에 영향을 받지 않습니다.
• 절제된 튜닝: H100에서는 BLOCK_H=8, num_warps=4가 조금 더 빨랐습니다. 그래도 차이가 4% 이내였기 때문에 arch별 분기를 넣지 않고 단일 정책을 택했습니다. 복잡도와 이득을 비교해 내린 합리적인 판단입니다.

일반적 교훈

1. "한 프로그램 = 가장 작은 논리 단위"는 함정이 될 수 있습니다. 프로그램 수가 많아도 각자 1 KiB만 옮긴다면 SM당 상주 CTA 수의 상한 때문에 memory parallelism이 부족해집니다. 프로그램당 작업량을 키워 in-flight 바이트를 늘려야 합니다.
2. latency-bound 여부를 진단하는 간단한 방법이 있습니다. 대역폭이 다른 두 GPU에서 같은 시간이 나온다면 메모리 대역폭이 병목이 아닙니다.
3. 이미 레지스터에 있는 데이터를 다시 load하지 마세요. offsets ^ 1 같은 gather 형태의 재로드는 tl.split/tl.join으로 레지스터 안에서 해결할 수 있습니다.
4. 블록 크기는 데이터 구조의 경계를 존중해야 합니다. gcd로 정하면 group 경계를 넘지 않는다는 보장을 공짜로 얻습니다.
5. decode와 prefill은 다른 config가 필요할 수 있습니다. 이때 JIT를 피하려면 warmup 단계에서 모든 variant를 미리 컴파일해 두어야 합니다.

참고 자료

• Triton tl.split — even/odd 분리에 사용
• Triton tl.join — 회전 결과 재결합에 사용
• Triton tl.reshape — (BLOCK_H, HEAD_DIM//2, 2) 변환
• Triton tl.where — RoPE 영역 선택
• Python math.gcd — BLOCK_H 계산

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