[vllm] GLM MTP Draft에서 쓰지 않는 Post-Attention 행 건너뛰기: CUDA Graph 캡처 전 인덱스 바인딩으로 4.2% 처리량 향상

PR 링크: vllm-project/vllm#67104 상태: Merged | 변경: +22 / -9

들어가며

Speculative decoding의 MTP(Multi-Token Prediction) draft 단계에서는 요청 하나당 여러 개의 row가 forward를 통과합니다. 예를 들어 num_speculative_tokens=5라면 draft의 prefill 단계에서 요청당 q = 6개 row가 들어옵니다. 그런데 최종적으로 다음 draft 토큰을 만들 때 필요한 것은 요청당 마지막 유효 토큰의 hidden state 하나뿐입니다.

Attention과 KV cache 기록은 모든 row에 대해 수행해야 합니다. 이후 단계의 KV가 필요하기 때문입니다. 하지만 attention 이후의 post_attention_layernorm, MoE, 최종 shared_head.norm은 버려질 row에 대해서도 계산되고 있었습니다. PR 설명의 예시는 다음과 같습니다.

Before: attention/KV on 24 rows -> norm/MoE/final norm on 24 rows -> select 4
After:  attention/KV on 24 rows -> select 4 -> norm/MoE/final norm on 4 rows

흥미로운 점은 row pruning 경로(rows is not None 분기)가 이미 specialized decode runtime 안에 존재했다는 것입니다. 다만 speculator가 last_token_indices를 draft 모델에 넘겨주지 않아서 이 경로가 실제로 활성화되지 않았습니다. 이 PR은 크게 두 가지를 합니다. 첫째, CUDA Graph 캡처 전에 인덱스 텐서를 바인딩해 pruning 경로를 켭니다. 둘째, 줄어든 row 수에 맞춰 MoE 커널을 미리 빌드하고 정책을 일관되게 선택하도록 고칩니다.

코드 분석

1. speculator.py — 캡처 전에 인덱스 텐서 바인딩

# After
draft_model = load_eagle_model(target_model, self.vllm_config)
if hasattr(draft_model.model, "set_prefill_rows"):
    # bind before capture; replay reads the indices updated in place
    # by prepare_prefill_inputs, including rejected and padded rows.
    draft_model.model.set_prefill_rows(self.last_token_indices)

핵심은 CUDA Graph의 정적 메모리 주소 특성을 이용한다는 점입니다. CUDA Graph는 캡처 시점의 텐서 주소를 그대로 기록하고, replay할 때도 같은 주소를 읽습니다. 그래서 매 스텝 새 인덱스 텐서를 만들어 넘기면 graph가 그 값을 보지 못합니다.

이 PR은 speculator가 이미 보유한 영속 버퍼 self.last_token_indices를 모델 로드 직후, 즉 캡처 이전에 한 번만 바인딩합니다. 이후 prepare_prefill_inputs가 이 버퍼를 in-place로 갱신하면, replay되는 graph가 자동으로 최신 인덱스를 읽습니다. 주석에 적혀 있듯 rejected row와 padded row도 이 버퍼 안에서 처리되므로 graph 형태는 고정된 채로 유지됩니다.

hasattr 가드를 둔 덕분에 set_prefill_rows를 구현하지 않은 다른 MTP 모델에는 영향이 없습니다.

2. draft.py — 줄어든 row에서도 원래 batch 기준으로 MoE 정책 선택

# Before
x, partial = self.mlp(layer, x)

# After
x, partial = self.mlp(
    layer, x, moe_policy_tokens=tokens if rows is not None else None
)

pruning 후 x.shape[0]는 tokens / q로 작아집니다. 그런데 MoE 커널 선택 정책(per_route vs grouped)은 원래 batch 크기 기준으로 짜여 있습니다. 그래서 pruning이 켜졌을 때(rows is not None)는 원래 tokens 수로 정책을 조회하도록 넘겨줍니다.

생성자에서도 draft의 q를 전달합니다.

self.mega = self._build_mega(prefill_q=self.q)

3. decode.py — pruned 크기용 mega 커널 사전 빌드

def _build_mega(self, prefill_q: int = 1) -> dict[int, Any]:
    ...
    if prefill_q > 1:
        # only prune when requests rows > 1, eg. MTP 5-> prefill_q = 6
        for sizes in (per_route_tokens, grouped_tokens):
            sizes.extend(m // prefill_q for m in tuple(sizes) if m % prefill_q == 0)
            sizes[:] = sorted(set(sizes))

mega MoE 커널은 token 수별로 미리 특수화되어 빌드됩니다. pruning 후 실제로 커널에 들어오는 row 수는 m // q입니다. 그래서 q로 나누어떨어지는 기존 크기마다 축소된 크기를 같은 정책 리스트에 추가합니다. 그 결과 batch 24는 정책상 "24 tokens"로 판단되면서도, 실제로는 4-row 커널이 준비되어 있게 됩니다. tuple(sizes)로 순회해 순회 중 변경 문제를 피하고, set으로 중복을 제거한 뒤 정렬합니다.

정책 선택 로직도 한 번만 계산해 전달하도록 정리했습니다.

# Before
if layer.layer_idx in self.mega and self.policy("moe", x.shape[0]):
    if partial:
        return self._mega_moe(layer, x), True
    if self.policy("moe", x.shape[0]) == "grouped":

# After
moe_tokens = x.shape[0] if moe_policy_tokens is None else moe_policy_tokens
choice = self.policy("moe", moe_tokens)
if layer.layer_idx in self.mega and choice:
    if partial:
        return self._mega_moe(layer, x, choice), True
    if choice == "grouped":

_mega_moe 내부도 self.policy("moe", x.shape[0])를 다시 호출하던 방식에서 choice 인자를 받는 방식으로 바뀌었습니다. 이렇게 하지 않으면 pruned x.shape[0]로 정책을 재조회해 바깥과 다른 커널 분기를 탈 수 있습니다. 결정을 한 곳에서 내리고 전파하는 것이 정합성 측면에서 훨씬 안전합니다.

왜 이게 좋은가

성능 수치 (GLM, TP4, MTP 5 tokens, concurrency 1)

지표	main	PR
Output throughput	620.58 tok/s	647.26~647.57 tok/s (+4.3%)
Mean ITL	6.36 ms	6.10 ms
Mean TPOT	1.58 ms	1.52 ms
Mean TTFT	30.25 ms	28.88 ms

Acceptance rate(60.50%)와 per-position acceptance는 완전히 동일합니다. draft 출력 자체는 바뀌지 않고 순수하게 계산량만 줄었다는 강력한 증거입니다. AIME25 정확도는 0.7333에서 0.7667로 나왔지만, stderr(약 0.08) 범위 안의 차이로 보는 것이 맞습니다.

MoE는 GLM 같은 모델에서 레이어 비용의 큰 몫을 차지합니다. q=6 기준으로 post-attention 연산 row를 1/6로 줄이면 draft 단계 비용이 눈에 띄게 줄어듭니다. 매 decode step마다 draft가 실행되므로 이 절감이 ITL에 그대로 반영됩니다.

일반적 교훈

1. "마지막에 select"는 최대한 앞당겨라. Attention처럼 모든 row가 필요한 연산 직후에 gather를 수행하면, 이후의 token-wise 연산(norm, MLP/MoE)은 필요한 row만 처리하면 됩니다. LLM serving에서 logits 계산 전 마지막 토큰만 고르는 것과 같은 원리입니다.
2. CUDA Graph와 동적 인덱스는 영속 버퍼 + in-place 갱신으로 공존시킨다. 캡처 전에 주소를 고정하고 값만 바꾸는 패턴은 graph 친화적인 동적 동작의 정석입니다.
3. 이미 만든 최적화가 실제로 켜져 있는지 확인하라. 이 PR의 핵심 변경은 4줄짜리 바인딩입니다. 존재하지만 연결되지 않은 fast path는 생각보다 흔합니다.
4. 정책 결정은 한 번, 전달은 명시적으로. shape가 변하는 최적화를 도입하면 shape 기반의 암묵적 재판단이 다른 분기를 타게 만들 수 있으므로, choice를 인자로 전파하는 방식이 안전합니다.

참고 자료

• CUDA Graphs - PyTorch — 캡처 시 고정되는 메모리 주소와 in-place 입력 갱신 패턴 설명
• vLLM Speculative Decoding — MTP를 포함한 vLLM speculative decoding 공식 문서

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