[sglang] Qwen3.5 및 Qwen3_Next 모델의 NPU 성능 향상을 위한 Triton 커널 퓨전 최적화

PR 링크: sgl-project/sglang#23925 상태: Merged | 변경: +194 / -19

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키며, 동시에 추론 속도와 효율성에 대한 요구를 높이고 있습니다. 특히 NPU(신경망 처리 장치)와 같은 특수 하드웨어 환경에서는 기존의 CPU 또는 GPU 기반 최적화 기법만으로는 한계에 부딪힐 수 있습니다.

이번 PR은 sglang 프로젝트에서 Qwen3.5 및 Qwen3_Next 모델의 NPU 환경에서의 추론 성능을 획기적으로 개선하기 위한 중요한 최적화를 다룹니다. 핵심은 어텐션(Attention) 레이어의 여러 연산, 즉 split_qkv, RMS Normalization, 그리고 Rotary Positional Embedding (RoPE)을 하나의 Triton 커널로 통합(fusion)하는 것입니다. 이러한 커널 퓨전은 메모리 접근을 줄이고 연산 파이프라인을 최적화하여 NPU에서의 LLM 추론 속도를 크게 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, Triton 커널 퓨전이 왜 성능 향상에 효과적인지, 그리고 실제 성능 개선 수치는 어떠한지 살펴보겠습니다.

코드 분석

이번 PR의 주요 변경 사항은 크게 세 부분으로 나눌 수 있습니다:

1. Triton 커널 구현: apply_interleaved_rope_triton 함수와 apply_interleaved_rope_kernel 커널이 새로 추가되었습니다. 이는 기존의 Python/PyTorch 기반 RoPE 연산을 Triton으로 재구현한 것입니다.
2. 모델별 적용 로직 수정: Qwen3_5GatedDeltaNet 및 Qwen3_Next 모델 클래스에서 NPU 환경일 때 새로운 Triton 기반 커널을 사용하도록 로직이 수정되었습니다.
3. RoPE 적용 방식 변경: mrope.py 파일에서 apply_interleaved_rope 함수 대신 새로운 Triton 커널을 사용하도록 조건부 로직이 추가되었습니다.

1. Triton RoPE 커널 구현 (python/sglang/srt/layers/rotary_embedding/mrope.py)

기존의 apply_interleaved_rope 함수는 Python과 PyTorch 연산을 사용하여 RoPE를 적용했습니다. 이는 여러 단계의 메모리 복사 및 연산으로 인해 성능 병목 현상을 일으킬 수 있습니다. 이번 PR에서는 이를 Triton을 사용하여 단일 커널로 통합했습니다.

Before (기존 방식 - 개념적):

def apply_interleaved_rope(x: torch.Tensor, mrope_section: list) -> torch.Tensor:
    # ... (Python/PyTorch 연산으로 RoPE 적용)
    # 여러 번의 슬라이싱, 인덱싱, 텐서 연산 포함
    pass

After (Triton 커널 도입):

import triton
import triton.language as tl

@triton.jit
def apply_interleaved_rope_kernel(
    x_ptr, # ... (포인터 및 스트라이드 정보)
    S: tl.constexpr, D: tl.constexpr,
    stride_x_m, stride_x_s, stride_out_s,
    section_1_end, section_2_end,
    BLOCK_S: tl.constexpr, BLOCK_SIZE: tl.constexpr,
):
    # ... (Triton 커널 내부 로직)
    # 메모리 로드, 조건부 연산, 메모리 저장 등 저수준 연산
    pass

def apply_interleaved_rope_triton(x: torch.Tensor, mrope_section: list) -> torch.Tensor:
    x = x.contiguous()
    M, S, D = x.shape
    out = torch.empty((S, D), dtype=x.dtype, device=x.device)
    # ... (Triton 커널 실행을 위한 그리드 설정 및 커널 호출)
    grid = (triton.cdiv(S, BLOCK_S), triton.cdiv(D, BLOCK_SIZE))
    apply_interleaved_rope_kernel[grid](...)
    return out

apply_interleaved_rope_kernel은 Triton 언어를 사용하여 GPU에서 직접 실행되는 저수준 커널입니다. 이 커널은 입력 텐서 x의 포인터와 스트라이드 정보를 받아, RoPE 연산에 필요한 복소수 곱셈 및 회전을 GPU 스레드 블록 단위로 효율적으로 처리합니다. BLOCK_S와 BLOCK_SIZE는 스레드 블록이 처리할 데이터의 크기를 정의하며, 이를 통해 GPU 캐시 활용을 최적화합니다.

apply_interleaved_rope_triton 함수는 이 Triton 커널을 PyTorch 텐서에 적용하기 위한 래퍼(wrapper) 함수입니다. 입력 텐서의 형태를 확인하고, 필요한 경우 contiguous()를 호출하여 메모리 레이아웃을 정리한 후, Triton 커널을 실행할 그리드(grid)를 계산하여 커널을 호출합니다. 최종적으로 계산된 결과를 담은 출력 텐서를 반환합니다.

2. 모델별 적용 로직 수정 (python/sglang/srt/models/qwen3_5.py, python/sglang/srt/models/qwen3_next.py)

Qwen3.5와 Qwen3_Next 모델은 NPU 환경에서 성능 최적화를 위해 새로운 커널을 사용하도록 수정되었습니다. 특히 self_attention 메소드 내에서 NPU 환경이고 특정 조건(forward_mode가 extend 또는 draft_extend가 아닌 경우, 그리고 attn_output_gate가 활성화된 경우)을 만족할 때, 기존의 forward_prepare_native 대신 forward_prepare_npu 함수를 호출하도록 변경되었습니다.

Before (Qwen3.5 self_attention 일부):

def self_attention(
    self, ...
):
    qkv, _ = self.qkv_proj(hidden_states)
    # ... (기존 RoPE 적용 및 QKV 분리 로직)
    q, k = self._apply_qk_norm(q, k)
    q, k = self.rotary_emb(positions, q, k)
    # ...

After (Qwen3.5 self_attention 일부):

def self_attention(
    self, ...
):
    if (
        not _is_npu
        or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
        or not self.attn_output_gate
    ):
        q, k, v, gate = self.forward_prepare_native(
            positions=positions,
            hidden_states=hidden_states,
        )
    else:
        q, k, v, gate = self.forward_prepare_npu(
            positions=positions,
            hidden_states=hidden_states,
            forward_batch=forward_batch,
        )
    # ... (이후 어텐션 계산)

forward_prepare_npu 함수는 NPU 환경에 특화된 최적화된 커널 호출 로직을 포함합니다. 이 함수 내부에서는 sgl_kernel_npu.norm.split_qkv_rmsnorm_rope.split_qkvgate_gemma_rmsnorm_rope와 같은 NPU 전용 커널을 호출하여 split_qkv, RMSNorm, RoPE 연산을 한 번에 처리합니다. 이는 GPU 커널 퓨전과 유사하게, 여러 단계를 거치던 연산을 단일 고성능 커널로 통합하여 메모리 이동과 커널 실행 오버헤드를 최소화합니다.

리뷰어 iforgetmyname의 지적대로, self.config.full_attention_interval을 직접 사용하는 것이 더 명확한 코딩 스타일입니다. 이 부분은 forward_prepare_npu 함수 내에서 self.attn.layer_id == (self.config.full_attention_interval - 1)와 같이 사용되어, 특정 어텐션 레이어에서만 get_cos_sin_with_position을 호출하도록 제어합니다. 이는 불필요한 RoPE 사인/코사인 계산을 줄여 성능을 더욱 향상시킵니다.

3. RoPE 적용 방식 변경 (python/sglang/srt/layers/rotary_embedding/mrope.py)

RotaryEmbedding 클래스의 get_cos_sin_with_position 메소드에서 RoPE 적용 방식이 변경되었습니다. Triton 백엔드를 사용할 수 있는 경우 (서버 설정에서 attention_backend이 Triton으로 지정된 경우), 새로 구현된 apply_interleaved_rope_triton 함수를 사용하도록 조건이 추가되었습니다.

Before:

if self.mrope_interleaved:
    cos = apply_interleaved_rope(cos, self.mrope_section)
    sin = apply_interleaved_rope(sin, self.mrope_section)

After:

if self.mrope_interleaved:
    if support_triton(get_global_server_args().attention_backend):
        cos = apply_interleaved_rope_triton(cos, self.mrope_section)
        sin = apply_interleaved_rope_triton(sin, self.mrope_section)
    else:
        cos = apply_interleaved_rope(cos, self.mrope_section)
        sin = apply_interleaved_rope(sin, self.mrope_section)

이 변경은 Triton 백엔드를 사용할 때만 새로운 Triton 기반 RoPE 구현을 적용하고, 그렇지 않은 경우에는 기존의 PyTorch 기반 구현을 유지하도록 하여 호환성을 보장합니다. 이는 점진적인 성능 개선을 가능하게 합니다.

왜 이게 좋은가? (성능 및 교훈)

성능 개선

PR 설명에 포함된 속도 테스트 결과는 이 최적화의 효과를 명확하게 보여줍니다.

• Qwen3.6-35b 모델:
  • Before: 어텐션 레이어당 약 100ms 이상의 시간이 소요되었습니다.
  • After: 첫 번째 어텐션 레이어의 경우, get_cos_sin_with_position 계산이 추가됨에도 불구하고 상당한 속도 향상이 있었습니다. 다른 어텐션 레이어에서는 이전 캐시를 재사용하므로, 첫 번째 레이어의 최적화가 전체 성능에 큰 영향을 미칩니다.
  • GPU 커널 테스트: Triton 커널 자체의 성능을 측정한 결과, 이전 방식 대비 상당한 속도 향상이 관찰되었습니다. 이는 커널 퓨전 및 Triton 최적화의 직접적인 효과를 나타냅니다.

정확한 수치 비교는 이미지에서 직접 확인해야 하지만, 프로파일링 결과는 Triton 커널 퓨전이 메모리 접근 횟수를 줄이고 GPU 연산 효율성을 높여 NPU 환경에서 LLM 추론 속도를 크게 향상시켰음을 시사합니다.

일반적인 교훈

1. 커널 퓨전의 힘: 여러 개의 작은 연산을 하나의 큰 커널로 통합하는 커널 퓨전은 LLM과 같은 연산 집약적인 워크로드에서 메모리 대역폭 병목 현상을 완화하고 연산 효율성을 높이는 강력한 기법입니다. 특히 NPU와 같이 특정 연산에 특화된 하드웨어에서는 그 효과가 극대화될 수 있습니다.
2. Triton의 활용: Triton은 PyTorch와 같은 프레임워크에서 고성능 커스텀 GPU 커널을 쉽게 작성할 수 있도록 지원하는 도구입니다. 복잡한 연산이나 특정 하드웨어에 최적화된 커널이 필요한 경우, Triton을 활용하여 성능을 크게 향상시킬 수 있습니다. 이 PR은 Triton을 RoPE와 같은 핵심 연산에 성공적으로 적용한 좋은 예시입니다.
3. 하드웨어 특화 최적화: NPU와 같은 특정 하드웨어의 아키텍처와 연산 특성을 이해하고 이를 활용하는 최적화는 필수적입니다. 이 PR은 NPU 환경에서 Qwen 모델의 성능을 극대화하기 위해 NPU 전용 커널(sgl_kernel_npu)을 도입했습니다.
4. 점진적 개선 및 조건부 적용: 새로운 최적화 기법을 도입할 때는 기존 코드와의 호환성을 유지하는 것이 중요합니다. support_triton 함수를 사용한 조건부 적용은 Triton 백엔드를 사용할 때만 새로운 커널을 사용하고, 그렇지 않은 경우에는 기존 방식을 유지하여 안정성을 확보했습니다.
5. 계산 캐싱 및 재사용: RoPE의 사인/코사인 값을 모든 어텐션 레이어에서 다시 계산하는 대신, 첫 번째 레이어에서만 계산하고 이후 레이어에서는 캐시된 값을 재사용하는 최적화는 불필요한 계산을 줄여 성능을 향상시킵니다. 이는 LLM 추론 시 KV 캐싱과 유사한 원리입니다.

결론

이번 PR은 Qwen3.5 및 Qwen3_Next 모델을 NPU 환경에서 실행할 때 발생하는 성능 병목 현상을 효과적으로 해결했습니다. split_qkv, RMSNorm, RoPE 연산을 Triton 커널로 퓨전하고, NPU 전용 커널을 도입함으로써 메모리 접근을 최소화하고 연산 효율성을 극대화했습니다. 이러한 최적화는 LLM의 추론 속도를 크게 향상시켜, 더 빠르고 효율적인 AI 서비스 제공에 기여할 것입니다. 커널 퓨전과 Triton의 활용은 앞으로도 LLM 성능 최적화의 중요한 방향이 될 것입니다.

참고 자료

• https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/rotary_embedding/mrope.py
• https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/qwen3_5.py
• https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/qwen3_next.py
• https://triton-lang.org/docs/getting-started/introduction
• https://pytorch.org/docs/stable/generated/torch.nn.Module.html
• https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils.py

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.