[sglang] CuTe-DSL 정규화 커널 최적화: 워프당 여러 청크 처리로 성능 향상

2026년 10월 26일수정: 2026년 10월 26일

PR 링크: sgl-project/sglang#46227 상태: Merged | 변경: +223 / -157

들어가며

최근 sglang 프로젝트의 diffusion 모델 관련 코드에서 성능 최적화 PR이 올라왔습니다. 이 PR의 핵심은 fused_norm_scale_shift 및 fused_scale_residual_norm_scale_shift와 같은 CuTe-DSL 기반 정규화 커널의 연산 방식을 개선하여 GPU 활용률을 높이고 전체적인 연산 시간을 단축하는 것입니다. 기존에는 각 워프(warp)가 한 번에 하나의 256개 컬럼 청크(chunk)만 처리했지만, 이 PR에서는 워프당 처리하는 청크 수를 늘려 GPU의 잠재력을 최대한 활용하고자 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 그 과정에서 얻을 수 있는 일반적인 최적화 교훈은 무엇인지 살펴보겠습니다.

코드 분석

1. `norm/cutedsl_common/norm_fusion.py` 및 `norm/cutedsl_common/reduce.py` 변경

이 파일들은 CuTe-DSL 커널의 공통 로직을 다룹니다. 주요 변경점은 다음과 같습니다.

apply_norm_cta -> row_norm_stats 및 normalize_chunk 분리: 기존의 apply_norm_cta 함수는 LayerNorm과 RMSNorm을 모두 처리하며, 각 워프가 하나의 256 컬럼 청크를 담당했습니다. 이를 row_norm_stats (행의 통계량 계산)와 normalize_chunk (청크 정규화)로 분리하여 각 함수의 역할을 명확히 했습니다.
row_norm_stats의 역할 변경: 이 함수는 이제 행 전체의 평균(mean)과 표준편차의 역수(1/std)를 계산하는 역할을 합니다. 이전에는 각 워프가 독립적으로 청크를 처리했지만, 이제는 여러 워프가 협력하여 행 전체의 통계량을 계산합니다. 특히, cta_reduce_sum 함수가 수정되어 여러 워프의 부분 합계를 더 효율적으로 통합하도록 변경되었습니다.

Before (기존 apply_norm_cta의 일부 로직):
```
# ... (rms norm logic) ...
val = warp_reduce_sum(val)
acc_sq = cta_reduce_sum(val, num_warps, tidx)
factor = cute.rsqrt(acc_sq / D + eps)
# ... (normalization) ...
```
After (새로운 row_norm_stats의 일부 로직):
```
# ... (rms norm logic for multiple chunks) ...
vals.append(warp_reduce_sum(val))
# ...
acc_sq = cta_reduce_sum(vals, num_chunks, warp_id, lane_id)
# ...
return cute.Float32(0.0), cute.rsqrt(acc_sq / D + eps)
```
cta_reduce_sum 함수 시그니처 변경:

Before:

def cta_reduce_sum(val: cute.Numeric, num_warps: cutlass.Constexpr, tidx: cutlass.Int32) ```

**After:**
```python

def cta_reduce_sum(vals: list, num_chunks: cutlass.Constexpr, warp_id: cutlass.Int32, lane_id: cutlass.Int32) ```

이 변경은 `cta_reduce_sum`이 이제 여러 워프의 부분 합계(`vals`)를 입력받아, `num_chunks`만큼의 청크에 대한 최종 합계를 계산하도록 확장되었음을 보여줍니다. 또한, `warp_id`와 `lane_id`를 사용하여 스레드 블록 내에서 더 세밀한 제어를 가능하게 합니다.

tensor_slice_for_bsfd의 청크 처리 로직 추가: 이 함수는 BSFD(Batch, Sequence, Feature, Dimension) 텐서에서 데이터를 슬라이싱하는 역할을 합니다. 이제 chunk_id와 chunk 크기 정보를 받아, 특정 워프가 처리해야 할 컬럼 범위를 명확히 지정합니다.

Before:
```
# ... (slice logic) ...
return tVgV, tVrV
```
After:
```
# ... (slice logic) ...
gV = cute.local_tile(gV, tiler=(chunk,), coord=(chunk_id,))
tVgV = thr_copy.partition_S(gV)
tVrV = cute.make_fragment_like(tVgV, tVgV.element_type)
return tVgV, tVrV
```
cute.local_tile 함수를 사용하여 chunk_id에 해당하는 chunk 크기의 데이터를 gV 텐서에서 추출하는 로직이 추가되었습니다. 이는 각 워프가 이전보다 더 큰 데이터 조각을 한 번에 다룰 수 있게 함을 의미합니다.

2. `norm/scale_residual_norm_cutedsl.py` 변경

이 파일은 실제 Scale-Residual Norm 커널 구현을 담당합니다. 주요 변경 사항은 다음과 같습니다.

chunks_per_warp 도입: 커널 클래스(ScaleResidualNormScaleShift)의 생성자에서 chunks_per_warp라는 새로운 파라미터가 추가되었습니다. 이 값은 각 워프가 처리할 청크의 수를 결정합니다. 이전에는 워프당 고정된 256개 컬럼 청크만 처리했지만, 이제는 이 값을 통해 워프당 처리량을 동적으로 조절할 수 있습니다.

Before (생성자):
```
def __init__(self, D: int, norm_type: str):
    self.D = D
    self.norm_type = norm_type  # "layer" or "rms"
    self.num_warps = self.D // 256  # num of warps per cta
    self.num_threads = self.num_warps * WARP_SIZE  # num of threads per cta
```
After (생성자):

def init(self, D: int, norm_type: str, chunks_per_warp: int): self.D = D self.norm_type = norm_type # "layer" or "rms" # a row is D // CHUNK chunks; each warp walks chunks_per_warp of them, so the # reduction order is fixed by D alone self.num_chunks = D // CHUNK self.chunks_per_warp = chunks_per_warp self.num_threads = self.num_chunks // chunks_per_warp * WARP_SIZE ```

`num_chunks`와 `chunks_per_warp`가 도입되면서 `num_threads` 계산 방식도 변경되었습니다. 이는 각 워프가 더 많은 청크를 처리하도록 스레드 구성을 최적화한 결과입니다.

커널 로직 수정: __call__ 메소드 내의 커널(kernel) 함수에서 데이터 처리 방식이 변경되었습니다. 이전에는 apply_norm_cta를 호출하여 각 워프가 독립적으로 청크를 처리했지만, 이제는 row_norm_stats를 호출하여 행 전체의 통계량을 계산하고, normalize_chunk를 사용하여 각 워프가 할당된 여러 청크를 처리합니다.

Before (커널 내부):
```
# ...
tNrN = apply_norm_cta(
    norm_type=self.norm_type,
    num_warps=self.num_warps,
    tidx=tidx,
    tXrX=tXrX,
    tWrW=tWrW,
    tBrB=tBrB,
    D=self.D,
    eps=eps,
)
# ...
```
After (커널 내부):
```
# ...
mean, factor = row_norm_stats(
    norm_type=self.norm_type,
    chunks=chunks,
    num_chunks=self.num_chunks,
    warp_id=warp_id,
    lane_id=lane_id,
    D=self.D,
    eps=eps,
)
# ...
tNrN = normalize_chunk(
    norm_type=self.norm_type,
    tXrX=tXrX,
    tWrW=tWrW,
    tBrB=tBrB,
    mean=mean,
    factor=factor,
)
# ...
```
apply_norm_cta 대신 row_norm_stats와 normalize_chunk를 순차적으로 호출하는 방식으로 변경되었습니다. 이는 행 전체의 통계량을 먼저 계산한 후, 각 워프가 자신의 청크들을 정규화하는 새로운 워크플로우를 반영합니다.
CHUNK 상수 정의: CHUNK = WARP_SIZE * 8 상수가 정의되었습니다. 이는 한 워프가 128비트 로드(Load)를 통해 처리할 수 있는 컬럼 수를 나타냅니다. 이 상수는 tensor_slice_for_bsfd 함수에서 사용되어, 각 워프가 처리할 데이터 청크의 크기를 결정하는 데 기여합니다.

왜 이게 좋은가?

이 PR의 핵심 목표는 GPU의 잠재력을 최대한 활용하여 정규화 연산의 성능을 높이는 것입니다. 주요 개선 사항과 그 효과는 다음과 같습니다.

향상된 GPU 활용률: 기존 커널은 각 워프가 한 번에 하나의 256 컬럼 청크만 처리했습니다. 이는 GPU의 병렬 처리 능력을 충분히 활용하지 못하는 병목 현상을 야기할 수 있습니다. 특히 D (hidden dimension)가 클 경우, 많은 수의 워프가 동시에 실행되지만 각 워프의 작업량이 적어 GPU 코어가 유휴 상태로 남는 경우가 발생합니다.

이 PR에서는 chunks_per_warp를 도입하여 각 워프가 여러 개의 청크(예: 5개 또는 10개)를 순차적으로 처리하도록 변경했습니다. 이를 통해 워프는 더 많은 데이터를 처리하게 되고, 결과적으로 GPU 코어의 활용률이 높아집니다. 이는 특히 D가 큰 경우에 더 큰 성능 향상을 가져옵니다.
메모리 접근 패턴 최적화: 워프가 여러 청크를 연속적으로 처리하면서, 캐시 활용도 및 메모리 대역폭 사용률이 개선될 수 있습니다. 연속된 메모리 영역에 대한 접근은 캐시 히트율을 높이고, 불필요한 메모리 접근 지연을 줄여줍니다.
성능 수치: PR 설명에 제시된 속도 테스트 결과는 이러한 최적화의 효과를 명확히 보여줍니다.
- B300, [1, 37800, 5120] bf16, fp32 modulation (capped):
  - norm1 (LayerNorm, scale/shift): 0.318 ms -> 0.218 ms (약 31% 향상)
  - self-attn residual (gate, fp32 affine): 0.476 ms -> 0.326 ms (약 31% 향상)
  - cross-attn residual (scale/shift): 0.370 ms -> 0.278 ms (약 25% 향상)
- 다른 Shape에서의 커널 시간:
  - 1 x 4608 x 3072 norm1: 15.1 us -> 13.6 us
  - 1 x 4608 x 3072 gated affine residual: 27.7 us -> 24.8 us
  - 1 x 4096 x 3840 norm1: 17.0 us -> 13.2 us (약 22% 향상)
  - 1 x 4096 x 3840 gated affine residual: 36.1 us -> 29.3 us (약 19% 향상)
- End-to-End (E2E) 성능: Wan2.2 T2V A14B 720p (4xB300)에서 이 PR과 다른 최적화(#46228)가 함께 적용되었을 때, 요청 처리 시간이 90.656초에서 90.187초로 약 0.47초 단축되었습니다. 이는 전체 파이프라인에서 정규화 연산이 차지하는 비중이 상당함을 시사합니다.
정확도 유지: 가장 중요한 점은 이러한 성능 향상이 모델의 정확도를 전혀 저하시키지 않는다는 것입니다. test_norm.py의 테스트 케이스들(test_cutedsl_norm_chunks_per_warp_keep_bits, 비트 단위 비교 등)을 통해 기존 커널과 새로운 커널의 출력이 비트 단위로 동일함을 검증했습니다. 이는 정확한 수치 계산과 재현성을 보장합니다.

일반적인 최적화 교훈

이 PR은 다음과 같은 일반적인 최적화 교훈을 제공합니다:

GPU 아키텍처 이해 및 활용: 워프(warp)와 스레드 블록(CTA)의 동작 방식을 깊이 이해하고, GPU의 병렬 처리 능력을 최대한 활용하도록 커널을 설계해야 합니다. 특히, 워프당 처리량을 늘리는 것은 GPU 활용률을 높이는 효과적인 방법 중 하나입니다.
데이터 처리 단위 최적화: 연산의 기본 단위를 적절히 조절하는 것이 중요합니다. 이 경우, 워프가 처리하는 데이터 청크 크기를 늘림으로써 전체적인 연산 효율성을 높였습니다.
메모리 접근 패턴 최적화: 연속적인 메모리 접근은 캐시 효율성을 높이고 대역폭 활용률을 개선합니다. 데이터 레이아웃과 처리 순서를 최적화하여 메모리 병목 현상을 완화해야 합니다.
정확도와 성능의 균형: 성능 최적화 과정에서 수치적 정밀도 손실이나 정확도 저하가 발생하지 않도록 주의해야 합니다. FP32 누적 합계 사용, 비트 단위 검증 등은 이를 보장하는 좋은 방법입니다.
모듈화 및 재사용성: CuTe-DSL과 같은 프레임워크를 사용하여 커널을 모듈화하고 재사용 가능한 컴포넌트(예: row_norm_stats, normalize_chunk)로 분리하면 코드의 유지보수성과 확장성을 높일 수 있습니다.

결론

이 PR은 CuTe-DSL 기반 정규화 커널의 성능을 크게 향상시키는 성공적인 최적화 사례입니다. 워프당 더 많은 데이터 청크를 처리하도록 커널 로직을 재설계함으로써 GPU 활용률을 높이고, 실제 연산 시간을 단축했습니다. 또한, 정확도를 완벽하게 유지하면서 이러한 성능 개선을 달성했다는 점에서 높은 가치를 지닙니다. 이는 대규모 딥러닝 모델의 추론 및 학습 성능을 개선하는 데 중요한 기여를 할 것입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer의 SM120 FP8 Attention 커널 최적화: Early PV Hand-off와 Next-unit Prefetch
현재글 : [sglang] CuTe-DSL 정규화 커널 최적화: 워프당 여러 청크 처리로 성능 향상
다음글 [flashinfer] NVFP4 GEMM 성능의 한계를 넘어서: L2 Weight Prefetch를 통한 파이프라인 최적화