본문으로 건너뛰기

[sglang] [NPU] GLM-4.7-Flash 성능 최적화: Fused Triton 커널로 연산 병목 해결하기

PR 링크: sgl-project/sglang#29509 상태: Merged | 변경: +37 / -15

들어가며

LLM 추론 엔진인 SGLang에서 NPU(Neural Processing Unit) 백엔드의 성능을 극대화하기 위한 중요한 업데이트가 진행되었습니다. 이번 PR은 GLM-4.7-Flash 모델의 추론 속도를 높이기 위해, 기존에 분리되어 있던 split 연산과 RMSNorm 연산을 하나의 Fused Triton Kernel로 통합한 최적화 사례를 담고 있습니다.

GLM-4.7-Flash는 DeepSeek-V2에서 제안된 MLA(Multi-head Latent Attention) 구조를 채택하고 있습니다. MLA는 KV 캐시를 압축하여 메모리 사용량을 줄이지만, 추론 과정에서 압축된 Latent 벡터를 다시 복원하고 정규화하는 추가적인 연산 오버헤드가 발생합니다. 특히 NPU와 같은 가속기 환경에서는 이러한 작은 연산들이 여러 번 호출될 때 발생하는 커널 런칭 오버헤드와 메모리 대역폭 낭비가 전체 성능의 병목이 됩니다.

본 글에서는 실제 코드 변경 사항을 통해 어떻게 연산을 통합하여 효율을 높였는지 분석해 보겠습니다.

코드 분석: MLA 연산의 효율화

1. Split + RMSNorm 파이프라인의 통합

가장 핵심적인 변경 사항은 deepseek_v2_attention_mla_npu.py 파일에서 발견됩니다. 기존에는 qkv_latent 벡터를 먼저 split한 뒤, 각각의 텐서에 대해 LayerNorm을 순차적으로 적용했습니다.

Before: 분리된 연산 호출

# 기존 방식: split 후 각각 LayerNorm 적용
q, latent_cache = (
    get_attn_tp_context()
    .fetch_qkv_latent()
    .split(
        [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
        dim=-1,
    )
)
k_nope = latent_cache[..., : m.kv_lora_rank]

q = m.q_a_layernorm(q)
# ... (중략)
k_nope = m.kv_a_layernorm(k_nope)

위 코드의 문제점은 split 연산이 새로운 텐서 뷰를 생성하거나 메모리 복사를 유발하고, 이후 q_a_layernorm과 kv_a_layernorm이 각각 별도의 커널로 실행된다는 점입니다. 이는 HBM(High Bandwidth Memory)에 대한 읽기/쓰기 횟수를 증가시킵니다.

After: Fused Kernel 도입

# 개선된 방식: fused_split_qk_norm 하나로 통합
if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp(forward_batch):
    q, k_nope, k_pe = fused_split_qk_norm(
        qkv_latent,
        m.q_a_layernorm,
        m.kv_a_layernorm,
        m.q_lora_rank,
        m.kv_lora_rank,
        m.qk_rope_head_dim,
        eps=m.q_a_layernorm.variance_epsilon,
    )
else:
    # Fallback logic for large batches or specific conditions
    q, latent_cache = qkv_latent.split(
        [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
        dim=-1,
    )
    # ... (기존 로직 유지)

fused_split_qk_norm 함수는 Triton으로 작성된 커널을 호출하여, 입력 텐서를 읽어 들이는 단 한 번의 과정에서 split 위치를 계산하고 곧바로 RMSNorm 연산을 수행합니다.

주목할 점은 조건부 실행입니다. qkv_latent.shape[0] < 65536일 때만 Fused 커널을 사용하는데, 이는 커널의 그리드 크기나 NPU의 하드웨어 특성에 따라 최적의 성능을 낼 수 있는 워크로드 범위를 설정한 것으로 보입니다. 너무 큰 배치 사이즈에서는 기존의 분할 방식이 더 안정적일 수 있기 때문입니다.

2. MoE TopK Scoring Function 최적화

topk.py에서는 MoE(Mixture of Experts) 레이어의 라우팅 로직이 수정되었습니다.

Before/After 비교

# Before
norm_type=(0 if topk_config.scoring_func == "softmax" else 1),

# After
norm_type=1,

기존에는 scoring_func 설정에 따라 softmax(0)와 sigmoid(1)를 선택했으나, 이번 변경을 통해 1(Sigmoid)로 고정되었습니다. GLM-4.7-Flash 모델의 특성상 MoE 게이팅에서 Sigmoid 기반의 정규화가 성능이나 정확도 측면에서 더 유리하거나, NPU 전용 커널인 fused_topk_npu가 Sigmoid 모드에서 최적화되어 있음을 시사합니다.

왜 이게 좋은가?

1. 메모리 대역폭 절약 (Memory Bound 해결)

현대적인 GPU/NPU에서 RMSNorm과 같은 Element-wise 연산은 연산량(Compute)보다 메모리 접근(Memory Access)에 더 많은 시간이 소요되는 Memory-bound 작업입니다. split과 norm을 합치면 중간 결과물을 HBM에 썼다가 다시 읽어오는 과정을 생략할 수 있어 이론적으로 2~3배 이상의 메모리 효율을 얻을 수 있습니다.

2. 커널 런칭 오버헤드 감소

딥러닝 모델 추론 시 수천 개의 커널이 실행됩니다. 각 커널이 실행될 때마다 호스트(CPU)에서 디바이스(NPU)로 명령을 내리는 오버헤드가 발생하는데, 3개의 커널을 1개로 줄이는 것은 특히 짧은 시퀀스를 처리하는 Decoding 단계에서 지연 시간(Latency)을 줄이는 데 결정적인 역할을 합니다.

3. 정확도와 속도의 트레이드오프 검증

PR 설명에 첨부된 Accuracy Test 결과에 따르면, 이러한 커널 최적화 이후에도 모델의 출력값이 기존과 동일하게 유지됨을 확인할 수 있습니다. 성능을 위해 정확도를 희생하지 않았다는 점이 이 PR의 완성도를 보여줍니다.

결론

이번 PR은 SGLang이 NPU 하드웨어의 잠재력을 끌어올리기 위해 얼마나 세밀한 최적화를 진행하고 있는지 잘 보여줍니다. 특히 DeepSeek-V2 계열의 MLA 구조처럼 복잡한 텐서 조작이 필요한 경우, Triton을 이용한 커널 퓨전은 선택이 아닌 필수입니다.

리뷰 과정에서 sglang-npu-bot과 Estrella-xx가 CI 테스트를 반복하며 안정성을 검증한 것 또한 오픈소스 프로젝트로서 신뢰도를 높이는 부분입니다. NPU 환경에서 LLM을 서비스하고자 하는 엔지니어라면, 이러한 커널 퓨전 기법을 자신의 모델에도 적용해 보시길 권장합니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글