본문으로 건너뛰기

[vllm] vLLM, ColBERT/ColPali 문서 랭킹 속도 혁신: Flash-MaxSim Triton 커널 통합

PR 링크: vllm-project/vllm#40337 상태: Merged | 변경: +508 / -4

들어가며

vLLM은 대규모 언어 모델(LLM) 서빙을 위한 고성능 추론 엔진으로, 특히 검색 증강 생성(Retrieval-Augmented Generation, RAG) 시스템에서 중요한 역할을 합니다. RAG 시스템은 외부 지식을 활용하여 LLM의 답변 정확성과 관련성을 높이지만, 관련 문서를 빠르고 정확하게 찾는 것이 성능의 핵심입니다. ColBERT와 ColPali와 같은 모델은 문서 랭킹(document scoring) 단계에서 쿼리와 문서 간의 유사도를 계산하는데, 이 과정이 병목 현상을 일으킬 수 있습니다.

이번 vLLM 프로젝트의 PR(#40337)은 이러한 문서 랭킹 성능을 획기적으로 개선하기 위해 Flash-MaxSim Triton 커널을 통합했습니다. 이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 왜 이러한 최적화가 중요한지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 기존의 CPU 기반 또는 일반적인 GPU 기반 MaxSim 스코어링 방식을 고도로 최적화된 Triton 커널로 대체하는 것입니다. 이를 통해 GPU의 연산 능력을 최대한 활용하고, 데이터 복사 및 파이썬 오버헤드를 최소화합니다.

1. Flash-MaxSim Triton 커널 도입 (vllm/v1/pool/flash_maxsim/)

PR은 vllm/v1/pool/flash_maxsim/ 디렉토리에 새로운 Triton 커널들을 추가했습니다. 이 커널들은 MaxSim 스코어링 계산을 GPU에서 매우 효율적으로 수행하도록 설계되었습니다.

  • flash_maxsim.py: 하드웨어별 최적화 설정을 포함하는 메인 포워드 커널입니다.
  • flash_maxsim_rerank.py: 제로-카피(zero-copy) 재랭킹 커널로, 데이터를 복사하지 않고 직접 GPU 메모리에 접근하여 처리합니다.
  • flash_maxsim_varlen.py: 가변 길이 시퀀스를 효율적으로 처리하기 위한 cu_seqlens 기반 커널입니다.
  • flash_maxsim_advanced.py: 쿼리 재사용(q-reuse) 및 K 분할(split-K)과 같은 고급 최적화 변형을 포함합니다.

이 커널들은 torch.compile과 같은 최신 기술을 활용하여 동적 컴파일 및 최적화를 수행합니다.

2. Late Interaction 스코어링 로직 변경 (vllm/v1/pool/late_interaction.py)

기존의 compute_maxsim_score_batched 함수는 이제 사용 가능한 경우 Flash-MaxSim 커널로 디스패치하도록 수정되었습니다. 이는 여러 스코어링 경로를 제공하며, 특정 조건에 따라 최적의 경로를 자동으로 선택합니다.

  • 제로-카피 경로 (기본값): GPU 사용, 임베딩 차원 d >= 16, 청크드 프리필(chunked prefill) 미사용, 기본 풀링 파라미터 사용 시 활성화됩니다. TokenEmbeddingPoolerHead.project_batch()를 한 번만 호출하여 모든 문서의 임베딩을 투영하고, 각 문서 요청은 이 투영된 텐서의 뷰(view)로 처리됩니다. 스코어 행렬은 GPU 메모리에 생성되지 않으며, 융합된 커널이 SRAM을 통해 직접 처리합니다.
--- a/vllm/v1/pool/late_interaction.py
+++ b/vllm/v1/pool/late_interaction.py
@@ -152,3 +152,107 @@ def test_invalid_query_uses_raises():
             req_ids=["query-req"],
             finished_mask=[True],
         )
+
+
+# ---------------------------------------------------------------------------
+# Fused Triton scoring path (PR #40337).
+# ---------------------------------------------------------------------------
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_score_flash_matches_fp64_oracle():
+    """Shared-query scoring through the fused kernel must match a per-pair
+    fp64 MaxSim oracle on ragged doc lengths, including length-1 docs."""
+    torch.manual_seed(0)
+    d = 128
+    q = torch.randn(32, d, device="cuda", dtype=torch.float16)
+    doc_lengths = [180, 37, 512, 1, 300, 64, 1030, 256]
+    docs = [
+        torch.randn(ld, d, device="cuda", dtype=torch.float16) for ld in doc_lengths
+    ]
+
+    runner = LateInteractionRunner()
+    assert runner._flash_enabled
+    scores = runner._score([q] * len(docs), docs)
+    assert runner._flash_enabled, "kernel path must not have fallen back"
+
+    for i, doc in enumerate(docs):
+        ref = (q.double() @ doc.double().T).max(dim=1).values.sum().float()
+        torch.testing.assert_close(
+            scores[i].to(torch.float32), ref, atol=5e-2, rtol=1e-3
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_score_multi_query_uses_reference_path():
+    """Distinct queries per pair route to the reference scorer and agree
+    with it exactly (no kernel dispatch for the N:N pattern)."""
+    torch.manual_seed(1)
+    d = 128
+    queries = [torch.randn(16, d, device="cuda") for _ in range(3)]
+    docs = [torch.randn(ld, d, device="cuda") for ld in (50, 3, 77)]
+
+    runner = LateInteractionRunner()
+    scores = runner._score(queries, docs)
+    ref = compute_maxsim_score_batched(queries, docs)
+    for s, r in zip(scores, ref):
+        torch.testing.assert_close(s, r)
+
+
+def test_score_flash_failure_disables_and_falls_back(monkeypatch):
+    """A kernel failure must serve the batch via the reference scorer and
+    disable the kernel path for the rest of the process."""
+    runner = LateInteractionRunner()
+    runner._flash_enabled = True
+
+    def _boom(query, docs):
+        raise RuntimeError("synthetic kernel failure")
+
+    monkeypatch.setattr(runner, "_score_flash_shared_query", _boom)
+    q = torch.randn(8, 32)
+    docs = [torch.randn(5, 32), torch.randn(7, 32)]
+    # CPU tensors already skip the kernel; force the dispatch condition.
+    monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
+    scores = runner._score([q, q], docs)
+    assert not runner._flash_enabled
+    ref = compute_maxsim_score_batched([q, q], docs)
+    for s, r in zip(scores, ref):
+        torch.testing.assert_close(s, r)
+
+
+def test_disabled_flag_uses_reference_path():
+    runner = LateInteractionRunner(enable_flash=False)
+    assert not runner._flash_enabled
+
+
+# ---------------------------------------------------------------------------
+# Autotune config pruning.
+# ---------------------------------------------------------------------------
+def _prune(monkeypatch, budget: int, named_args: dict):
+    from vllm.v1.pool.flash_maxsim import _common
+
+    monkeypatch.setattr(_common, "_smem_budget", lambda: budget)
+    return _common._prune_configs(_common._get_configs(), named_args)
+
+
+def _est(cfg, d_pad: int) -> int:
+    bq, bd = cfg.kwargs["BLOCK_Q"], cfg.kwargs["BLOCK_D"]
+    return (bq * d_pad + bd * d_pad) * 2 + bq * bd * 4
+
+
+def test_prune_configs_uses_padded_dim(monkeypatch):
+    """d=513 pads to 1024: survivors must fit the budget at d_pad, not at
+    the un-padded d (which would admit ~2x-oversized tiles)."""
+    budget = 166_912  # A100 opt-in
+    survivors = _prune(monkeypatch, budget, {"Lq": 1024, "d": 513, "d_pad": 1024})
+    assert survivors
+    assert all(_est(c, 1024) <= budget for c in survivors)
+
+
+def test_prune_configs_no_rejected_fallback(monkeypatch):
+    """When nothing fits, the fallback is the single smallest-footprint
+    config — never a slice of rejected configs."""
+    from vllm.v1.pool.flash_maxsim import _common
+
+    survivors = _prune(monkeypatch, 10_000, {"Lq": 1024, "d_pad": 1024})
+    assert len(survivors) == 1
+    smallest = min(_common._get_configs(), key=lambda c: _est(c, 1024))
+    assert survivors[0].kwargs == smallest.kwargs

  • 플래시-팩 경로: 제로-카피 경로가 비활성화되거나 특정 풀링 파라미터(예: Matryoshka)를 사용할 때 사용됩니다. 모든 문서들을 torch.cat으로 연결하고 flash_maxsim_packed 커널을 cu_seqlens와 함께 호출합니다. 패딩 토큰은 cu_seqlens를 통해 건너뜁니다.

  • 바닐라 경로: CPU 텐서, 임베딩 차원 d < 16 (Triton tl.dot 최소 요구사항), 또는 VLLM_FORCE_VANILLA_MAXSIM=1 환경 변수가 설정된 경우 사용됩니다. 이는 기존의 서브-배치된 패딩된 bmm 연산을 그대로 사용합니다.

Triton의 자동 튜닝(autotune)은 런타임 시 성능을 최적화하기 위해 다양한 블록 크기(BLOCK_Q, BLOCK_D) 조합을 탐색합니다. 이 PR에서는 LateInteractionRunner.__init__에서 72가지의 버킷 조합을 미리 컴파일하여 첫 요청 시 자동 튜닝이 발생하는 오버헤드를 줄였습니다.

3. 풀링 헤드 투영 로직 변경 (vllm/model_executor/layers/pooler/tokwise/heads.py)

TokenEmbeddingPoolerHead.project_batch() 메서드가 수정되었습니다. 이 메서드는 제로-카피 경로에서 핵심적인 역할을 하며, 입력 hidden_states를 투영하고 정규화(normalization)합니다. 리뷰어의 지적에 따라, 이 함수는 이제 입력 정밀도(input precision)가 아닌 head_dtype으로 업캐스팅한 후 투영을 수행하여 정확도를 높였습니다.

--- a/vllm/model_executor/layers/pooler/tokwise/heads.py
+++ b/vllm/model_executor/layers/pooler/tokwise/heads.py
@@ -60,7 +60,12 @@
         # Note: If the head is not a Linear layer, we assume it has a similar
         # projection mechanism.
         if isinstance(self.head, (torch.nn.Linear, torch.nn.Identity)):
-            projected_batch = self.head(hidden_states)
+            # Upcast to head_dtype before projection to match forward_chunk
+            # semantics and ensure bit-exactness with the reference path.
+            # This is done in chunks to avoid large intermediate allocations.
+            projected_batch = torch.empty_like(hidden_states, dtype=self.head_dtype)
+            for i in range(0, hidden_states.shape[0], 16384):
+                projected_batch[i:i+16384] = self.head(hidden_states[i:i+16384].to(self.head_dtype))
         else:
             # Fallback for non-Linear heads: assume they handle dtype internally.
             projected_batch = self.head(hidden_states)

4. 환경 변수 및 기능 토글

  • VLLM_DISABLE_ZEROCOPY=1: 제로-카피 경로를 비활성화하고 플래시-팩 경로를 사용합니다.
  • VLLM_FORCE_VANILLA_MAXSIM=1: 플래시 커널을 완전히 비활성화하고 기존 바닐라 bmm을 사용합니다.
  • PoolerConfig.enable_flash_late_interaction: API 서버 시작 시 --no-enable-flash-late-interaction 플래그를 통해 제어할 수 있으며, 엔진의 PoolerConfig에도 반영되어 플래시 커널 사용 여부를 결정합니다.

왜 이게 좋은가?

1. 압도적인 성능 향상

PR에서 제공된 결과는 Flash-MaxSim 커널 통합이 가져온 성능 향상을 명확히 보여줍니다.

커널 레벨 성능 (가변 길이 문서):

Workload Vanilla Flash-Packed Zero-Copy
B=1K, Ld=10-100 42.6ms 0.58ms 0.05ms (~788×)
B=5K, Ld=10-100 220.3ms 2.16ms 0.08ms (~2,830×)
B=10K, Ld=10-100 437.6ms 4.04ms 0.14ms (~3,034×)
B=10K, Ld=50-500 472.9ms 4.22ms 0.53ms (~892×)
ColPali B=1K, Ld≈1030 40.4ms 0.72ms 0.22ms (~184×)

특히 제로-카피 경로는 기존 바닐라 방식 대비 수백 배에서 수천 배까지 빠른 성능을 보여줍니다. 이는 파이썬 측의 패딩, FP32 캐스팅, 할당 오버헤드를 제거하고 GPU 커널에서 직접 처리하기 때문입니다.

엔드-투-엔드 /v1/score 성능:

Docs/req Conc Flash TPS Vanilla TPS ΔTput ΔP95
500 8 6.2/s 5.2/s +19% +19%
500 16 6.0/s 5.2/s +17% +7%
1000 4 3.0/s 2.6/s +16% +13%
2000 4 1.6/s 1.3/s +23% +24%
5000 1 0.5/s 0.4/s +22% +18%
5000 2 0.6/s 0.5/s +15% +19%

엔드-투-엔드 성능 역시 최대 23%의 처리량(throughput) 증가와 24%의 P95 지연 시간 감소를 보여줍니다. 이는 특히 문서 수가 많을 때(예: 2000개 이상) 더욱 두드러집니다.

2. 정확성 보장

성능 향상과 더불어 정확성 또한 철저히 검증되었습니다.

  • 커널과 FP32 bmm 간의 최대 절대 오차는 4e-6에 불과합니다.
  • 5000개의 실제 ColBERT 문서에 대한 엔드-투-엔드 테스트에서 최대 절대 차이는 0.001 미만이며, 상위 3개 랭킹이 동일하게 유지되었습니다.

이는 FP16 연산의 비결정성으로 인한 미세한 순위 변동(4-5번째 순위)을 제외하고는 기존 로직과 동일한 결과를 보장함을 의미합니다.

3. 메모리 효율성 및 유연성

제로-카피 경로는 데이터를 GPU 메모리 내에서 직접 처리하므로 불필요한 데이터 복사를 피하고 메모리 사용량을 줄입니다. 또한, 다양한 환경과 제약 조건에 맞춰 제로-카피, 플래시-팩, 바닐라 경로 간의 자동 전환 및 명시적 제어 기능을 제공하여 유연성을 높였습니다.

4. 기술적 교훈

  • 커널 최적화의 중요성: 고성능 컴퓨팅에서는 Triton과 같은 DSL(Domain-Specific Language)을 활용한 커널 수준의 최적화가 전체 시스템 성능에 지대한 영향을 미칠 수 있습니다.
  • 제로-카피 및 데이터 재사용: 데이터를 GPU 메모리 내에서 최대한 재사용하고 복사를 최소화하는 설계는 성능 향상의 핵심입니다.
  • 다중 경로 전략: 다양한 하드웨어 및 사용 사례에 맞춰 여러 구현 경로를 제공하고, 런타임에 최적의 경로를 선택하는 전략은 견고성과 성능을 동시에 확보하는 데 유용합니다.
  • 정확성 검증: 성능 최적화 과정에서 정확성을 보장하기 위한 철저한 테스트(단위 테스트, 비교 벤치마크, FP64 오라클 비교 등)가 필수적입니다.
  • 점진적 배포 및 롤백: 새로운 최적화 기능이 실패할 경우, 기존의 안정적인 경로로 자동 전환되고 해당 기능이 비활성화되는 메커니즘은 시스템 안정성을 유지하는 데 중요합니다.

리뷰 피드백 반영

리뷰 과정에서 몇 가지 중요한 피드백이 반영되었습니다:

  • project_batch의 정규화 및 dtype 처리: 초기에는 project_batch가 정규화를 수행하지 않거나 head_dtype으로의 업캐스팅이 누락되어 정확성 문제가 발생할 수 있다는 지적이 있었습니다. PR은 forward_chunk와 동일한 head_dtype 업캐스팅 및 정규화 로직을 project_batch에 적용하고, 이를 검증하는 테스트 케이스를 추가했습니다.
  • 메모리 사용량 최적화: 제로-카피 경로에서 project_batch의 전체 출력을 한 번에 할당하는 대신, 청크 단위로 처리하여 최대 메모리 사용량을 줄였습니다.
  • GPU-CPU 동기화 제거: PoolingCursor에서 CPU 미러를 사용하여 GPU-CPU 동기화 없이 필요한 메타데이터에 접근하도록 개선했습니다.
  • 런타임 폴백 및 비활성화: 커널 실행 중 오류가 발생하면, 해당 요청은 기존의 바닐라 스코어러로 처리되고, 이후 동일한 프로세스에서는 Flash-MaxSim 기능이 영구적으로 비활성화되도록 하여 안정성을 높였습니다.
  • 자동 튜닝 최적화: Triton 자동 튜닝 시 공유 메모리(shared memory) 사용량 계산 시 패딩된 차원(d_pad)을 고려하고, 실제 장치의 공유 메모리 제한을 동적으로 쿼리하며, 최적의 설정이 없을 경우 가장 작은 설정을 반환하도록 개선했습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글