[vllm] vLLM DeepSeek-V4 성능 최적화: Eager Break 시 Workspace 재사용 개선

2026년 7월 31일수정: 2026년 7월 31일

PR 링크: vllm-project/vllm#49236 상태: Merged | 변경: +354 / -30

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 하지만 모델의 크기가 커질수록 추론 성능, 특히 응답 속도에 대한 요구는 더욱 높아지고 있습니다. vLLM은 LLM 추론을 위한 고성능 라이브러리로, 지속적인 최적화를 통해 이러한 요구에 부응하고 있습니다. 이번 PR(#45861)은 vLLM에서 DeepSeek-V4 모델의 특정 상황, 즉 'Eager Break' 시 발생하는 Workspace 재사용 로직을 최적화하여 전체 추론 시간(End-to-End Time To First Token, TTFT)을 약 3.9% 개선하는 것을 목표로 합니다.

본 글에서는 이 PR이 어떤 문제를 해결하려 했는지, 코드 변경사항은 무엇이며, 왜 이러한 변경이 성능 향상으로 이어지는지 상세히 분석하고, 관련 기술적 교훈을 공유하고자 합니다.

코드 변경 분석

이번 PR의 핵심 변경사항은 csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu 파일에 집중되어 있습니다. 기존에는 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 함수가 내부적으로 q_out 텐서를 동적으로 할당했지만, 이를 _out 접미사가 붙은 별도의 함수로 분리하고, 호출하는 측에서 q_out 텐서를 미리 할당하여 전달받도록 변경했습니다.

1. `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out` 함수 도입

가장 중요한 변경은 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 함수의 시그니처 변경과 _out 버전의 도입입니다.

Before:

-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
     torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
+    torch::stable::Tensor& q_out,                // [N, q_head_padded, 512]
     torch::stable::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
     torch::stable::Tensor const& slot_mapping,   // [N] int64
     torch::stable::Tensor const& position_ids,   // [N] int64
@@ -970,8 +971,16 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
   STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
                   "q_in and kv dtype must match");
+  STD_TORCH_CHECK(q_out.device() == q_in.device() && q_out.is_contiguous(),
+                  "q_out must be contiguous and on the same device as q_in");
+  STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(),
+                  "q_out dtype must match q_in");
   STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
                   "q_head_padded must be >= q_in.size(1) (num_heads_q)");
+  STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) &&
+                      q_out.size(1) == q_head_padded &&
+                      q_out.size(2) == q_in.size(2),
+                  "q_out shape [N, q_head_padded, 512]");
   STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
                   "k_cache must be uint8");
   STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
@@ -999,11 +1008,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
       q_in.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
 
-  // Allocate the padded q output.  The kernel writes every element (live
-  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
-
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
@@ -1020,6 +1024,20 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 num_heads_q_padded, cache_block_size_i, kv_block_stride,
                 stream);
       });
+}
+
+torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
+    torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& position_ids,
+    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
+    double eps, int64_t cache_block_size) {
+  auto q_out = torch::stable::new_empty(
+      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
+  fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
+      q_in, kv, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache,
+      q_head_padded, eps, cache_block_size);
   return q_out;
 }

변경 전에는 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 함수 내부에서 q_out 텐서를 torch::stable::new_empty를 사용하여 동적으로 할당했습니다. 하지만 변경 후에는 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 함수가 q_out을 인자로 받아 직접 수정(in-place)하는 방식으로 변경되었고, 기존의 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 함수는 이 _out 함수를 호출하면서 q_out을 동적으로 할당하는 역할을 그대로 수행합니다.

이 변경의 핵심은 호출하는 측에서 q_out 텐서의 할당을 관리하도록 위임한 것입니다. 이는 vLLM의 메모리 관리 전략, 특히 Eager Break 시점에 Workspace 재사용을 최적화하기 위한 조치입니다. Eager Break는 특정 조건에서 연산 중간 결과를 즉시 반환하고 다음 연산을 준비하는 메커니즘인데, 이 과정에서 불필요한 메모리 할당 및 해제를 줄이는 것이 중요합니다.

2. `torch_bindings.cpp` 업데이트

C++ 바인딩 파일인 torch_bindings.cpp에서도 새로운 _out 함수에 대한 정의가 추가되었습니다.

Before (Implicitly, as _out version didn't exist):

-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-      "Tensor q_in, Tensor kv, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
+

After:

+  ops.def(
+      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
+      "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
+      "int q_head_padded, float eps, int cache_block_size) -> ()");
+
   // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
   // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.

새로운 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 함수는 반환 값이 Tensor가 아닌 () (void)로 명시되어, 인자로 받은 q_out 텐서를 직접 수정함을 나타냅니다. 이는 C++ 레벨에서 함수의 인터페이스를 명확히 하고, Python/PyTorch 코드에서 이 함수를 호출할 때 기대하는 동작을 정의합니다.

3. 테스트 코드 변경

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py 파일에서도 테스트 방식이 변경되었습니다.

Before:

-    q_out = _call_fused(
-        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+    q_out = torch.empty(num_tokens, padded_heads, HEAD_DIM, dtype=dtype, device=device)
+    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
+        q,
+        kv,
+        q_out,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        padded_heads,
+        eps,
+        bs,
     )

테스트 함수 _call_fused 내부에서 q_out을 할당하고 커널을 호출하는 대신, 테스트 함수 외부에서 torch.empty로 q_out을 미리 할당한 후, 새로 도입된 torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 함수를 직접 호출하도록 변경되었습니다. 이는 실제 코드 변경을 반영하고, _out 버전 함수가 어떻게 사용되는지를 보여줍니다.

또한, tests/kernels/test_fused_indexer_q_rope_quant.py에서도 output_buffers 인자를 통해 미리 할당된 버퍼를 전달하는 테스트 로직이 추가되었는데, 이는 fused_indexer_q_rope_quant 함수에서도 유사한 'output buffer' 관리 패턴이 적용되었거나 적용될 수 있음을 시사합니다. 리뷰어의 피드백(

참고 자료

https://pytorch.org/docs/stable/generated/torch.empty.html

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

PR Analysis 의 다른글

이전글 [sglang] [성능 최적화] 불필요한 Tree Mask Fill 제거를 통한 Speculative Decoding 가속화
현재글 : [vllm] vLLM DeepSeek-V4 성능 최적화: Eager Break 시 Workspace 재사용 개선
다음글 [onnxruntime] ONNX Runtime WebGPU EP: 디바이스 없는 오프라인 컴파일 지원

[vllm] vLLM DeepSeek-V4 성능 최적화: Eager Break 시 Workspace 재사용 개선

들어가며

코드 변경 분석

1. `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out` 함수 도입

2. `torch_bindings.cpp` 업데이트

3. 테스트 코드 변경

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

들어가며

코드 변경 분석

1. fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 함수 도입

2. torch_bindings.cpp 업데이트

3. 테스트 코드 변경

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

1. `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out` 함수 도입

2. `torch_bindings.cpp` 업데이트