CustomOp

CustomOp (커스텀 연산)

CustomOp는 여러 연산의 forward 메서드를 적절한 백엔드로 디스패치하는 데 쓰는 추상 클래스예요. vLLM과 OOT(Out-Of-Tree, 저장소 밖) 플러그인이 각자의 커스텀 연산을 등록할 수 있는 메커니즘도 제공해요.

이 문서에서는 CustomOp가 vLLM에서 어떻게 동작하는지, 그리고 새 CustomOp를 어떻게 구현하는지 소개할게요.

vLLM에서 CustomOp가 동작하는 방식

CustomOp는 클래스 레벨에서 모든 커스텀 연산(즉 등록된 이름으로 색인되는 연산 클래스)의 딕셔너리 두 개를 관리해요. 하나는 vLLM용, 하나는 OOT 플러그인용이에요.

@CustomOp.register("op_name") 데코레이터로 연산 클래스를 CustomOp 시스템에 등록할 수 있어요. 등록하면 op_name과 그 클래스가 op_registry 딕셔너리에 추가돼요. OOT 연산은 @CustomOp.register_oot("op_name")로 등록할 수 있어요. 이 메커니즘은 뒤에서 자세히 다룰게요.

CustomOp가 호출되면(즉 forward() 메서드가 호출되면), 활성화돼 있다면(--compilation_config.custom_ops '["+op_name"]'로 활성화), current_platform에 따라 forward 메서드를 알맞은 백엔드로 자동 디스패치해요. 비활성화 상태라면 forward_native()만 호출해서 이 forward 메서드의 PyTorch 네이티브 구현을 써요.

  • CPU 플랫폼: forward_cpu()로 디스패치
  • CUDA 플랫폼: forward_cuda()로 디스패치
  • ROCm 플랫폼: forward_hip()로 디스패치. forward_hip()이 구현돼 있지 않으면 forward_cuda()를 폴백으로 사용
  • XPU 플랫폼: forward_xpu()로 디스패치
  • TPU 플랫폼: forward_tpu()로 디스패치
  • OOT 플랫폼: forward_oot()로 디스패치. OOT 플랫폼에서만 호출됨
  • 기본(Default): 모든 플랫폼에서 최종 폴백으로 forward_native()로 디스패치

!!! note 클래스 상속 때문에 디스패치 로직이 절대적이지 않을 수 있어요. 파생 클래스가 동작을 재정의할 수 있거든요.

게다가 vLLM은 compilation_config.custom_ops에 따라 CustomOp를 활성화할지 비활성화할지 결정해요. 구체적으로 compilation_config.custom_ops에 등록되지 않은 CustomOp(즉 기본 설정 사용)는, compilation_config.custom_ops에 all이 포함돼 있으면 활성화되고 none이 포함돼 있으면 비활성화돼요.

!!! note all과 none은 compilation_config.custom_ops에 공존할 수 없어요.

기본적으로 compilation_config.backend == "inductor"이고 compilation_config.mode != CompilationMode.NONE이면 none이 compilation_config.custom_ops에 추가되고, 그 외에는 all이 추가돼요. 즉 일부 플랫폼(예: torch.compile의 기본 백엔드를 inductor로 쓰는 곳)에서는 torch compile 모드로 실행할 때 CustomOp가 비활성화된다는 뜻이에요. 이 경우 Inductor가 비활성화된 커스텀 연산을 위한 (퓨전된) Triton 커널을 생성해요.

!!! note 멀티모달 모델의 경우 vLLM은 ViT 부분에서 더 나은 성능을 위해 장치별로 깊게 최적화된 커널을 사용하도록 일부 커스텀 연산(예: MMEncoderAttention, ApplyRotaryEmb)을 강제로 활성화해요. CustomOp의 __init__() 메서드에 enforce_enable=True 파라미터를 넘겨 객체 레벨에서 스스로 활성화를 강제할 수도 있어요.

이 `enforce_enable` 메커니즘은 멀티모달 부분용 별도 `compilation_config`를 추가한 뒤 제거될 예정이에요.

CustomOp 설정을 커스터마이즈하는 방법

vLLM은 서버를 띄울 때 --compilation_config.custom_ops '["..."]'를 직접 넘겨, 어떤 커스텀 연산을 활성화/비활성화할지 세밀하게 제어할 수 있는 옵션도 제공해요.

예를 들어:

  • --compilation_config.custom_ops '["all"]' → 모든 커스텀 연산 활성화
  • --compilation_config.custom_ops '["none"]' → 모든 커스텀 연산 비활성화
  • --compilation_config.custom_ops '["all,-op1"]' → op1을 제외한 모든 커스텀 연산 활성화(- 접두사는 "비활성화"를 뜻함)
  • --compilation_config.custom_ops '["none,+op1,+op2"]' → op1과 op2만 활성화(+ 접두사는 "활성화"를 뜻함)

vLLM에서 지원하는 CustomOp 유형

1. Attention (어텐션):

  • multi_head_latent_attention (vllm/model_executor/layers/mla.py)

2. Activation (활성화):

  • silu_and_mul, mul_and_silu, gelu_new, gelu_fast, quick_gelu, gelu_and_mul, gelu_and_mul_sparse, relu2, xielu, swigluoai_and_mul, fatrelu_and_mul (vllm/model_executor/layers/activation.py)

3. MM-Conv (멀티모달 컨볼루션):

  • conv2d, conv3d (vllm/model_executor/layers/conv.py)

4. Embedding (임베딩):

  • vocab_parallel_embedding, parallel_lm_head (vllm/model_executor/layers/vocab_parallel_embedding.py)

5. Linear (선형):

  • row_parallel_linear, column_parallel_linear, replicated_linear (vllm/model_executor/layers/linear.py)

6. Logits Processor (로짓 프로세서):

  • logits_processor (vllm/model_executor/layers/logits_processor.py)

7. Mamba:

  • mamba_mixer (vllm/model_executor/layers/mamba/mamba_mixer.py), mamba_mixer2, mixer2_gated_rms_norm (mamba_mixer2.py), short_conv (short_conv.py)

8. MoE (Mixture of Experts):

  • fused_moe (vllm/model_executor/layers/fused_moe/layer.py), modular_fused_moe (fused_moe_modular_method.py), unquantized_fused_moe (unquantized_fused_moe_method.py), transformers_fused_moe (vllm/model_executor/models/transformers/moe.py), grouped_topk (router/grouped_topk_router.py)

9. Norm (정규화):

  • rms_norm, rms_norm_gated, gemma_rms_norm (vllm/model_executor/layers/layernorm.py)

10. Quantization (양자화):

  • quant_fp8 (vllm/model_executor/layers/quantization/input_quant_fp8.py)

11. Rope (회전 위치 임베딩):

  • rotary_embedding (vllm/model_executor/layers/rotary_embedding/base.py), dual_chunk_rotary_embedding (dual_chunk_rope.py), apply_rotary_emb (common.py)

12. Encoder (인코더):

  • qwen2_decoder (vllm/model_executor/models/deepencoder2.py), mm_encoder_attn (vllm/model_executor/layers/attention/mm_encoder_attention.py), rel_pos_attention (vllm/model_executor/models/deepencoder.py)

새 CustomOp 구현 가이드라인

vLLM에서 새 CustomOp 구현하기

이 부분은 vLLM에서 새 CustomOp를 구현하는 튜토리얼이에요.

단계:

  1. CustomOp 기본 클래스에서 확장한 새 연산 클래스를 구현
  2. 이 연산 클래스에 @CustomOp.register("op_name") 데코레이터를 붙여 CustomOp 시스템에 등록
  3. 필요에 따라 다른 forward_xxx() 메서드를 구현

MMEncoderAttention을 예로 들면:

@CustomOp.register("mm_encoder_attn")
class MMEncoderAttention(CustomOp):

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float | None = None,
        num_kv_heads: int | None = None,
        prefix: str = "",
        multimodal_config: MultiModalConfig | None = None,
    ) -> None:
        super().__init__()
        # Init...

    def forward_native(self, query, key, value, cu_seqlens=None, max_seqlen=None):
        # TORCH_SDPA 구현 호출...

    def forward_cuda(self, query, key, value, cu_seqlens=None, max_seqlen=None):
        # FA(Flash Attention) 또는 TORCH_SDPA 구현 호출...

    def forward_cpu(self, query, key, value, cu_seqlens=None, max_seqlen=None):
        # TORCH_SDPA 구현 호출...

    def forward_xpu(self, query, key, value, cu_seqlens=None, max_seqlen=None):
        # FA 구현 호출...

    def forward_tpu(self, query, key, value, cu_seqlens=None, max_seqlen=None):
        # PALLAS 구현 호출...

(max_seqlen은 Flash Attention에서만 사용돼요.)

OOT 장치 플러그인에서 새 CustomOp 등록하기

현재 vLLM의 하드웨어 플러그인 메커니즘 덕분에 다양한 OOT 장치 플러그인이 등장해서 vLLM이 서로 다른 하드웨어에서 매끄럽게 동작하도록 해주고 있어요. 이 메커니즘에 대한 자세한 내용은 Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU에서도 볼 수 있어요.

이 경우 CustomOp는 하드웨어 제조사가 OOT CustomOp를 등록하고 forward_oot() 메서드를 구현하기만 하면, 런타임에 vLLM의 연산을 특정 장치용으로 깊게 최적화된 커널로 매끄럽게 교체할 수 있게 해줘요.

이제 장치 플러그인의 OOT CustomOp를 등록하는 방법을 보여줄게요.

MMEncoderAttention을 예로 들면:

  1. MMEncoderAttention에서 확장한 CustomMMEncoderAttention 클래스를 구현하고 forward_oot() 메서드를 구현
  2. 등록한 CustomMMEncoderAttention을 vLLM에 등록해서 MMEncoderAttention을 교체
from vllm.model_executor.layers.attention import MMEncoderAttention
from vllm.model_executor.custom_op import CustomOp


@CustomOp.register_oot("MMEncoderAttention")
class CustomMMEncoderAttention(MMEncoderAttention):

    def __init__(...):
        super().__init__(...)

    def forward_oot(...):
        # 장치별로 최적화된 커널 호출
        ...

이 경우 {"MMEncoderAttention": CustomMMEncoderAttention} 항목이 op_registry_oot에 추가돼요. MMEncoderAttention 연산 객체를 초기화할 때 클래스 이름(MMEncoderAttention)이 op_registry_oot의 키에 포함돼 있으면, vLLM이 등록된 클래스(CustomMMEncoderAttention)로 교체해 인스턴스화해요.

그 후 이 MMEncoderAttention 연산이 호출되면, 활성화돼 있다면 여러분의 forward_oot()가 호출돼요. 이렇게 해서 vLLM을 직접 수정하지 않고도 여러분 하드웨어에서 기대하는 성능을 얻을 수 있어요.

또한 더 나은 관리를 위해 모든 CustomOp를 한곳에 등록할 수도 있어요.

from vllm.model_executor.custom_op import CustomOp

REGISTERED_CUSTOM_OPS = {
    "CustomOP1": YourCustomOp1,
    "CustomOP2": YourCustomOp2,
    "CustomOP3": YourCustomOp3,
}

for op_name, op_cls in REGISTERED_CUSTOM_OPS.items():
    CustomOp.register_oot(_decorated_op_cls=op_cls, name=op_name)

더 알아보기