[SGLang] Hardware Backends: MLX, NPU, XPU 하드웨어 추상화

2026년 4월 15일수정: 2026년 4월 15일

들어가며

SGLang은 NVIDIA GPU 외에도 Apple Silicon(MLX), Huawei Ascend NPU, Intel XPU 등 다양한 하드웨어를 지원한다. hardware_backend/ 디렉토리는 각 하드웨어에 특화된 모델 실행기, 메모리 관리, 어텐션 커널, 양자화 등을 제공한다. 이 포스트에서는 각 백엔드의 구조와 CUDA 기본 경로와의 차이점을 분석한다.

구조도

hardware_backend/
├── mlx/                      ── Apple Silicon (MLX)
│   ├── model_runner.py        ── MLX 전용 모델 실행기
│   ├── model_runner_stub.py   ── 스텁 (미지원 기능)
│   └── tp_worker.py           ── TP 워커
└── npu/                      ── Huawei Ascend NPU
    ├── utils.py               ── 서버 기본값, 메모리 설정
    ├── allocator_npu.py       ── NPU 메모리 할당기
    ├── memory_pool_npu.py     ── NPU KV 캐시 풀
    ├── attention/             ── NPU 어텐션 커널
    ├── graph_runner/          ── NPU 그래프 실행기
    ├── modules/               ── NPU 모듈
    ├── moe/                   ── NPU MoE
    ├── quantization/          ── NPU 양자화
    ├── cmo.py                 ── 캐시 메모리 최적화
    └── utils.py               ── NPU 유틸리티

[추상화 레벨]
MultiPlatformOp (layers/utils/)
    │
    ├── forward_cuda()  ── NVIDIA CUDA
    ├── forward_npu()   ── Huawei Ascend NPU
    ├── forward_cpu()   ── Intel CPU (AMX)
    ├── forward_xpu()   ── Intel XPU
    └── forward_hip()   ── AMD ROCm

핵심 코드 분석

MLX Backend: Apple Silicon 전용

MLX 백엔드는 PyTorch MPS를 우회하고 MLX 프레임워크에서 직접 모델을 실행한다. 이는 MPS-MLX 텐서 브릿지 오버헤드를 제거한다.

class MlxModelRunner:
    """Model runner that executes the entire model in MLX.
    This avoids the MPS<->MLX tensor bridge overhead by keeping all
    computation within MLX."""

MLX의 mlx_lm 라이브러리를 사용하여 모델을 로드하고, 배치 KV 캐시를 관리한다.

import mlx.core as mx
from mlx_lm import load as mlx_lm_load
from mlx_lm.models.cache import (
    BatchKVCache, BatchRotatingKVCache, KVCache, RotatingKVCache,
    make_prompt_cache)

@dataclass
class MlxRequestState:
    token_ids: list[int]
    cache: list       # List of KVCache per layer
    generated_tokens: int = 0

KV 캐시 병합과 추출은 MLX의 BatchKVCache.merge와 extract를 사용한다.

def _merge_kv_caches(caches_list):
    for layer_idx in range(num_layers):
        layer_caches = [caches[layer_idx] for caches in caches_list]
        if isinstance(layer_caches[0], KVCache):
            batch_cache = BatchKVCache.merge(layer_caches)
        elif isinstance(layer_caches[0], RotatingKVCache):
            batch_cache = BatchRotatingKVCache.merge(layer_caches)
    return merged

NPU Backend: Huawei Ascend 설정

NPU 백엔드는 서버 시작 시 기본값을 자동으로 설정한다. 어텐션 백엔드는 ascend로 고정되고, 메모리 설정은 NPU 사양에 따라 조정된다.

def set_default_server_args(args: "ServerArgs"):
    args.attention_backend = "ascend"
    args.prefill_attention_backend = "ascend"
    args.decode_attention_backend = "ascend"
    if args.page_size is None:
        args.page_size = 128

    npu_mem = get_npu_memory_capacity()
    if npu_mem <= 32 * 1024:
        # Ascend 910B4 - 메모리 제한 설정
        ...

NPU는 자체 텐서 포맷(ACL_FORMAT_FRACTAL_NZ)을 사용하며, MoE 연산 모드도 별도로 정의한다.

class NPUACLFormat(IntEnum):
    ACL_FORMAT_UNDEFINED = -1
    ACL_FORMAT_ND = 2
    ACL_FORMAT_FRACTAL_NZ = 29

class FusedMoEMode(IntEnum):
    FUSED_DEEP_MOE = 1
    DISPATCH_FFN_COMBINE = 2

MultiPlatformOp: 통합 디스패치

layers/utils/의 MultiPlatformOp은 하드웨어 감지를 기반으로 적절한 forward_* 메서드를 자동 선택한다. 각 레이어(RoPE, Activation 등)는 이를 상속하여 플랫폼별 최적 구현을 제공한다.

# activation.py에서의 플랫폼별 분기 예시
class SiluAndMul(MultiPlatformOp):
    def forward_cuda(self, x):   # NVIDIA
        silu_and_mul(x, out)
    def forward_npu(self, x):    # Ascend NPU
        torch_npu.npu_swiglu(x)
    def forward_cpu(self, x):    # Intel AMX
        torch.ops.sgl_kernel.silu_and_mul_cpu(x)
    def forward_xpu(self, x):    # Intel XPU
        silu_and_mul(x, out)

NPU 전용 모듈

NPU 백엔드는 자체 어텐션 구현, 그래프 실행기, 양자화 방식을 갖고 있다.

npu/
├── attention/      ── NPU Ascend 어텐션 (npu_mrope 등)
├── graph_runner/   ── NPU 전용 CUDA Graph 대체
├── modules/        ── NPU 최적화 모듈
├── moe/            ── NPU MoE (FUSED_DEEP_MOE)
└── quantization/   ── NPU 양자화 (GPTQLinearAscendMethod 등)

하드웨어별 비교

구분	CUDA (NVIDIA)	MLX (Apple)	NPU (Ascend)	XPU (Intel)
프레임워크	PyTorch CUDA	MLX native	PyTorch + torch_npu	PyTorch XPU
어텐션	FlashInfer/FA3	MLX 내장	Ascend 전용	sgl_kernel
GEMM	DeepGEMM/CUTLASS	MLX matmul	NPU GEMM	oneDNN
KV 캐시	페이지드	BatchKVCache	페이지드 (128)	페이지드
양자화	FP8/INT8/INT4	MLX quantize	GPTQ/AWQ Ascend	AWQ
그래프	CUDA Graph	-	NPU Graph	-

XPU 지원 (layers 수준)

Intel XPU는 hardware_backend/ 아래에 별도 디렉토리가 없지만, 각 레이어의 forward_xpu 메서드와 sgl_kernel의 XPU 바인딩을 통해 지원된다.

# rotary_embedding/base.py
def forward_xpu(self, positions, query, key, ...):
    return torch.ops.sgl_kernel.rotary_embedding(
        positions, query, key, self.head_size,
        self.cos_sin_cache, self.is_neox_style)

참고

소스 코드: python/sglang/srt/hardware_backend/
Apple MLX: github.com/ml-explore/mlx
Huawei Ascend CANN 문서
Intel oneAPI 문서

SGLang 의 다른글

이전글 [SGLang] Reasoning & Code Completion Parser: 추론 및 코드 파서
현재글 : [SGLang] Hardware Backends: MLX, NPU, XPU 하드웨어 추상화
다음글 없음