연산자 개발

연산자 개발 (Operator Development)

Ascend NPU용 커스텀 연산자(Ascend C / Triton)를 개발하고 이를 SGLang 추론 엔진에 통합하는 방법을 다룹니다.

출처: 문서

본문

개요 (Overview)

SGL-Kernel-NPU는 SGLang 프레임워크가 Ascend NPU용으로 제공하는 공식 연산자 라이브러리입니다. 두 가지 유형의 연산자 구현을 포함합니다:

Ascend C 연산자: Ascend C로 작성된 고성능 C++ 커널로, libsgl_kernel_npu.so로 컴파일되고 PyTorch의 커스텀 연산자 메커니즘(TORCH_LIBRARY_FRAGMENT)으로 등록됩니다. SGLang에서 torch.ops.npu.<op_name>()으로 호출됩니다.
Triton 연산자: Triton으로 작성되고 Ascend NPU에 맞게 적응된 Python 커널입니다. from sgl_kernel_npu.xxx import ...로 직접 호출됩니다.

SGLang이 NPU 디바이스를 감지하면 sgl_kernel_npu를 자동으로 로드하고 GPU 대응물 대신 그 연산자를 사용해 Ascend 하드웨어에서 최적화된 추론을 제공합니다.

디렉터리 구조 (Directory Structure)

이 가이드의 식별자 `sgl_kenel_npu_ops.h`, `KernalHelloworld`, `retrive_*`(예: `retrive_index`, `retrive_next_token`, `retrive_next_sibling`)는 업스트림 [sgl-kernel-npu](https://github.com/sgl-project/sgl-kernel-npu) 저장소의 철자를 그대로 따르며 일관성을 위해 원문 그대로 유지됩니다.

sgl-kernel-npu/
├── csrc/                          # Ascend C operator C++ sources
│   ├── CMakeLists.txt             # Build configuration
│   ├── pytorch_extensions.cpp     # PyTorch op registration (core integration file)
│   └── <op_name>/                 # One directory per operator
│       ├── op_host/               # Host-side code (validation, tiling, launch)
│       │   ├── <op_name>.cpp
│       │   └── tiling/            # Optional: tiling data
│       └── op_kernel/             # Device-side code (Ascend C kernel on AI Core)
│           └── <op_name>_kernel.cpp
├── include/
│   └── sgl_kenel_npu_ops.h        # C++ interface declarations
├── python/
│   └── sgl_kernel_npu/
│       └── sgl_kernel_npu/
│           ├── __init__.py        # Loads libsgl_kernel_npu.so
│           ├── attention/         # Triton attention kernels
│           ├── norm/              # Triton normalization kernels
│           ├── activation/        # Triton activation kernels
│           ├── fla/               # Triton linear attention kernels
│           ├── mamba/             # Triton Mamba kernels
│           ├── moe/               # Triton MoE kernels
│           └── sample/            # Triton speculative decoding kernels
├── tests/
│   └── python/sgl_kernel_npu/     # One test file per operator
├── build.sh                       # Build script
└── CMakeLists.txt                 # Root CMake configuration

Ascend C 연산자 개발 (Developing Ascend C Operators)

완전한 Ascend C 연산자는 두 부분으로 구성됩니다:

디바이스 부분 (Device part): NPU AI Core에서 실행되며 실제 연산을 담당하는 커널 코드. Ascend C API로 작성.
호스트 부분 (Host part): CPU에서 실행되며 파라미터 검증, 데이터 전처리, tiling, 커널 런치를 담당.

두 텐서에 대한 요소별 덧셈을 수행하는 간단한 연산자인 helloworld 예시부터 시작하는 것을 권장합니다.

1단계: 연산자 디렉터리와 파일 생성

csrc/ 아래에 op_host/ + op_kernel/ 구조를 따라 새 연산자 디렉터리를 만듭니다:

csrc/<op_name>/
├── op_host/
│   └── <op_name>.cpp
└── op_kernel/
    └── <op_name>_kernel.cpp

2단계: 디바이스 측 커널 작성 (Write the Device-side Kernel (op_kernel))

디바이스 측 코드는 AI Core에서 실행되며 Ascend C 프로그래밍 모델을 따릅니다. 핵심 구조는 Init()과 Process() 메서드와 extern "C" 엔트리 함수를 가진 클래스입니다.

helloworld를 예로 들면:

// csrc/helloworld/op_kernel/kernel_helloworld.cpp
#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2;

class KernalHelloworld {
public:
    __aicore__ inline KernalHelloworld() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
    {
        // Compute workload for current block
        this->blockLength = totalLength / AscendC::GetBlockNum();
        this->tileNum = 8;
        this->tileLength = this->blockLength / this->tileNum / BUFFER_NUM;

        // Set global memory buffers
        xGm.SetGlobalBuffer((__gm__ half *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        yGm.SetGlobalBuffer((__gm__ half *)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        zGm.SetGlobalBuffer((__gm__ half *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);

        // Initialize pipeline queues
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(half));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(half));
    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);    // Move data from Global Memory to Local Memory
            Compute(i);   // Compute on Local Memory
            CopyOut(i);   // Move results back to Global Memory
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress) { /* data copy-in... */ }
    __aicore__ inline void Compute(int32_t progress) { /* core computation... */ }
    __aicore__ inline void CopyOut(int32_t progress) { /* data copy-out... */ }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<half> xGm, yGm, zGm;
    uint32_t blockLength, tileNum, tileLength;
};

// Entry function: the compile tool auto-generates aclrtlaunch_<op_name>.h from this name
extern "C" __global__ __aicore__ void helloworld(
    GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength)
{
    KernalHelloworld op;
    op.Init(x, y, z, totalLength);
    op.Process();
}

핵심 포인트:

클래스 메서드는 __aicore__로 표시해 AI Core에서 실행됨을 나타내야 합니다.
AscendC::TPipe + AscendC::TQue를 사용해 데이터 이동과 연산을 겹치는 파이프라인을 만듭니다.
엔트리 함수는 extern "C" __global__ __aicore__로 선언해야 합니다. 컴파일 도구가 함수 이름에서 호스트 호출 가능한 런치 헤더 aclrtlaunch_<func_name>.h를 생성합니다.
단순 연산자(예: helloworld, cache_assign, lora)는 추가 workspace 메모리가 필요 없습니다. 복잡한 연산자(예: mla_preprocess, alloc_extend, build_tree)는 임시 workspace 메모리가 필요하며 CMakeLists.txt에서 별도로 컴파일됩니다.

더 심층적인 Ascend C 프로그래밍 지식은 Ascend C Kernel Development Guide를 참조하세요.

3단계: 호스트 측 코드 작성 (Write the Host-side Code (op_host))

호스트 측 코드는 PyTorch 텐서를 커널에 전달하고 런치하는 역할을 합니다. 핵심 매크로는 EXEC_KERNEL_CMD(csrc/utils/torch_helper.h에 위치)입니다.

// csrc/helloworld/op_host/helloworld.cpp
#include "defines.h"                   // Provides HOST_API macro
#include "torch_helper.h"              // Provides EXEC_KERNEL_CMD macro
#include "aclrtlaunch_helloworld.h"    // Auto-generated by compile tool

namespace sglang {
namespace npu_kernel {

HOST_API at::Tensor helloworld(const at::Tensor &x, const at::Tensor &y)
{
    // Create output tensor
    at::Tensor z = at::empty_like(x);

    // Define block count
    uint32_t blockDim = 8;

    // Compute total element count
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }

    // Launch kernel via EXEC_KERNEL_CMD macro
    EXEC_KERNEL_CMD(helloworld, blockDim, x, y, z, totalLength);
    return z;
}

} // namespace npu_kernel
} // namespace sglang

핵심 포인트:

네임스페이스는 sglang::npu_kernel이어야 합니다.
함수 시그니처는 at::Tensor <op_name>(const at::Tensor &input, ...) 패턴을 따릅니다.
출력이 여러 개인 연산자는 std::tuple<at::Tensor, at::Tensor, ...> 또는 non-const 레퍼런스 파라미터를 사용합니다.

4단계: C++ 인터페이스 선언 (include/sgl_kenel_npu_ops.h)

include/sgl_kenel_npu_ops.h에 연산자 함수 선언을 추가합니다:

// include/sgl_kenel_npu_ops.h
namespace sglang {
namespace npu_kernel {

at::Tensor helloworld(const at::Tensor &x, const at::Tensor &y);

} // namespace npu_kernel
} // namespace sglang

5단계: PyTorch 커스텀 연산자 등록 (pytorch_extensions.cpp)

csrc/pytorch_extensions.cpp에서 두 단계로 연산자를 등록합니다: 스키마 정의와 구현 바인딩.

// csrc/pytorch_extensions.cpp
namespace {

// 1. Define operator schema (used by torch.compile, etc.)
TORCH_LIBRARY_FRAGMENT(npu, m)
{
    m.def("helloworld(Tensor x, Tensor y) -> Tensor");
    // ... other operator schemas ...
}

// 2. Bind implementation for the PrivateUse1 device (i.e., NPU)
TORCH_LIBRARY_IMPL(npu, PrivateUse1, m)
{
    m.impl("helloworld", TORCH_FN(sglang::npu_kernel::helloworld));
    // ... other operator implementations ...
}

} // namespace

스키마 규칙:

네임스페이스는 npu로 고정됩니다. SGLang에서 연산자는 torch.ops.npu.<op_name>()으로 호출됩니다.
출력 텐서 파라미터는 Tensor(a!) 변형(mutating) 어노테이션을 사용합니다.
선택적 파라미터는 Tensor? 어노테이션을 사용하며, impl에서 c10::optional<T>로 처리합니다.
자세한 스키마 문법은 PyTorch Schema Reference를 참조하세요.

구현 바인딩 규칙:

디바이스 이름은 PrivateUse1(PyTorch NPU 백엔드 식별자)로 고정됩니다.
구현 함수 바인딩에는 TORCH_FN 매크로를 사용합니다.
선택적 파라미터가 있는 복잡한 연산자는 lambda 표현식으로 인자를 풀어냅니다.

6단계: 빌드 구성 업데이트 (csrc/CMakeLists.txt)

csrc/CMakeLists.txt에 새 연산자의 소스 파일을 추가합니다:

workspace가 필요 없는 연산자(단순 연산자)는 커널 소스를 no_workspace_kernel에 추가합니다:

ascendc_library(no_workspace_kernel STATIC
    # ... existing kernel files ...
    ${PROJECT_OP_SRC_BASE}/<op_name>/op_kernel/<op_name>_kernel.cpp
)

workspace가 필요한 연산자(복잡한 연산자)는 -DHAVE_WORKSPACE -DHAVE_TILING 컴파일 플래그로 workspace_kernel에 커널 소스를 추가합니다:

ascendc_library(workspace_kernel STATIC
    # ... existing kernel files ...
    ${PROJECT_OP_SRC_BASE}/<op_name>/op_kernel/<op_name>_kernel.cpp
)

호스트 소스 파일을 OP_SRCS에 추가:

FILE(GLOB OP_SRCS
    # ... existing host files ...
    ${PROJECT_OP_SRC_BASE}/<op_name>/op_host/<op_name>.cpp
)

7단계: 빌드 (Build)

python/sgl_kernel_npu/README.md의 단계에 따라 빌드합니다:

cd sgl-kernel-npu

# Build all modules
bash build.sh

# Install the sgl_kernel_npu wheel
pip install output/sgl_kernel_npu*.whl

컴파일된 libsgl_kernel_npu.so는 python/sgl_kernel_npu/sgl_kernel_npu/lib/에 복사되고 Python 패키지에 의해 로드됩니다.

Triton 연산자 개발 (Developing Triton Operators)

Triton 연산자는 python/sgl_kernel_npu/sgl_kernel_npu/ 아래에 기능 범주별로 정렬되어 있습니다:

python/sgl_kernel_npu/sgl_kernel_npu/
├── attention/     # Attention (decode_attention, sinks_attention)
├── norm/          # Normalization (rmsnorm, fused_qk_norm, l1_norm)
├── activation/    # Activation (swiglu_oai, swiglu_quant)
├── fla/           # Linear attention (chunk, cumsum, wy_fast)
├── mamba/         # Mamba-related (causal_conv1d, state_update)
├── moe/           # MoE-related (mul_add, zero_experts)
└── sample/        # Speculative decoding (verify_tree_greedy)

개발 단계:

적절한 범주 하위 디렉터리에 새 .py 파일을 만듭니다.
같은 디렉터리의 기존 연산자를 템플릿으로 사용해 Triton 언어로 커널을 작성합니다.
필요하면 해당 __init__.py에서 함수를 내보냅니다.
tests/python/sgl_kernel_npu/ 아래에 테스트를 작성합니다.

참고: 많은 Triton 연산자가 SGLang의 GPU Triton 커널에서 적응되었습니다(예: fla/utils.py의 주석이 원본 소스를 표시). 적응 시 NPU와 GPU의 차이를 특히 주의하세요.

연산자를 SGLang에 통합 (Integrating Operators into SGLang)

Ascend C 연산자 통합

위의 1-7단계(커널 작성, torch op 등록, 빌드)를 완료하고 sgl-kernel-npu 휠을 설치한 후, SGLang에서 다음과 같이 연산자를 호출합니다:

import sgl_kernel_npu  # Loading the library auto-triggers libsgl_kernel_npu.so loading

# Call the operator
result = torch.ops.npu.helloworld(x, y)

SGLang에서의 실제 사용(sglang/srt/speculative/eagle_utils.py에서):

torch.ops.npu.build_tree_kernel_efficient(
    parent_list, selected_index, verified_seq_len, tree_mask,
    positions, retrive_index, retrive_next_token,
    retrive_next_sibling, topk, depth, draft_token_num, tree_mask_mode
)

Triton 연산자 통합

Python으로 직접 임포트·호출:

from sgl_kernel_npu.attention.decode_attention import decode_attention_fwd
from sgl_kernel_npu.norm.rmsnorm_bias import rmsnorm_bias
from sgl_kernel_npu.mamba.causal_conv1d import causal_conv1d_fwd

# Direct function call
output = decode_attention_fwd(q, k, v, ...)

SGLang에서의 실제 사용(sglang/srt/models/llama.py에서):

from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope

sgl-kernel-npu 휠 업데이트 프로세스

SGLang과 sgl-kernel-npu는 별도의 Python 패키지이므로 의존성 업데이트는 multi-PR 워크플로가 필요합니다:

sgl-kernel-npu PR 제출: sgl-kernel-npu 저장소에서 연산자를 추가/수정하고 모든 테스트가 통과하는지 확인합니다.
sgl-kernel-npu 버전 올리기: sgl-kernel-npu의 버전 번호를 업데이트합니다. 병합 시 자동 PyPI 릴리스가 트리거됩니다.
SGLang에서 새 버전 참조:
- SGLang의 python/pyproject.toml에서 sgl-kernel-npu 버전 요구사항을 업데이트합니다.
- SGLang 코드에서 새 연산자를 사용합니다.

긴급하지 않으면 일반 릴리스(보통 1주 이내)를 기다릴 수 있습니다.

단위 테스트 작성 (Writing Unit Tests)

각 연산자는 tests/python/sgl_kernel_npu/ 아래에 해당 단위 테스트가 필요하며, Python의 unittest 프레임워크를 사용합니다.

테스트 파일 이름 규칙: test_<op_name>.py

# tests/python/sgl_kernel_npu/test_helloworld.py
import unittest
import torch
import sgl_kernel_npu

class TestHelloworld(unittest.TestCase):
    def test_helloworld_basic(self):
        x = torch.randn(2048, dtype=torch.float16, device="npu")
        y = torch.randn(2048, dtype=torch.float16, device="npu")

        z = torch.ops.npu.helloworld(x, y)
        expected = x + y

        torch.testing.assert_close(z, expected)

if __name__ == "__main__":
    unittest.main()

테스트 실행:

python tests/python/sgl_kernel_npu/test_helloworld.py

테스트 체크리스트:

전형적인 입력 형태(2의 거듭제곱 크기와 비표준 크기)를 다룹니다.
다른 데이터 타입(bf16 / fp16 등)을 다룹니다.
in-place 동작이 있는 연산자는 출력 텐서의 정확성을 검증합니다.
PyTorch 네이티브 연산과 비교해 정확성을 검증합니다.

코드 스타일 (Code Style)

Pre-commit 검사

sgl-kernel-npu는 일관된 코드 스타일을 위해 pre-commit을 사용합니다:

pip3 install pre-commit
cd sgl-kernel-npu
pre-commit install
pre-commit run --all-files

참고: pre-commit run --all-files가 처음에 실패하면 다시 실행해 모든 린트 오류가 자동 수정되는지 확인하세요. 모든 코드는 PR 제출 전에 검사를 통과해야 합니다.

C++ 코드 스타일

C++17 표준을 사용합니다.
모든 연산자 구현을 sglang::npu_kernel 네임스페이스 아래에 두세요.
기존 코드 스타일을 따르고 .clang-format으로 포맷하세요.
오류 검사에는 TORCH_CHECK를 사용합니다(표준 GE OP_ADD 매크로 사용 금지).
불필요한 GE 등록 코드(예: OP_ADD() 매크로)를 포함하지 마세요.

Python 코드 스타일

PEP 8을 따릅니다.
파일·함수 이름에 snake_case를 사용합니다.
SGLang GPU 코드에서 적응할 때는 파일 헤더에 원본 소스를 명시합니다.

일반 원칙

코드 중복을 피하세요: 5줄 이상 반복되는 코드 블록은 공용 함수로 추출합니다.
디바이스 동기화를 최소화하세요: tensor.item()이나 tensor.cpu() 같은 CPU-NPU 동기화 연산을 줄입니다.
함수를 순수하게 유지하세요: in-place 인자 수정을 피합니다.
파일을 간결하게 유지하세요: 2,000줄을 초과하는 파일은 분할합니다.

PR 제출 (Submitting a PR)

저장소 포크: GitHub에서 sgl-kernel-npu를 포크한 후 로컬에 클론합니다.
브랜치 생성: main에서 새 브랜치를 만듭니다. 예: feature/add-my-op.
개발 및 테스트: 위 단계에 따라 연산자를 개발하고 테스트를 작성합니다. 모든 테스트가 통과하는지 확인합니다.
pre-commit 실행: 코드 포맷 준수를 확인합니다.

커밋 및 푸시:

git add .
git commit -m "feat: add <op_name> operator"
git push origin feature/add-my-op

PR 생성: GitHub에서 브랜치에서 sgl-project/sgl-kernel-npu:main으로 Pull Request를 엽니다.
CI와 리뷰 대기: CI 검사에는 린팅, 컴파일, 연산자 테스트가 포함됩니다. 통과 후 메인테이너 리뷰와 병합을 기다립니다.

연산자 개발