[vllm] vLLM의 Mamba 블록 테이블 가더(Gather) 최적화: 7개 커널을 1개로 통합하여 6배 속도 향상

PR 링크: vllm-project/vllm#58737 상태: Merged | 변경: +273 / -12

들어가며

vLLM의 Mamba 모델 구현체에서 mamba_get_block_table_tensor의 'align' 브랜치는 추론 과정에서 매우 빈번하게 호출되는 핵심 경로입니다. 기존에는 PyTorch 연산들을 조합하여 구현되어 있었는데, 이 과정에서 무려 7개의 커널이 순차적으로 실행되는 비효율이 발생했습니다. 특히 단일 요청(single-request) 디코딩 시에는 커널 실행 오버헤드가 실제 연산 시간보다 훨씬 커지는 문제가 있었습니다. 본 PR은 이 7개의 커널을 하나의 최적화된 Triton 커널로 통합하여, 커널 수준에서 약 6배의 성능 향상을 달성했습니다.

코드 분석

1. Triton 커널 구현 (vllm/v1/attention/backends/utils.py)

기존의 PyTorch 연산은 여러 단계의 커널 호출을 유발했습니다. 이를 하나의 _mamba_align_block_table_kernel로 통합했습니다.

Before (PyTorch):

start_indices = (seq_lens - 1) // kv_cache_spec.block_size
start_indices.clamp_(min=0)
offsets = torch.arange(1 + kv_cache_spec.num_speculative_blocks, ...)
indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
return torch.gather(block_table, 1, indices_to_gather)

After (Triton):

@triton.jit(do_not_specialize_on_alignment=["out_ptr", "block_table_ptr", "seq_lens_ptr"])
def _mamba_align_block_table_kernel(
    out_ptr, block_table_ptr, block_table_stride_0: tl.int64, ...
):
    req_idx = tl.program_id(0)
    seq_len = tl.load(seq_lens_ptr + req_idx * seq_lens_stride)
    start = tl.maximum((seq_len - 1) // block_size, 0)
    # ... (생략) ...
    vals = tl.load(block_table_ptr + req_idx * block_table_stride_0 + (start + offs) * block_table_stride_1, ...)
    tl.store(out_ptr + req_idx * n_out + offs, vals, mask=mask)

핵심 개선점은 do_not_specialize_on_alignment를 사용하여 MRV1 DBO 경로에서 발생하는 불필요한 JIT 재컴파일을 방지하고, block_table_stride_0을 tl.int64로 명시하여 인덱스 오버플로우 문제를 해결한 점입니다.

2. 테스트 및 검증 (tests/v1/attention/test_mamba_align_block_table.py)

단순히 성능만 개선한 것이 아니라, 비연속적(non-contiguous) 메모리 레이아웃이나 엣지 케이스(seq_len=0, num_reqs=0)에서도 기존 PyTorch 구현과 동일한 결과를 보장하도록 55개의 테스트 케이스를 추가했습니다.

왜 이게 좋은가

1. 커널 오버헤드 제거: 7개의 커널 호출을 1개로 줄임으로써, 특히 배치 사이즈가 작을 때 발생하는 커널 런칭 오버헤드를 획기적으로 줄였습니다. 측정 결과 10.81us에서 1.78us로 약 6배의 속도 향상을 보였습니다.
2. 정확성 보장: 기존 torch.gather가 지원하던 비연속적 메모리 레이아웃을 Triton 커널에서도 stride를 명시적으로 전달받아 처리함으로써, 공유 유틸리티로서의 안정성을 확보했습니다.
3. 일반적 교훈: 커널 수준의 마이크로 벤치마크가 작아 보일지라도, 모델 전체의 디코딩 루프에서 반복적으로 호출되는 경우 전체 처리량(throughput)에 유의미한 영향을 줄 수 있습니다. 특히 JIT 컴파일 오버헤드를 방지하기 위한 do_not_specialize_on_alignment와 같은 세밀한 튜닝이 실무 환경에서 얼마나 중요한지 보여줍니다.

참고 자료

• Triton JIT Documentation — 이 PR에서 사용된 Triton 커널 컴파일러 공식 문서

참고 자료

• https://triton-lang.org/main/index.html

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