vLLM IR: 함수형 중간 표현

vLLM IR: 함수형 중간 표현 (Functional Intermediate Representation)

vLLM IR은 저수준 torch 연산과 RMSNorm·양자화 연산 같은 vLLM 레이어 사이의 간극을 메우는 함수형 중간 표현(IR)입니다. 연산의 의미(semantics)를 구현·디스패치와 분리함으로써 컴파일과 커널 등록·디스패치를 동시에 단순화합니다. torch FX 표현 위의 다이얼렉트(dialect)로 동작해 "일반" torch 연산·커스텀 torch 연산/커널과 완전히 상호운용되며, 기존 CustomOp 방식에서 부분적으로(조각별) 이전할 수도 있습니다.

출처: 문서

본문

동기 (Motivation)

vLLM IR은 저수준 torch 연산과 RMSNorm·양자화 연산 같은 vLLM 레이어 사이의 간극을 메우는 함수형 IR입니다. 연산의 의미론을 구현·디스패치와 분리함으로써 컴파일과 커널 등록·디스패치를 동시에 단순화합니다. torch FX 표현 위의 다이얼렉트로 동작해 "일반" torch 연산·커스텀 torch 연산/커널과 완전히 상호운용되며, 기존 CustomOp 방식에서 조각별 이전도 가능합니다.

주요 설계 원칙:

• Eager-컴파일 일관성: eager 모드와 컴파일 모드에서 (사소한 수치 차이를 제외하고) 동일한 동작.
• 단순하고 투명하지만 강력한 커널 선택: 가시성과 제어가 좋아 디버깅이 쉽다.
• 관례가 설정보다 우선(Convention over configuration): 연산·구현을 등록하는 데 필요한 보일러플레이트가 거의 0.
• 확장성: 연산과 구현을 in-tree든 out-of-tree든 어디서나 등록 가능.
• 상호운용성: "일반" torch 연산·커스텀 torch 연산/커널과 완전 호환 — 개발자 마찰을 줄이고 조각별 이전을 가능하게 함.

깔끔한 의미론/구현 분리는 통합적이고 확장 가능한 디스패치 메커니즘을 가능하게 해, 플랫폼당 여러 커널과 강력한 커널 선택을 허용합니다. 또한 이 분리는 테스트·벤치마킹을 더 깔끔하게 만들어 레거시 접근법에서 흔한 보일러플레이트를 크게 줄입니다.

커널 선택을 컴파일 과정 후반까지 미룸으로써 컴파일러는 더 고수준의 표현 위에서 동작할 수 있으며, 주요 이점은 다음과 같습니다.

• 퓨전/변환 패스의 패턴 매칭이 연산당 단일 단순 패턴만 필요.
• OOT(out-of-tree) 컴파일러 백엔드가 고수준 표현에서 내려올(lower) 수 있음(진행 중).
• 컴파일러가 가용 구현들에 대해 자동 튜닝(autotune) 가능(향후 기능).

빠른 개요 (Quick Overview)

IR 연산 선언 (Declaring an IR Operation)

IR 연산은 네이티브 PyTorch 구현과 함께 @register_op 데코레이터로 선언되며, 이 구현이 연산의 의미론을 정의합니다.

# vllm/ir/ops/layernorm.py
from torch import Tensor
from vllm.ir import register_op

@register_op
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
    """Weighted root-mean-square layer normalization"""
    orig_dtype = x.dtype
    x = x.to(torch.float32)
    x_var = x if variance_size is None else x[..., :variance_size]
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
    x = x * torch.rsqrt(variance + epsilon)
    x = x.to(orig_dtype)
    if weight is not None:
        x = x * weight
    return x

네이티브 구현은 세 가지 목적을 제공합니다.

• 의미론 정의: 형상·스트라이드를 포함한 연산의 정확한 의미론을 명시.
• 기본 구현: 더 나은 구현이 없을 때 사용.
• 테스트용 참조: 다른 구현들은 이 의미론과 일치해야 함.

구현 등록 (Registering Implementations)

커널 구현은 IR 연산 객체의 register_impl 데코레이터로 등록합니다.

# vllm/kernels/vllm_c.py
from vllm import ir

rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None

@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike())
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
    output = torch.empty_like(x)
    torch.ops._C.rms_norm(output, x, weight, epsilon)
    return output

구현은 다음을 지정할 수 있습니다.

• supported: 이 구현을 쓸 수 있는지 나타내는 정적 불리언.
• supports_args: 특정 인자를 지원하는지 확인하는 함수.
• inplace: 이 구현이 출력에 입력 메모리를 재사용하는지 여부.

모델에서 IR 연산 사용하기 (Using IR Operations in Models)

IR 연산은 모델 코드에서 직접 import·호출합니다.

# vllm/model_executor/layers/layernorm.py
from vllm import ir

class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x: Tensor, residual: Tensor | None = None):
        if residual is None:
            return ir.ops.rms_norm(x, self.weight, self.variance_epsilon)

        # Use maybe_inplace overload to allow implementation to reuse input memory for outputs
        # (using x or residual after this call is undefined behavior)
        return ir.ops.fused_add_rms_norm.maybe_inplace(
            x, residual, self.weight, self.variance_epsilon
        )

커널 선택 설정하기 (Configuring Kernel Selection)

커널 선택은 설정의 우선순위 목록으로 제어됩니다. 우선순위 목록은 구현이 고려되는 순서를 지정하며, 첫 번째로 지원되는 구현이 선택됩니다. 여기에는 정적 지원 검사(supported=...)와 동적 인자 지원 검사(supports_args=...)가 포함됩니다.

커맨드라인 설정 (Command Line Configuration)

--ir-op-priority.<op_name>=<provider1>,<provider2>,...를 사용합니다.

# CUDA: Use vllm_c implementation for rms_norm
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=vllm_c

# ROCm: Try aiter first, fall back to vllm_c, then native
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=aiter,vllm_c,native

# Configure multiple operations
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=vllm_c \
  --ir-op-priority.fused_add_rms_norm=vllm_c

Python 설정 (Python Configuration)

from vllm import LLM
from vllm.config import VllmConfig, KernelConfig

llm = LLM(
    model="meta-llama/Llama-3.2-1B",
    vllm_config=VllmConfig(
        kernel_config=KernelConfig(
            ir_op_priority={
                "rms_norm": ["vllm_c", "native"],
                "fused_add_rms_norm": ["vllm_c", "native"],
            }
        )
    )
)

플랫폼 기본값 (Platform Defaults)

각 플랫폼은 자동 적용되는 기본 우선순위 목록을 제공합니다.

# CUDA/XPU/ROCm platform defaults (when compiling with Inductor)
{
  "rms_norm": ["native"],  # Native torch is default
  "fused_add_rms_norm": ["native"],
}

# CUDA platform defaults (eager or Dynamo-only)
{
  "rms_norm": ["vllm_c", "native"],
  "fused_add_rms_norm": ["vllm_c", "native"],
}

# ROCm platform defaults (future - currently same as CUDA)
{
    "rms_norm": ["aiter", "vllm_c", "native"],
    "fused_add_rms_norm": ["aiter", "vllm_c", "native"],
}

# XPU platform defaults (eager or Dynamo-only)
{
    "rms_norm": ["xpu_kernels", "native"],
    "fused_add_rms_norm": ["xpu_kernels", "native"],
}

사용자가 지정한 우선순위는 플랫폼 기본값 앞에 붙습니다. 그래서 순서가 벗어난 구현만 지정하면 되고, 나머지 구현은 자동으로 뒤에 추가됩니다.

컴파일 파이프라인 (Compilation Pipeline)

vLLM IR은 torch.compile 기반 컴파일 과정을 크게 커스터마이즈해, 커스텀 컴파일 패스가 고수준 IR 위에서 동작하면서도 최종적으로는 효율적인 저수준 코드를 산출하게 합니다. 컴파일 파이프라인은 여러 단계로 구성됩니다.

1. Dynamo 트레이싱 (Dynamo Tracing)

torch.compile이 모델의 forward pass를 트레이싱하면 vLLM IR 연산은 vllm_ir torch 라이브러리의 커스텀 연산으로 나타납니다. 이 연산은 Dynamo에게 불투명opaque)하므로, 분해 없이 FX 그래프에 직접 나타납니다.

# Python code (epsilon=1e-5)
x1 = ir.ops.rms_norm(x, weight, epsilon)
x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon)

# FX graph after Dynamo tracing
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None
x2 = out[0]
residual_out = out[1]

2. AOTAutograd와 기능화 (AOTAutograd and Functionalization)

AOTAutograd는 그래프를 기능화(functionalize)해 변경(mutating) 연산을 함수형 등가물로 변환합니다. maybe_inplace 오버로드를 가진 vLLM IR 연산은 AOTAutograd 전에 pre-grad 커스텀 패스 훅을 사용해 함수형 default 오버로드로 변환합니다.

# After functionalization
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None
x2 = out[0]
residual_out = out[1]

이 패스는 또한 어떤 입력이 "기증(donated)"되었는지(maybe_inplace에 전달됨) 추적해 이 정보를 clone 제거에 쓸 수 있도록 vLLM의 PassContext에 저장합니다.

3. IR 퓨전·변환 패스 (IR Fusion and Transformation Passes)

기능화 후, 커스텀 vLLM 패스는 고수준 IR 연산을 담은 함수형 FX 그래프 위에서 동작합니다. 이 패스들은 퓨전, 시퀀스 병렬화를 위한 연산 분배, 기타 변환을 수행할 수 있습니다.

# Example: Sequence Parallelism (see SequenceParallelismPass)
# Before SP pass

all_reduce = torch.ops.vllm.all_reduce(x, "tp:0")
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)

# after SP pass
reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0")
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
all_gather = torch.ops.vllm.all_gather(x, "tp:0")

퓨전 패스는 고수준 표현의 이점을 얻습니다. 저수준 PyTorch 연산과 매칭할 필요도, 서로 다른 커널 구현을 따로 처리할 필요도, 커스텀 커널의 기능화를 다룰 필요도 없기 때문입니다.

4. IR 하향 변환 (IR Lowering)

하향 변환 패스(VllmIRLoweringPass)는 각 vLLM IR 연산을 선택된 구현으로 교체합니다. 구현은 우선순위 목록과 지원 술어(support predicate)로 선택되며, 연산 인자 대신 그래프 메타데이터의 fake 텐서를 사용합니다.

# Implementation selection, same in eager dispatch and compile lowering
def dispatch(*args) -> IrOpImpl:
  for provider in priority_list:  # e.g., ["vllm_c", "native"]
    impl = ir_op.impls[provider]
    if not impl.supported:
      continue
    if impl.supports_args and not impl.supports_args(*args):
      continue
    return impl

# make_fx uses torch.fx.symbolic_trace
impl_graph = make_fx(selected_impl.impl_fn)
# Replace IR op node with impl_graph's nodes
match.replace_by_example(selected_impl.impl_fn, node.args)

예를 들어 rms_norm을 vllm_c 구현으로 하향 변환하면:

# Before lowering (IR op)
rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5)

# After lowering (vllm_c implementation traced)
# Note: Lowering does not currently functionalize, this will likely change in the future.
empty =  torch.ops.aten.empty.memory_format(x.shape, ...)
rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5)

입력을 변경하는 구현(inplace=True)을 하향 변환할 때, 하향 패스는 함수형 의미론을 보존하기 위해 clone을 삽입합니다.

# vllm_c implementation for fused_add_rms_norm mutates its first two arguments
# Lowered with clones for safety
clone_default = torch.ops.aten.clone.default(x)
clone_default_1 = torch.ops.aten.clone.default(residual)
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5)

5. Clone 정리 (Clone Cleanup)

하향 변환 후, clone 제거 패스(UnsafeCloneEliminationPass)가 변환 중 도입된 불필요한 clone을 제거합니다. 이 패스는 maybe_inplace로 in-place 커널을 쓸 때 zero-copy 동작을 달성하는 데 필수적입니다. 패스는 다음 경우에 clone을 제거합니다.

• clone된 입력이 그래프에서 생성되었고 그래프에서 다시 사용되지 않는 경우.
• clone된 입력이 그래프 파라미터이며 donated로 표시된 경우.

# After cleanup (donated inputs, no subsequent uses)
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5)

inplace 기능화(donated 입력 추적)와 clone 정리의 조합 덕분에 컴파일러는 중복 복사를 추가하거나 메모리 사용을 늘리지 않고도 in-place 커널을 안전하게 사용할 수 있습니다.

6. Inductor 최적화와 코드 생성 (Inductor Optimization and Codegen)

IR 하향 변환·정리 후 그래프에는 표준 PyTorch 연산과 플랫폼별 커스텀 연산만 남습니다. 그러면 Inductor가 표준 코드 생성으로 이어갑니다.

• Inductor 하향 변환과 pointwise 퓨전: 요소별 연산·리덕션 등의 퓨전.
• 메모리 계획: 버퍼 할당·재사용 결정.
• 커널 생성: 퓨전된 연산을 위한 Triton 또는 C++ 코드 생성.
• 자동 튜닝: 최적 커널 설정 선택.

파이프라인 요약 (Pipeline Summary)

Model Forward Pass
    ↓
[Dynamo Tracing] → FX Graph with vllm_ir.* ops
    ↓
[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs
    ↓
[AOTAutograd] → Functionalization
    ↓
[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant)
    ↓
[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed)
    ↓
[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info
    ↓
[Inductor] → Pattern matching, fusion, memory planning, codegen
    ↓
Compiled Code

핵심 vLLM IR 개념 (Core vLLM IR Concepts)

연산 선언 (Operation Declaration)

연산은 @register_op 데코레이터로 선언되며, 이 데코레이터가 IrOp 객체를 만듭니다.

@register_op(
    name=None,           # Operation name (defaults to function name)
    activations=None,    # List of activation parameters (defaults to params starting with 'x')
    allow_inplace=False, # Whether to create a maybe_inplace overload
)
def op_name(...):
    ...

파라미터:

• activations: "활성화"로 간주되는 파라미터 이름 목록(전형적으로 maybe_inplace에 의해 소비됨). x로 시작하는 파라미터가 기본값.
• allow_inplace: 메모리 효율 실행을 위한 maybe_inplace 오버로드를 생성(아래 참조).

maybe_inplace 오버로드 (The maybe_inplace Overload)

maybe_inplace 오버로드는 LLM 추론의 메모리 효율성을 위한 핵심 기능입니다. 호출자가 연산 후 활성화 입력을 보존할 필요가 없음을 알려, in-place 구현이 입력 메모리를 출력에 재사용할 수 있게 합니다.

의미론과 사용법 (Semantics and Usage)

# Standard usage: inputs are preserved
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
# x and residual are unchanged, out and res_out are new tensors

# maybe_inplace: inputs may be modified
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
# x and residual may be modified (undefined behavior to use them after this)
# out and res_out may alias x and residual

maybe_inplace에 전달한 뒤 활성화 입력을 사용하는 것은 정의되지 않은 동작(undefined behavior)입니다.

# WRONG: Using x after donating it
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
result = out + x  # ERROR: x was donated!

입력을 보존해야 한다면 default 오버로드를 쓰거나 직접 clone하세요.

# Option 1: Use default overload
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
result = out + x  # OK: x is preserved

# Option 2: Clone before maybe_inplace
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon)
result = out + x  # OK: x is preserved, clone was donated

컴파일 동작 (Compilation Behavior)

컴파일 동안 inplace 기능화 패스는 donated 입력이 다시 사용되지 않는지 검증하고 maybe_inplace를 함수형 default 오버로드로 변환합니다.

# Inplace functionalization pass (pre-grad)
for node in graph.nodes:
    if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace:
        # Check that activation inputs aren't used after this node
        for activation_arg in activation_inputs:
            for user in activation_arg.users:
                if user appears after node:
                    raise ValueError(f"Input {activation_arg} donated but used again")

        # Convert to default overload
        node.target = torch.ops.vllm_ir.fused_add_rms_norm.default

        # Track donated graph inputs for later clone elimination
        for i, arg in enumerate(node.args):
            if arg.op == "placeholder" and i in activation_indices:
                pass_context.donated_input_ids.add(node_to_idx[arg])

donated 입력 정보는 이후 clone 정리 패스가 in-place 커널을 하향 변환할 때 불필요한 복사를 제거하는 데 사용됩니다.

Eager 모드 동작 (Eager Mode Behavior)

eager 모드(torch.compile 없음)에서 maybe_inplace는 IR 연산이 in-place 구현으로 직접 디스패치하게 해 메모리 효율성을 극대화합니다.

# Eager dispatch logic for maybe_inplace
impl: IrOpImpl = ir_op.dispatch(*args)
return impl.impl_fn(*args)

# Eager dispatch logic for default:
impl: IrOpImpl = ir_op.dispatch(*args)
if impl.inplace:
  args = [
    arg.clone() if i in ir_op.activations else arg
    for i, arg in enumerate(args)
  ]
return impl.impl_fn(*args)

모델 코드의 maybe_inplace와 in-place 커널 구현의 조합은 eager와 컴파일 모드 모두에서 최적의 메모리 효율을 제공하며, 두 경우 모두 의미론은 동일합니다.

메모리 절약 예시 (Memory Savings Example)

잔차 연결이 있는 트랜스포머 레이어를 생각해 봅시다.

# Without maybe_inplace (2 allocations per layer)
hidden_states = self.attention(input)
normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps)
# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new)

# With maybe_inplace (0 allocations per layer when using in-place kernel)
hidden_states = self.attention(input)
normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps)
# Memory: normed (reuses hidden_states), residual (reuses input)

구현 등록 (Implementation Registration)

구현은 register_impl 메서드로 등록합니다.

@ir.ops.op_name.register_impl(
    provider="provider_name",  # Unique identifier (e.g., "vllm_c", "aiter", "triton")
    supported=True,            # Static availability check
    supports_args=None,        # Dynamic argument support check
)
def impl_fn(...):
    ...

Provider 명명 규칙:

• native: 네이티브 torch 구현(@register_op으로 선언됨)을 위해 예약.
• vllm_c: torch.ops._C를 통한 C++/CUDA 커널.
• aiter: AMD AITER 라이브러리.
• xpu_kernels: vllm-xpu-kernels에 구현된 SYCL/SYCLTLA 커널.
• triton_*: Triton 커널.
• 기타 구현을 위한 플랫폼/라이브러리 이름.

지원 검사:

• supported: 정적 불리언. import 시점에 한 번 검사(예: HAS_TRITON, is_cuda_alike()).
• supports_args: 인자 호환성을 검사하는 함수 (*args, **kwargs) -> bool. 컴파일 중에는 fake 텐서로 호출되어 비용 0 검사.
• eager 모드 디스패치 중에는 실텐서(real tensor)로 호출.
• 배치 크기를 검사하거나 값에 근거한 가드를 추가해서는 안 됨.

지원 술어 예시:

def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None):
    # Check dtype (OK: doesn't depend on batch size)
    if x.dtype not in [torch.float16, torch.bfloat16]:
        return False
    # Check optional parameter (OK: static check)
    if variance_size is not None:
        return False
    return True

@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports)
def rms_norm(...):
    ...

배치 불변(batch-invariant) 커널은 VLLM_BATCH_INVARIANT=1이 설정되면 자동으로 선택됩니다.

Eager 모드 vs 컴파일 모드 (Eager Mode vs Compile Mode)

vLLM IR 연산은 eager 모드와 컴파일 모드에서 동일하게 동작합니다.

Eager 모드:

• 우선순위 목록에 따라 구현으로 직접 디스패치.
• 실텐서 인자로 지원 검사.
• 최소 오버헤드(필요하면 더 최적화 가능).

컴파일 모드:

• IR 연산이 FX 그래프에 torch.ops.vllm_ir.* 커스텀 연산으로 나타남.
• 하향 변환이 fake 텐서로 구현을 선택.
• Inductor 최적화와 완전 통합.

이 일관성이 가능하게 하는 것들:

• eager 모드에서 자신 있게 프로토타이핑.
• 컴파일을 꺼서 디버깅.
• eager에서 컴파일 실행으로 점진적 이전.

기타 주제 (Other Topics)

아웃오브트리 구현 (Out-of-Tree Implementations)

외부 플랫폼은 vLLM을 수정하지 않고 구현을 등록할 수 있습니다.

# In external package
from vllm import ir

@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform())
def rms_norm(x, weight, epsilon, variance_size=None):
    return my_platform.rms_norm(x, weight, epsilon)

그런 다음 우선순위를 설정해 자신의 구현을 사용합니다.

class MyPlatform(Platform):
  def get_default_ir_op_priority(self):
    return IrOpPriorityConfig(rms_norm=['my_platform', 'native'])

# Users can still override priority in the same way
llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel']))

디버깅과 관측성 (Debugging and Observability)

Note: 관측성 개선 아이디어가 있다면 알려주세요(use-case 기반).

커널 선택을 보려면 디버그 로깅을 켜세요.

VLLM_LOGGING_LEVEL=DEBUG vllm serve ...

로그 내용:

• 각 연산에 대해 어떤 구현이 선택되었는지.
• 구현이 거부된 이유(지원 안 됨, 인자 미지원).
• 컴파일 캐시 히트/미스.
• IR 하향 변환 통계.

컴파일된 그래프의 선택 구현 확인:

# After compilation, inspect the lowering pass
lowering_pass = backend.lowering_pass
print(lowering_pass.selected_impls)
# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}}

CustomOp에서 이전하기 (Migration from CustomOp)

vLLM IR은 CustomOp와 공존하며 점진적으로 대체하도록 설계되었습니다.

• 연산 선언: CustomOp 클래스 PluggableLayer를 변환하고 forward_native를 @register_op 함수로 이동.
• 구현 등록: 메서드 오버라이드 대신 @ir.ops.op_name.register_impl 사용.
• 레이어 사용: self.op(...)를 ir.ops.op_name(...)으로 교체.
• 설정: --compilation-config.custom-ops를 --ir-op-priority로 이전.

이전은 한 번에 한 연산씩 점진적으로 할 수 있습니다.

함께 보기 (See Also)

• torch.compile 통합 — 일반 컴파일 인프라.
• Fusions — vLLM의 커스텀 퓨전·변환 패스.
• Custom Operations — 레거시 커스텀 연산 시스템.

더 알아보기 (Learn more)

• torch.compile 설계 문서
• --ir-op-priority.<op> / KernelConfig.ir_op_priority 설정