커스텀 커널¶

개요¶

vLLM RBLN은 성능 향상을 위해 커스텀 커널을 지원합니다. 현재 RBLN NPU는 Triton을 사용하여 사용자 정의 커널을 작성할 수 있습니다.

설치¶

Triton을 사용하기 위해 별도의 설치는 필요하지 않습니다. Triton을 사용하여 커스톰 커널을 만드는 기능은 RBLN 컴파일러(rebel-compiler)에 내장되어 있습니다. 툴체인만 설치하면 충분합니다.

설치 방법은 설치 가이드를 참조하세요.

컴파일 과정¶

커스텀 커널을 작성하고 실행하면, Triton은 커널을 RBLN IR 아티팩트로 컴파일합니다. 그런 다음 RBLN 컴파일러가 해당 아티팩트를 타겟 바이너리로 컴파일하여 디바이스에서 실행합니다.

컴파일 파이프라인은 다음과 같이 구분됩니다:

타겟 독립적인 분석과 최적화를 수행하는 프론트엔드.
Triton IR을 RBLN 전용 형태로 낮추고 rebel-compiler가 사용할 수 있는 아티팩트를 생성하는 RBLN 백엔드.

컴파일 과정은 다음과 같습니다:

프론트엔드 (타겟 독립적)
- 커스텀 커널을 파싱하고 의미를 검증합니다 (컴파일러가 요구하는 형태, 포인터 사용, 제약 조건).
- 커널을 Triton Tensor IR(TTIR)로 변환합니다. 컴파일을 위한 타겟 독립적인 IR입니다.
- 특정 타겟에 독립적인 최적화를 적용합니다 (예: 정규화 및 단순화).
RBLN 백엔드 (타겟 특정)
- RBLN NPU에서 프로그램을 실행 가능하게 만들기 위한 RBLN 특정 컴파일러 패스를 실행합니다.
- TTIR을 RBLN 컴파일러가 사용하는 RBLN IR로 변환합니다.
- RBLN 제약 조건에 맞추기 위해 하드웨어 특정 최적화를 적용합니다 (예: 레이아웃 및 스케줄링)는
- 툴체인이 타겟 바이너리로 컴파일하고 실행할 수 있는 IR을 생성합니다.

기본 커스텀 커널 작성 방법¶

이번 섹션에서는 vLLM RBLN에서 사용하고 있는 flash_attention_naive_prefill을 예제로 실제 사용법을 보여줍니다. 관련된 코드는 flash attention에서 확인하실 수 있습니다.

커스텀 커널 사용 방법¶

1. Triton으로 커스텀 커널 정의¶

커널은 @triton.jit으로 정의합니다. 커널 동작은 해당 함수 내부에 구현하시면 됩니다.

@triton.jit
def flash_attention_naive_prefill(
    query,
    key,
    ...

2. 커스텀 커널 컴파일 (`warmup`)¶

Triton으로 작성된 커널은 Triton 컴파일러에 의해 rebel-compiler에서 인식할 수 있는 RBLN IR로 변환됩니다. 이를 지원하려면 Triton에서 제공하는 warmup을 실행하십시오.

def warmup(func, *args):
    ...

3. PyTorch 연산자 등록 (`triton_op`, `register_fake`)¶

작성된 커스텀 커널을 PyTorch에서 인식하려면 PyTorch에 등록해야 합니다.

@triton_op("rbln_triton_ops::flash_attention_naive_prefill", mutates_args=())
def _(
    ...

@register_fake("rbln_triton_ops::flash_attention_naive_prefill")
def _(
    ...

4. 연산자 호출 (`torch.ops.rbln_triton_ops`)¶

실제 모델에서 사용 시에는 torch.ops.rbln_triton_ops.<user_kernel_name>로 사용합니다. 이 예제에서 <user_kernel_name>은 flash_attention_naive_prefill입니다.

    flash_attention_naive_prefill = torch.ops.rbln_triton_ops.flash_attention_naive_prefill

    attn_output = flash_attention_naive_prefill(
        .....
    )

triton 커널 작성 방법¶

커널 작성 시 유의점¶

루프에 tl.static_range 사용
- 루프는 고정된 횟수로 동작합니다. tl.static_range(...)와 tl.constexpr 매개변수를 사용하세요.
- 현재 tl.range는 지원하지 않습니다.
리덕션(tl.max, tl.sum)에 keep_dims=True 설정
- 많은 attention 커널이 리덕션된 값을 원래 rank로 다시 브로드캐스트합니다.
- keep_dims=True를 사용하여 차원을 유지하고 브로드캐스팅을 명확하게 만드세요.
이 예제에서 warmup grid는 컴파일 타임 전용
- 이 예제에서 warmup은 컴파일 및 아티팩트 생성을 위해 사용되며, 런타임 실행은 RBLN 툴체인이 처리합니다.
- grid=(1,)와 같은 간단한 플레이스홀더를 사용하세요.

지원하는 연산¶

현재 지원 연산 관련해서는 지원하는 연산을 참조하세요. 지원 범위는 차후 확대 예정입니다.

triton 표준 연산 외에 RBLN NPU 하드웨어에 특화된 추가 연산도 지원합니다. 자세한 내용은 특화 연산을 참조하세요.

vLLM RBLN 커널 예제¶

현재 vLLM RBLN에서는 댜음 커스텀 커널을 지원하고 있습니다.

참조¶

Triton 공식문서