본문으로 건너뛰기

[vllm] vLLM DeepSeek v4 Fused Indexer Q 양자화 커널 최적화: CuteDSL을 활용한 성능 향상

PR 링크: vllm-project/vllm#41428 상태: Merged | 변경: +0 / -0

들어가며

대규모 언어 모델(LLM)의 추론 성능은 GPU 메모리 대역폭에 크게 좌우됩니다. 특히 vLLM과 같은 고성능 추론 엔진에서는 어텐션 메커니즘의 핵심 부분인 쿼리(Q) 처리 및 양자화(Quantization) 과정에서 발생하는 메모리 접근 패턴이 전체 성능에 지대한 영향을 미칩니다. 본 PR은 vLLM의 DeepSeek v4 모델을 위한 _fused_indexer_q_rope_mxfp4_kernel 커널의 성능 병목을 해결하고, GPU 메모리 대역폭 활용률을 극대화하기 위한 최적화 방안을 제시합니다.

기존 Triton 커널은 256비트 로드(load)를 효율적으로 활용하지 못하여 메모리 바운드(memory-bound) 작업에서 성능 한계가 있었습니다. 이 PR은 NVIDIA의 CuteDSL(CUDA Template Library for Expressing Tensor Operations)을 사용하여 커널을 재작성함으로써, 더 넓은 메모리 접근 패턴(256비트 로드)을 가능하게 하고, 결과적으로 DeepSeek v4 모델의 추론 처리량을 크게 향상시키는 것을 목표로 합니다.

코드 분석: CuteDSL로의 전환과 최적화

이번 PR의 핵심은 기존 Triton 기반의 _fused_indexer_q_rope_mxfp4_kernel을 CuteDSL 기반의 fused_indexer_q_rope_quant_mxfp4_cutedsl로 대체하는 것입니다. 이 변경은 단순히 커널 구현 언어를 바꾸는 것을 넘어, GPU 하드웨어의 메모리 접근 방식을 최적화하여 성능을 끌어올리는 중요한 개선입니다.

tests/kernels/test_fused_indexer_q_rope_quant.py

테스트 코드에서는 새로운 num_tokens 값 1023이 추가되었습니다. 이는 다양한 시퀀스 길이에 대한 커널의 견고성과 성능을 검증하기 위함입니다. 특히 1023과 같이 2의 거듭제곱이 아닌 값은 엣지 케이스 테스트에 유용합니다.

--- a/tests/kernels/test_fused_indexer_q_rope_quant.py
+++ b/tests/kernels/test_fused_indexer_q_rope_quant.py
@@ -122,7 +122,7 @@ def _reference(
         return q_fp8, weights_out
 
 
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023])
 @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
 @pytest.mark.parametrize("use_fp4", [False, True])
 @torch.inference_mode()

vllm/utils/import_utils.py

has_cutedsl 함수가 추가되어 cutlass 패키지의 가용성을 확인합니다. 이는 CuteDSL 커널이 특정 환경에서만 사용될 수 있도록 하는 유연성을 제공하며, cutlass가 없는 환경에서는 기존 Triton 커널로 폴백(fallback)할 수 있게 합니다.

--- a/vllm/utils/import_utils.py
+++ b/vllm/utils/import_utils.py
@@ -469,3 +469,8 @@ def has_mori() -> bool:
 def has_fbgemm_gpu() -> bool:
     """Whether the optional `fbgemm_gpu` package is available."""
     return _has_module("fbgemm_gpu")
+
+
+def has_cutedsl() -> bool:
+    """Whether the optional `cutelass` package is available."""
+    return _has_module("cutlass")

vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py

이 파일은 fused_indexer_q_rope_quant 함수의 핵심 로직을 변경합니다. has_cutedsl()을 통해 CuteDSL 커널의 사용 가능 여부를 확인하고, 가능하다면 새로운 fused_indexer_q_rope_quant_mxfp4_cutedsl을 호출하고, 그렇지 않다면 기존 _fused_indexer_q_rope_mxfp4_kernel Triton 커널을 사용하도록 변경되었습니다. 이는 호환성을 유지하면서 점진적인 전환을 가능하게 합니다.

또한, CuteDSL 커널 호출 시 인자 전달 방식이 간소화되었습니다. 기존 Triton 커널은 stride와 같은 세부적인 메모리 레이아웃 정보를 명시적으로 전달해야 했지만, CuteDSL은 cute.Tensor 객체를 통해 이러한 정보를 내부적으로 관리하므로 더 깔끔한 인터페이스를 제공합니다.

--- a/vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
+++ b/vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
@@ -342,30 +344,49 @@ def fused_indexer_q_rope_quant(
             dtype=torch.uint8,
             device=index_q.device,
         )
-        _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
-            positions,
-            index_q,
-            index_q.stride(0),
-            index_q.stride(1),
-            index_q_cos_sin_cache,
-            index_q_cos_sin_cache.stride(0),
-            index_q_cos_sin_cache.shape[-1] // 2,
-            index_q_packed,
-            index_q_packed.stride(0),
-            index_q_packed.stride(1),
-            index_q_scale,
-            index_q_scale.stride(0),
-            index_q_scale.stride(1),
-            index_q_head_dim,
-            MXFP4_BLOCK_SIZE,
-            index_weights,
-            index_weights.stride(0),
-            index_weights_softmax_scale,
-            index_weights_head_scale,
-            index_weights_out,
-            index_weights_out.stride(0),
-            num_warps=1,  # TODO: Tune this
-        )
+        if has_cutedsl():
+            # lazily import, otherwise some tests fail due to CUDA driver init failure.
+            from .fused_indexer_q_cutedsl import (
+                fused_indexer_q_rope_quant_mxfp4_cutedsl,
+            )
+
+            fused_indexer_q_rope_quant_mxfp4_cutedsl(
+                positions,
+                index_q,
+                index_q_cos_sin_cache,
+                index_weights,
+                index_weights_softmax_scale,
+                index_weights_head_scale,
+                index_q_packed,
+                index_q_scale,
+                index_weights_out,
+            )
+        else:
+            _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
+                positions,
+                index_q,
+                index_q.stride(0),
+                index_q.stride(1),
+                index_q_cos_sin_cache,
+                index_q_cos_sin_cache.stride(0),
+                index_q_cos_sin_cache.shape[-1] // 2,
+                index_q_packed,
+                index_q_packed.stride(0),
+                index_q_packed.stride(1),
+                index_q_scale,
+                index_q_scale.stride(0),
+                index_q_scale.stride(1),
+                index_q_head_dim,
+                MXFP4_BLOCK_SIZE,
+                index_weights,
+                index_weights.stride(0),
+                index_weights_softmax_scale,
+                index_weights_head_scale,
+                index_weights_out,
+                index_weights_out.stride(0),
+                num_warps=1,  # TODO: Tune this
+            )

vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q_cutedsl.py (신규 파일)

이 파일은 새로운 CuteDSL 커널 IndexerQMxFp4Kernel을 정의합니다. 주요 특징은 다음과 같습니다:

  1. CuteDSL 유틸리티 함수: _recast_val, _fp32x2_to_bf16x2, _bf16x2_to_fp32, _bf16x2_abs, _bf16x2_max, _fp32x8_to_fp4x8와 같은 저수준 CUDA inline assembly를 활용하는 dsl_user_op들이 정의되어 있습니다. 이들은 데이터 타입 변환 및 특정 연산을 GPU 하드웨어에 최적화된 방식으로 수행합니다. 특히 _fp32x8_to_fp4x8는 8개의 FP32 값을 4개의 E2M1x2 바이트로 패킹하여 하나의 32비트 레지스터에 저장하는데, 이는 256비트 로드를 효율적으로 활용하기 위한 핵심적인 부분입니다.

  2. IndexerQMxFp4Kernel 클래스: 이 클래스는 커널의 파라미터(head_dim, rope_dim 등)를 캡슐화하고, __call__ 메서드를 통해 커널을 런칭합니다. coarsen 파라미터는 num_tokens 값에 따라 동적으로 선택되는 휴리스틱(heuristic)으로, 작은 시퀀스에는 1, 큰 시퀀스에는 4를 사용하여 RoPE(Rotary Positional Embedding) 로드 비용을 상각(amortize)합니다. 이는 다양한 입력 크기에 대한 성능 최적화를 가능하게 합니다.

  3. kernel 메서드: 실제 GPU 커널 로직을 포함합니다. block_id, tid를 사용하여 스레드 인덱싱을 수행하고, global_subwarp_id, token_id, head_tile_id를 계산하여 각 스레드가 처리할 데이터 범위를 결정합니다. 특히 subwarp_size = head_dim // 16은 32B 로드(16 BF16 요소)를 활용하기 위한 설계로, 8개의 스레드가 하나의 head_dim을 처리하도록 합니다. 이는 256비트 로드를 통해 메모리 대역폭을 최대한 활용하려는 의도입니다.

# vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q_cutedsl.py
# ... (생략)

@dsl_user_op
def _fp32x8_to_fp4x8(
    vals: cute.Tensor,
    offset: cutlass.Constexpr[int],
    *,
    loc=None,
    ip=None,
) -> Uint32:
    # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32.
    assert vals.element_type is Float32
    out = llvm.inline_asm(
        T.i32(),
        [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)],
        "{\n\t"
        ".reg .b8 x0, x1, x2, x3;\n\t"
        "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t"
        "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t"
        "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t"
        "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t"
        "mov.b32 $0, {x0, x1, x2, x3};\n\t"
        "}\n",
        "=r,f,f,f,f,f,f,f,f",
        has_side_effects=False,
        is_align_stack=False,
    )
    return Uint32(out)


class IndexerQMxFp4Kernel:
    """Eight-thread subwarps process one ``(token, head)`` row."""

    def __init__(
        self,
        head_dim: int = 128,
        rope_dim: int = 64,
        num_heads: int = 64,
        cos_sin_dtype: type[cutlass.Numeric] = Float32,
        coarsen: int = 4,
    ):
        self.head_dim = head_dim
        self.rope_dim = rope_dim
        self.nope_dim = head_dim - rope_dim
        self.num_heads = num_heads
        self.cos_sin_dtype = cos_sin_dtype

        # process multiple heads at the same time to armotize RoPE load costs
        assert num_heads % coarsen == 0
        self.coarsen = coarsen

        # later we will use 32B load = 16 BF16 elems
        # thus, head_dim=128 requires 8 threads to handle.
        # let's call subwarp = 8 threads.
        self.subwarp_size = head_dim // 16
        self.tb_size = 128
        self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size

    # ... (생략)

    @cute.kernel
    def kernel(
        self,
        positions: cute.Tensor,
        q: cute.Tensor,
        cos_sin_cache: cute.Tensor,
        weights: cute.Tensor,
        q_fp4: cute.Tensor,
        q_scale: cute.Tensor,
        weights_out: cute.Tensor,
        scale: Float32,
    ):
        block_id, _, _ = cute.arch.block_idx()
        tid, _, _ = cute.arch.thread_idx()

        num_token_heads = q.shape[0] * self.num_heads
        global_tid = block_id * self.tb_size + tid

        global_subwarp_id = global_tid // self.subwarp_size
        sublane = tid % self.subwarp_size

        token_id = global_subwarp_id // (self.num_heads // self.coarsen)
        head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen)
        head_start = head_tile_id * self.coarsen

        # NOTE: token_id may

왜 이게 좋은 최적화/개선인가?

이 PR은 GPU 커널 최적화의 모범 사례를 보여줍니다. 핵심적인 개선점과 그 이유는 다음과 같습니다.

  1. 256비트 로드 활용: 기존 Triton 커널은 128비트 로드에 최적화되어 있었을 가능성이 높습니다. 최신 GPU 아키텍처(특히 GB200과 같은 Hopper 세대)는 256비트 또는 그 이상의 메모리 트랜잭션을 효율적으로 처리할 수 있습니다. CuteDSL을 사용하여 _fp32x8_to_fp4x8와 같은 저수준 인라인 어셈블리 코드를 직접 작성함으로써, 8개의 FP32 값을 256비트 레지스터에 패킹하고 한 번에 로드/저장하는 것이 가능해집니다. 이는 메모리 대역폭 활용률을 극대화하여 데이터 전송 시간을 단축시킵니다.

  2. CuteDSL의 유연성: CuteDSL은 Triton보다 더 세밀한 GPU 하드웨어 제어를 가능하게 합니다. 이를 통해 개발자는 특정 아키텍처의 특성(예: 256비트 로드, 특정 데이터 타입의 효율적인 처리)을 최대한 활용하는 커널을 작성할 수 있습니다. Triton은 높은 생산성을 제공하지만, 때로는 최저 수준의 최적화에는 한계가 있을 수 있습니다.

  3. RoPE 로드 비용 상각: coarsen 파라미터를 통해 여러 헤드를 동시에 처리하는 전략은 RoPE 계산에 필요한 cos_sin_cache 로드 비용을 여러 헤드에 걸쳐 분산시킵니다. 이는 특히 rope_dim이 head_dim보다 작은 경우, RoPE 데이터 로드가 병목이 될 수 있는 상황에서 효율적입니다.

  4. 성능 수치: 마이크로벤치마크 결과는 이 최적화의 효과를 명확하게 보여줍니다. num_tokens가 증가함에 따라 Before 대비 After의 속도 향상(Speedup)이 최대 9.12배에 달하며, 대역폭(Bandwidth) 또한 크게 증가했습니다. 이는 커널이 메모리 바운드에서 컴퓨트 바운드로 전환되거나, 메모리 접근 패턴이 훨씬 효율적으로 개선되었음을 의미합니다.

    • T=16384일 때, Before (us): 576.05 -> After (us): 63.14 (9.12배 속도 향상)
    • T=16384일 때, Before BW (GB/s): 604.6 -> After BW (GB/s): 5549.2 (약 9.17배 대역폭 향상)

    E2E 벤치마크에서도 전체 처리량(Throughput)이 66877 tok/s에서 67910 tok/s로 소폭 상승했습니다. 마이크로벤치마크만큼 극적인 변화는 아니지만, 이는 이 커널이 전체 워크로드에서 차지하는 비중을 고려할 때 유의미한 개선입니다.

  5. 폴백 메커니즘: has_cutedsl()을 통한 조건부 로딩 및 실행은 cutlass 라이브러리가 없는 환경(예: ROCm)에서도 vLLM이 계속 작동하도록 보장합니다. 이는 소프트웨어의 안정성과 호환성을 높이는 좋은 설계 결정입니다.

일반적 교훈

  • 하드웨어 친화적 프로그래밍: GPU 커널 개발 시 대상 하드웨어의 메모리 아키텍처(캐시, 레지스터, 메모리 트랜잭션 크기)를 이해하고 이를 최대한 활용하는 것이 중요합니다. 256비트 로드와 같은 넓은 메모리 접근은 대역폭 활용률을 높이는 핵심 요소입니다.
  • 저수준 최적화의 필요성: Triton과 같은 고수준 DSL은 개발 생산성을 높이지만, 극한의 성능이 요구되는 경우에는 CuteDSL이나 CUDA C++와 같은 저수준 언어를 통해 직접 하드웨어를 제어하는 것이 필요할 수 있습니다.
  • 점진적 개선과 폴백: 새로운 최적화 기법을 도입할 때는 기존 시스템과의 호환성을 유지하고, 새로운 기술이 지원되지 않는 환경을 위한 폴백 메커니즘을 마련하는 것이 중요합니다.
  • 벤치마크의 중요성: 마이크로벤치마크와 E2E 벤치마크를 모두 수행하여 최적화의 실제 효과를 정량적으로 측정하고, 병목 지점을 정확히 파악하는 것이 중요합니다.

결론

이 PR은 vLLM의 DeepSeek v4 Indexer Q 양자화 커널을 CuteDSL로 성공적으로 전환하여 GPU 메모리 대역폭 활용률을 극대화하고, 결과적으로 최대 9배 이상의 성능 향상을 달성했습니다. 이는 LLM 추론 엔진의 성능을 한 단계 끌어올리는 중요한 기여이며, 저수준 GPU 커널 최적화의 중요성을 다시 한번 상기시켜 줍니다. 앞으로 vLLM에서 더 많은 CuteDSL 커널이 도입되어 성능이 더욱 향상될 것을 기대합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글