Perplexity에서의 CuTeDSL

Perplexity는 사내 런타임 최적화 서빙 엔진(ROSE)을 사용하여 임베딩부터 조 단위 파라미터 LLM에 이르는 다양한 모델을 서빙합니다. Sonar, 검색, 임베딩을 포함한 모든 API를 뒷받침하는 당사의 런타임은 연구원들이 가장 적합한 모델과 아키텍처를 탐색하고 최고의 후보를 신속하게 프로덕션에 배포할 수 있도록 구축되었습니다. 이를 달성하기 위해 우리는 최첨단 아키텍처에 대한 추론을 빠른 속도로 구현, 배포 및 최적화할 수 있어야 합니다. 이 블로그 게시물에서는 추론 엔진에 통합된 CuTeDSL 커널이 어떻게 NVIDIA GPU에서 모델을 최고 성능으로 끌어올리는 데 필요한 GPU 커널을 구축할 수 있게 하는지 살펴봅니다.

Perplexity의 다양한 팀이 구축한 커스텀 모델은 AI 추론 팀이 만든 ROSE 추론 엔진을 사용하여 NVIDIA Hopper 및 Blackwell GPU에서 사내 호스팅됩니다. ROSE는 원래 언어 디코딩 및 분류를 위해 커스터마이징된 Llama 모델을 더 효과적으로 서빙하기 위해 구축되었으며, NVIDIA Triton Inference Server와 호환되는 인터페이스를 갖추고 있습니다. 시간이 지남에 따라 이 엔진은 대규모 LLM뿐만 아니라 순위 지정, 분류, 점수 산정 및 기타 트랜스포머 기반 모델을 서빙하여 Perplexity의 제품을 구동하는 광범위한 팀을 지원하는 완전한 추론 엔진으로 성장했습니다. ROSE의 설계 목표는 단순함입니다. 시간이 지남에 따라 최고 성능을 달성하기 위해 노력을 투자하지만, 주요 우선순위는 우리 모델의 끊임없이 진화하는 아키텍처와 사용 사례에 적응하는 것입니다.

데이터 유형, 헤드 차원, 컴파일 시 상수, 실행 설정 전반에 걸친 커널 구성 공간.

상위 수준에서 ROSE는 모델을 클라이언트가 상호 작용할 수 있는 인터페이스에 맞게 조정하는 엔진으로 구성됩니다. 이 엔진은 재사용성이 뛰어나 요청 스케줄링 및 배칭을 처리하고, 기본 모델에서 디코딩된 토큰, 임베딩 또는 점수를 검색합니다. 장치 초기화, 프로세스 간 및 노드 간 통신 채널, 가중치 로딩은 모두 엔진이 담당합니다. LLM의 경우, 엔진은 배칭, 청킹, 샘플링을 담당합니다. 또한 전체 및 선형 어텐션 모두에 대한 KV 스토리지 할당과 접두사 매칭도 책임집니다. 임베딩 엔진은 온라인 배칭만 수행하므로 더 간단합니다. 실행은 엔진이 가속기에서 모델이 실행되는 동안 다음 작업 라운드를 병렬로 준비하도록 최적화되어 있습니다.

ROSE는 모델링을 위한 다양한 커스텀 레이어 모음을 제공하며, CuTeDSL, Triton, CUDA, CUTLASS 및 cuBLAS로 구축된 GPU 커널을 래핑합니다. 행렬 곱셈 및 어텐션 구현을 위해 당사는 주로 CUTLASS 및 cuBLAS의 NVIDIA 커널에 의존합니다. 임베딩, 정규화, MoE 및 활성화 커널을 포함한 파이프라인의 나머지 부분은 ROSE의 커널 라이브러리의 일부입니다. 추론을 실행하는 하드웨어를 최상으로 지원하기 위해 ROSE는 MoE 라우팅, 디스패치 및 결합을 위한 자체 장치 간 통신 기본 요소를 갖추고 있습니다. 다른 커널들은 비교적 간단하지만, 실행되는 모든 다양한 숨겨진 차원(hidden dimensions)과 장치에서 최고 성능을 달성하려면 철저히 전문화되어야 합니다. 또한, 컴파일 타임 전문화를 통해 다양한 모델의 모든 매개변수와 설정을 수용합니다(RMS 정규화의 가중치 편향 유무, 활성화 함수 등). 결과적으로, 우리는 이들을 구현하기 위해 DSL을 사용하기로 결정했습니다. 초기에 ROSE는 Triton을 사용하여 구축되었으나, 이후 GPU 프로그래밍을 위한 기본 환경으로 CuTeDSL을 채택했습니다.

수학적 관점에서 추론 커널은 대개 단순합니다. 따라서 커널을 개발할 때 복잡성은 일반적으로 어떤 계산을 수행할지에 대한 표현이 아니라 하드웨어가 정확히 어떻게 수행해야 하는지에 있습니다. 최고의 커널을 구축하려면 기본 하드웨어에 대한 적절한 추상화를 제공하는 언어가 필요합니다. 올바른 언어를 사용하면 계산을 수행하기 위한 상위 수준 작업을 지정하고, 광범위한 구성 스페이스 전체에서 이를 공격적으로 전문화하고 최적화할 수 있습니다.

CuTeDSL은 CuTe 레이아웃 대수와 MLIR을 기반으로 구축된 NVIDIA GPU용 Python 기반 DSL입니다. 이 언어는 저수준 하드웨어 기본 요소에 대한 제어 권한을 유지하면서 최적화된 PTX로 JIT(Just-In-Time) 컴파일됩니다. 개발 속도를 저해하는 긴 컴파일 시간 없이 CUTLASS C++와 동일한 기능과 성능을 제공하는 것을 목표로 합니다. 우리가 조사하거나 작업한 언어(CUDA, cuTile, TVM IR 및 Triton 포함) 중에서 우리는 이것이 커널 개발에 가장 적합하다고 판단합니다.

전반적으로 Triton이나 cuTile만큼 상위 수준이거나 간결하지는 않지만, CuTe 레이아웃 대수와 그 구성 요소(cute.Tensor를 통해 노출됨)가 합리적인 대체제임을 확인했습니다. 또한 CuTeDSL은 하드웨어 기본 요소에 대한 제한 없는 액세스를 제공하여 모든 엣지 케이스에서 최고 성능을 달성하는 데 필요한 수준의 제어를 제공합니다. CUDA의 경우 컴파일 타임 전문화 기회는 제한적이었습니다. 매개변수를 추가할 때마다 컴파일 시간에 불합리한 비용이 발생하거나 빌드 시스템의 복잡성이 증가했기 때문입니다. 반면, JIT 컴파일 DSL을 사용하면 반복 속도를 희생하지 않고도 더 공격적으로 전문화할 수 있습니다. CUDA를 중간 단계로 사용하는 TVM IR과 달리, 우리는 PTX로 로딩되는 MLIR 기반 파이프라인이 CUDA를 중간 단계로 배출하는 파이프라인보다 더 강력하고 디버깅하기 쉽다는 것을 발견했습니다. 또한 PTX 명령어를 래핑하는 커스텀 연산을 사용하여 기본 하드웨어의 모든 기능을 완전히 활용할 수 있음을 확인했습니다.

CUDA 커널의 컴파일 비용 중 상당 부분은 일반적으로 커널의 템플릿 확장 단계에서 발생합니다. PTX 단계 이후의 최적화 및 컴파일은 몇 초 만에 완료될 수 있기 때문입니다. 컴파일 타임 매개변수가 많을 경우, 이 컴파일 단계는 복잡한 커널의 경우 수 분으로 늘어날 수 있습니다. 이러한 병목 현상은 C++ 템플릿 처리가 순차적으로 이루어지며, 병렬 컴파일 기회가 IR 및 PTX 수준 이후에야 발생한다는 사실로 인해 더욱 강조됩니다. 이를 해결하기 위한 기술들이 존재하지만, 템플릿 인스턴스화를 별도의 컴파일 단위로 이동하는 방식은 빌드 시스템을 크게 수정해야 하며 다루기 힘든 매크로를 사용해야 합니다. 반면, JIT로 컴파일되는 CuTeDSL 커널은 복잡한 빌드 규칙이 필요하지 않으며 컴파일 시간을 저해하지 않습니다. 개발 주기는 구성 스페이스의 단일 지점에 집중할 수 있고, 이후 정밀한 단위 테스트를 통해 커널의 전체 설계 스페이스를 다룰 수 있기 때문입니다.

JIT 컴파일이 비용을 개발 시간에서 런타임으로 옮기기는 하지만, 그 오버헤드는 충분히 수용 가능합니다. 모델이 배포되면 전체 구성 스페이스 중 일부 포인트만 활성화되므로 추론 엔진의 시작/웜업 중에 초 단위의 오버헤드가 발생합니다. 반면, 커널의 모든 구성을 사전 컴파일하면 수 분이 소요될 수 있습니다. 또한 JIT 컴파일된 아티팩트를 캐싱하여 오버헤드를 더욱 완화할 수 있습니다.

CUDA의 무거운 컴파일 비용은 CuTe 및 CUTLASS의 경우와 같이 템플릿 기반 디스패치 및 전문화가 사용될 때 오류의 기원을 추적하기 어렵다는 점으로 인해 가중됩니다. DSL을 사용하면 오류를 생성한 코드 라인을 가리킬 수 있어 디버깅이 크게 쉬워집니다. 또한 중간 MLIR 단계를 통해 생성된 코드를 DSL 구문과 매핑하거나 생성된 PTX에 첨부된 디버깅 정보를 활용하는 것이 훨씬 간단합니다. DSL의 이러한 측면 덕분에 우리는 커널을 개발할 때 생산성이 훨씬 향상되었습니다.

커널은 모델 간에 많은 유사점을 공유하지만 여전히 성능에 영향을 줄 수 있는 다양한 구성 노브(knob)를 갖추어야 합니다. 관련 매개변수는 다음과 같습니다.

숨겨진 차원/헤드 차원/어텐션 헤드 수

입력/출력 데이터 타입

양자화 체계

컴파일 타임 상수(편향, 활성화 함수 등)

컴파일 타임 전문화는 성능과 디버깅 모두에 유용합니다. 특히 GPU와 같은 파이프라인 프로세서에서는 항상 실행되거나 전혀 실행되지 않는 분기보다 존재하지 않는 분기가 더 낫습니다. 또한 특정 구성 노브 세트에 대해 커널을 컴파일하면 PTX/SASS 명령어가 줄어들어 조사 중에 cuda-gdb를 사용할 때 분석 및 디버깅이 더 쉬워집니다.

주어진 배치 크기 및 구성에 대해 DRAM으로의 예상 트래픽과 산술 연산 횟수를 추정할 수 있습니다. 이후 최적의 그리드 할당, 동기화 및 특정 하드웨어 기능 사용에 대해 사례별로 최적의 결정을 내릴 수 있습니다. CuTeDSL을 사용하면 허용 불가능한 컴파일 타임 오버헤드를 발생시키지 않으면서 심층적인 전문화를 위해 구성 요소를 정의하고 재사용할 수 있습니다.

일반적으로 추론 커널의 성능은 두 가지 별도의 사용 사례에 최적화되어야 합니다.

프리필: 작은 배치 크기(소수의 청크, 시퀀스), 시퀀스당 많은 토큰. 임베딩과 LLM의 입력 토큰 모두에 대해 처리량을 최대화해야 함. 지연 시간은 수십에서 수백 밀리초.

디코드: 소규모에서 대규모 배치 크기, 시퀀스당 1개(추측성 디코딩 없음)에서 수 개의 토큰(드래프트-타겟 추측, 다중 토큰 예측). 지연 시간은 커널 실행 지연 시간과 근접.

GEMM 및 어텐션의 경우, 프리필 단계는 많은 수의 토큰에서 실행되며 성능은 장치의 연산 처리량에 의해 제한됩니다. 디코딩은 소규모 배치 크기에서 동일한 작업을 포함하며, 실행 시간은 DRAM에서 스트리밍 멀티프로세서 및 텐서 코어로 활성화 값과 가중치를 로드하는 비용이 지배하며, 대신 장치의 메모리 대역폭에 의해 제한됩니다. 배치 크기나 배치 내 전체 토큰 수가 비교적 명확한 지표이지만, 메모리 대역폭은 모델의 숨겨진 차원과 활성화 데이터 타입에 의해서도 결정되기 때문에 프리필과 디코드 간의 전환 임계값은 더 미묘합니다. 결과적으로 알고리즘 차이나 그리드/블록 할당이 어느 한 쪽에 이점이 있다면 커널은 프리필과 디코드 구현으로 분할됩니다.

프리필과 디코드 커널 간에 보편적으로 다른 매개변수는 블록 내 워프(warp)의 수(스레드 수, 워프당 32스레드인 워프 크기로 나눔)입니다. 이 매개변수를 컴파일 타임 상수로 만드는 것이 레지스터 할당에 도움이 되므로 바람직합니다. CUDA에서는 일반적으로 __launch_bounds__를 통해 스레드 수를 지정하고 템플릿 매개변수로 워프 수를 제공하여 이를 수행합니다.

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

블록 내 모든 스레드에서 사용할 수 있는 레지스터 수는 Hopper 및 Blackwell에서 64K로 하드웨어 제약이 있습니다. 스레드 수를 미리 알고 있으면 컴파일러는 각 스레드에 할당할 수 있는 레지스터 수를 정확히 알 수 있으며 더 공격적으로 최적화된 커널을 성공적으로 컴파일할 수 있습니다. 그러나 이는 이제 커널을 다른 워프 크기에 대해 다시 컴파일해야 하므로 컴파일 시간 비용이 발생합니다. CuTeDSL을 사용하면 전체 구성 스페이스에 대한 컴파일 비용을 미리 지불할 필요가 없으므로 더 이상 문제가 되지 않습니다. DSL은 또한 블록 크기가 상수 표현식일 것으로 예상하며, CUDA의 __latch_bounds__가 의미하는 최적화를 암묵적으로 수행합니다.

워프 수를 상수로 설정하는 것 외에도, prefill과 decode 사이에서 이를 달리하는 것이 바람직합니다. GPU의 메모리 대역폭을 완전히 활용하려면 블록당 항상 많은 수의 워프를 할당하는 것만으로는 충분하지 않습니다. 사용 가능한 Streaming Multiprocessors(SM)를 가능한 한 많이 사용해야만 전체 대역폭을 달성할 수 있기 때문입니다. 따라서 주어진 토큰 수와 알려진 데이터 타입 및 히든 차원에 대해, 먼저 각 스레드에 대한 이상적인 벡터 차원(bfloat16의 경우 8, float32의 경우 4)을 선택하여 생성된 코드가 128비트 load/store 명령어(Blackwell에서는 256비트)를 사용하도록 합니다. 벡터 크기를 알면 커널에 필요한 워프 수를 계산할 수 있습니다. 먼저 SM 수를 최대화하려고 시도합니다. 워프 수가 SM 수보다 적으면 블록당 워프 1개를 선택합니다. 이는 일반적으로 커널의 decode 변형입니다. 워프가 충분하면 워프 수를 8, 16 또는 32까지 점진적으로 늘리며, 그 이후에는 일반적으로 성능 향상이 더 이상 보이지 않는 상한을 벤치마킹합니다.

Triton이 이러한 최적화 중 일부를 내부적으로 수행했지만, 상위 수준 컴파일러의 단점에 얽매이지 않고 CuTe 레이아웃 대수를 신중하게 사용하여 유사한 결과를 얻을 수 있음을 발견했습니다. 특히 서브 워프 수평 축소(horizontal reduction)가 포함된 경우, CuTeDSL이 최적의 코드를 생성하는 데 필요한 세밀한 제어 기능을 제공함을 확인했습니다.

일부 상황에서는 가능한 한 많은 블록에 작업을 분산하기 위해 cooperative_groups::this_grid().sync()를 사용하여 참여하는 모든 블록을 동기화할 수 있습니다. 제약 사항이 있습니다. 일단 블록이 실행을 시작하면 SM에서 스케줄 해제될 수 없으므로, 전진 진행을 보장하기 위해 블록 수는 장치에서 사용 가능한 SM 수를 초과할 수 없습니다. 이는 언제든지 모든 블록이 동기화와 관련된 전역 메모리 위치를 업데이트하고 폴링할 수 있도록 보장하기 위함입니다. 이 제한은 cudaLaunchCooperativeKernel에 의해 강제됩니다. 그리드 배리어의 비용은 Hopper에서 약 3us이며, 일부 상황에서는 암시적 메모리 배리어를 통해 이전의 모든 메모리 작업을 플러시하는 비용이 추가됩니다. 전반적으로 이는 커널을 두 개의 실행으로 분할하는 오버헤드와 비슷하며, 성능에 민감한 설계 결정을 내리게 합니다. Hopper부터는 PDL(Programmatic Dependent Launch)이 실행 오버헤드를 줄일 수 있지만, 디코드 커널의 경우 오버헤드를 피하는 것이 가치가 있습니다.

워크로드 중 하나가 SM 수보다 많은 블록을 필요로 하는 경우, 커널 실행의 암시적 배리어에 의존하여 커널을 분할해야 합니다. 이는 주로 프리필의 경우입니다. 모든 워크로드에 이상적인 블록 수가 그리드에 들어맞는 경우, 작은 수의 토큰에서 작동하는 디코드 커널의 경우와 같이 그리드 배리어가 선호됩니다.

CuTeDSL에서는 호스트 측 실행 코드에서 이러한 결정을 내릴 수 있으며, 아래 예시와 같이 필요에 따라 커널을 구성하고 그리드 배리어와 인터리브할 수 있습니다. 이러한 하이브리드 전략을 사용하면 디코드 커널을 최소 지연 시간에 도달하도록 최적화하고 프리필 커널을 최대 처리량에 최적화할 수 있습니다. CUDA 기반 솔루션과 달리 이 로직을 더 간결하고 컴파일 타임 오버헤드를 줄이면서 정의할 수 있습니다.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

아래 자세히 설명된 MoE all-to-all 커널의 경우, 그리드 배리어 유무에 관계없이 커널을 구축하고 컴파일 타임에 최적화하는 유연성을 확보함으로써 CuTeDSL을 통해 상당한 성능 향상을 달성했습니다.

CuTeDSL에서 처음 구축한 커널은 작은 차원에 최적화된 전문화된 RMS 정규화 커널이었습니다. 최근 Qwen 모델의 셀프 어텐션 레이어에 사용된 QK 정규화는 64, 128 또는 256개 요소에 대해 헤드별로 적용됩니다. 수천을 초과하는 숨겨진 차원에서 작동하며 공유 메모리와 레지스터를 통한 수평 축소가 필요한 일반 RMS 정규화와 달리, QK 정규화는 워프 간의 축소만 필요하며 여러 헤드를 단일 워프에 압축할 가능성이 있습니다. 이것이 바로 모든 헤드 차원과 데이터 타입에 대해 최상의 토큰-워프 할당을 찾기 위해 컴파일 타임에 공격적으로 최적화할 수 있는 사용 사례가 되었습니다.

RMS 정규화는 벡터의 제곱 노름을 계산하기 위해 수평 합계가 필요하지만 상당히 간단한 작업입니다.

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

어텐션 헤드는 일반적으로 작고 스레드당 작업은 매우 벡터화되어 있으므로 단일 워프가 제곱 노름을 계산할 수 있습니다. 메모리 대역폭을 최대화하려면 각 스레드가 128비트의 입력을 처리해야 하며, 이는 8개의 bfloat16 값으로 변환됩니다. 결과적으로 하프 워프는 헤드를 처리하기에 충분하며, 제곱 노름은 shuffle_sync_bfly로 수행되는 수평 축소를 사용하여 결정할 수 있습니다. 헤드 차원과 데이터 타입에서 시작하여 워프당 헤드 수를 결정하고 컴파일 시점에 인수와 축소 라운드 수를 지정할 수 있습니다.

그 후 자동으로 벡터화된 로드/스토어를 활용하고 워프 축소를 위해 낮은 수준의 아키텍처 특정 작업을 사용하여 비교적 적은 코드 라인으로 전문화된 커널을 구축할 수 있습니다. 커널은 작업할 헤드(Q 또는 K)를 결정하고 적용할 해당 가중치를 식별합니다. 헤드에서 전체 벡터를 로드하고 먼저 로컬에서 합산합니다. 그 후 헤드당 스레드 수에 따라 제곱 노름이 계산됩니다. 마지막으로 자동으로 벡터화된 RMS 정규화가 결정되어 메모리에 다시 기록됩니다. WEIGHT_BIAS 및 EPS와 같은 모델 특정 매개변수는 상수 폴딩됩니다.

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

일반 RMS 정규화 커널과 비교하여 이 전문화된 변형은 프리필과 디코드 모두에 대해 2~3배 더 나은 처리량과 지연 시간을 제공합니다.

MoE 디스패치 및 결합은 입력 토큰을 라우팅된 전문가에게 섞고, MoE GEMM 커널이 관련 행렬 곱셈을 수행한 후 가중 누적을 수행합니다. 다양한 제품이 다양한 크기의 MoE 모델에 의존하기 때문에 성능은 특히 관심사입니다. 배포 유형에 따라 최적화할 다른 커널을 구분합니다.

TP=1: 단일 GPU용 로컬 커널. 디스패치는 토큰을 관련 전문가에게 섞고 결합은 토큰을 누적합니다. 주로 H100의 Qwen3 30B 배포에 사용됩니다.

DP=1, TP ≠ 1: 텐서 병렬, 단일 노드 커널. 입력이 복제되고, 디스패치가 현재 순위에 할당된 토큰을 선택하며, 결합은 모든 순위에 걸쳐 토큰을 누적합니다. NVLink 또는 NVLS를 통한 누적. 대형 프리필러에 사용됩니다.

DP ≠ 1, NVLink: 선택적 텐서 병렬 처리가 있는 데이터 병렬 커널. 각 순위는 특정 전문가를 호스팅하는 다른 모든 순위로 토큰을 보냅니다. 라우팅 정보가 집계되고 토큰은 NVIDIA NVLink를 통해 전송됩니다. 데이터 병렬 H200 또는 GB200 배포에 사용됩니다.

DP ≠ 1, InfiniBand: NVLink와 NVIDIA Quantum InfiniBand 모두에서 작동하는 피어 투 피어 데이터 병렬 및 텐서 병렬 커널. 다중 노드 Hopper 배포에 사용됩니다.

로컬 커널은 라우팅 테이블을 기반으로 각 전문가에게 할당된 토큰 수를 계산하고 각 전문가에게 속한 토큰 목록 내에서 각 토큰에 고유한 인덱스를 할당합니다. 라우팅 단계 후, 토큰은 MoE GEMM 커널로 전달되기 위해 선택적으로 패딩된 버퍼에 복사됩니다. 토큰 위치를 사용하여 결합 커널은 레지스터에서 누적을 수행하고 최종 토큰을 출력 버퍼에 기록합니다.

디스패치 단계는 프리필과 디코드에 따라 다릅니다. 토큰 수가 적을 때는 읽을 라우팅 정보가 거의 없으므로 단일 블록에서 처리하고 그리드의 나머지에 신호를 보내며 개별 블록이 토큰 청크를 복사합니다. 토큰 할당은 공유 메모리의 카운터를 원자적으로 증가시켜 결정됩니다. 프리필은 더 많은 라우팅 정보를 처리해야 하므로 라우팅은 별도의 커널에서 여러 블록으로 분할됩니다. 인덱스는 먼저 공유 메모리의 카운터를 사용하여 로컬에서 할당되고, 각 토큰의 위치를 결정하기 위해 전역 메모리의 모든 블록에 걸쳐 집계됩니다. 이 커널의 그리드는 그리드 단위 배리어를 사용할 수 있도록 SM 수에 고정됩니다. 후속 커널은 사용 가능한 메모리 대역폭을 완전히 활용하기 위해 훨씬 더 많은 수의 블록과 워프에 걸쳐 시작되어 복사를 수행합니다.

결합 커널은 벡터화된 누적을 수행하고 라우팅 정보와 디스패치에 의해 계산된 오프셋에 의존하므로 비교적 간단합니다. CuTeDSL은 계산을 모든 데이터 타입과 숨겨진 차원에 대해 효율적인 코드로 컴파일되는 간단한 CuTe 텐서를 통해 표현할 수 있었기 때문에 여기에서 가치가 입증되었습니다.

텐서 병렬 배포에서 라우팅 정보와 입력 토큰은 복제됩니다. 로컬 디스패치 커널은 현재 순위에 위치한 전문가에게 할당된 토큰만 라우팅하고 선택하도록 약간 변경하여 완전히 재사용할 수 있습니다. 일부 매개변수와 컴파일 타임 전문화를 통해 로컬 디스패치 커널이 재사용됩니다.

순위 수가 일반적으로 토큰이 라우팅되는 전문가 수보다 적기 때문에 평균적으로 각 토큰은 각 순위로 라우팅되는 복제본 중 최소 하나를 가지고 있다고 가정할 수 있습니다. 이 상황에서 로컬 결합 커널은 토큰의 로컬 복사본을 집계하는 데 재사용될 수 있으며, 이후 참여 순위에 걸쳐 합계를 구하기 위해 all-reduce가 이어집니다. 프리필 커널의 경우 별도의 커널을 시작하지만 디코드는 두 작업을 단일 실행으로 융합합니다.

All-reduce의 경우 두 순위만 참여하면 일반 NVLink를 사용하고, 4개 이상의 순위가 관련될 때는 처리량이 10~20% 증가하는 NVLS로 전환합니다. DSL 덕분에 MoE와 일반 all-reduce 커널 간의 축소 로직과 피어 버퍼 교환의 대부분을 공유할 수 있습니다. NVLS의 경우 multimem.ld_reduce 및 multimem.st를 통한 DSL 작업이 필요하므로 하드웨어에 대한 긴밀한 제어가 유익한 상황입니다. 축소 구현 자체는 주로 NCCL 테스트의 all_reduce 예제를 따릅니다.

데이터 병렬 커널은 NVLink에 의존하여 여러 장치의 토큰을 집계합니다. 디스패치 단계에서 토큰을 교환하기 전에 피어는 먼저 로컬에서 라우팅하여 보내는 각 토큰에 대해 각 순위의 대상 오프셋을 식별합니다. 그 후 충돌을 피하기 위해 각 발신자에게 할당된 수신 버퍼로 토큰을 복사하기 시작합니다. 완료 신호를 보내기 위해 토큰별 릴리스-스토어 쓰기가 사용됩니다. 수신기 측에서는 MoE GEMM을 위해 모든 토큰을 단일 버퍼로 압축할 수 있도록 라우팅 정보가 집계됩니다. 수신기는 쓰기 완료를 기다리기 위해 획득-로드(acquire-load)를 사용하여 피어가 기록한 플래그를 폴링합니다.

결합 측에서 토큰은 NVLink를 통해 소스 순위로 복사되며 토큰별 플래그는 수신을 나타냅니다. 결합 커널은 로컬 및 단일 노드 변형과 유사하며 주소 지정 및 일부 추가 동기화 로직에 약간의 차이가 있습니다.

이 커널을 구현하기 위해 시스템 및 장치 수준에서 범위가 지정된 릴리스-획득 로드 및 스토어를 노출하는 PTX 래퍼 라이브러리를 구축했습니다. 동기화 플래그의 세분성을 개선하고 NVLink를 통해 토큰을 보내고 받기 위해 더 많은 블록을 할당함으로써 이전 CUDA 커널보다 2배의 개선을 달성했습니다.

우리의 피어 투 피어 커널은 이전에 ConnectX-7 InfiniBand 어댑터와 페어링된 Hopper GPU에서 최첨단 성능을 달성했습니다. 이후 더 나은 성능 향상을 위해 결합 커널을 CuTeDSL로 포팅했습니다. CUDA 버전은 이미 컴파일 타임 전문화에 크게 의존했기 때문에 우리 저장소에서 컴파일 타임 병목 현상이 발생하여 일부 매개변수를 생략해야 했습니다. JIT 컴파일 DSL로 마이그레이션함으로써 더 많은 매개변수를 전문화하고 그리드 할당을 조정하여 프리필과 디코드 간의 차이를 더 도입하여 지연 시간을 약 3% 향상시켰습니다.

ROSE 커널의 관점에서 Blackwell 지원은 아키텍처 불가지론(architecture-agnostic)과 아키텍처 특정 커널이라는 두 가지 범주로 나뉩니다. 텐서 코어에 의존하는 아키텍처 특정 커널은 적응하기 위해 명시적인 작업이 필요합니다. 그러나 DSL 사용 덕분에 아키텍처 불가지론 커널은 사소하게 적응합니다.

Blackwell GPU는 Hopper의 132개와 달리 최대 160개의 SM을 가지고 있습니다. 대부분의 커널에 이것은 컴파일 타임 상수이며 호스트 측 실행 코드의 로직은 실행할 최적(프리필 대 디코드) 커널을 결정할 때 이를 고려합니다. 또한 Blackwell은 256비트 로드 명령을 도입하여 벡터 너비(CuTe 텐서의 마지막 차원)를 늘려 메모리 대역폭 활용도를 높임으로써 쉽게 배출될 수 있습니다.

Blackwell은 상당수의 양자화 체계(MXFP8, MXFP4, NVFP4)에 대한 하드웨어 지원을 도입합니다. 이에 대한 지원은 주로 FP8 블록별 스케일링을 위해 기존 커널을 조정하여 구현되었습니다. JIT 지원은 컴파일 타임 구성 스페이스를 폭발시키지 않고 새로운 모드를 지원할 수 있으므로 여기에서 매우 중요합니다. 과거에 우리는 Triton이 워프의 절반이나 쿼터에 걸쳐 몇 개의 요소에 대한 수평 축소를 표현하는 데 어색하다는 것을 알았습니다. NVFP4 및 MXFP8/MXFP4에 대해 각각 16개 또는 32개 요소에 걸쳐 절대 최대값을 계산할 때 최적의 성능을 달성하려면 축소에 대한 더 세밀한 제어가 중요했습니다.

비교적 새로운 언어로서 LLM은 사전 훈련 중에 CuTeDSL 코드에 노출될 기회가 거의 없었습니다. 또한 활성화 및 정규화에 최적화된 커널을 포함한 현재 우리 코드의 스타일은 TMA 및 아키텍처 특정 행렬 곱셈 명령에 최적화된 기존 CUTLASS 또는 Quack 예제와 약간 달랐습니다. LLM이 커널 작성에 도움을 줄 수 있는 지점에 도달하는 것은 점진적인 과정이었습니다.

에이전트에게 커널 변환을 의존할 수 있기 전에 일관된 패턴을 따라 커널과 유틸리티 함수의 초기 코퍼스를 수동으로 구축했습니다. 코퍼스가 진행됨에 따라 상당한 수동 재작업을 거쳐 CUDA에서 CuTeDSL로 대략적인 변환을 수행할 수 있었습니다. 특히 모델은 포인터-텐서 변환과 벡터화에 어려움을 겪었으며, 테스트 컴파일을 목표로 할 때 타입 오류를 수정하려다 막히는 경우가 많았습니다. 모델은 스레드와 그리드 동기화를 혼동하며 커널의 그리드 기반 특성에 대한 이해가 거의 없었습니다. 가장 실망스러운 점은 자체 CuTeDSL 대체제로 구축하는 대신 존재하지 않는 CUDA 특정 라이브러리의 번역을 환각하려는 유혹이 있었다는 것입니다. 시간이 지나 CuTeDSL 커널 코퍼스가 성장함에 따라 LLM이 대부분의 스켈레톤 코드를 구축하고 빈칸을 채울 수 있어 우리가 상위 수준 설계에 집중할 수 있게 되었습니다.

CUDA와 비교하여 우리는 CuTeDSL이 더 나은 컴파일 타임 성능을 제공하여 더 많은 매개변수를 전문화하고 성능 향상을 달성할 수 있음을 확인했습니다. Triton에 비해 특정 추상화가 부족하여 커널이 약간 더 장황하지만, 하드웨어에 대한 세밀한 제어를 통해 더 빠른 커널을 구축할 수 있으므로 복잡성은 정당화된다고 생각합니다.

이 글을 작성하는 시점에서 CUTLASS C++와 CuTeDSL 커널 간의 성능 격차는 미미하며, 고성능 GEMM 커널을 장치 간 통신, 정규화 또는 활성화와 구성 및 융합할 수 있는 흥미로운 기회가 열리고 있습니다. 이 DSL을 사용함으로써 우리는 장치에서 더 많은 성능을 추출하여 임베딩부터 LLM까지 모든 워크로드에 걸쳐 처리량을 개선할 수 있을 것입니다.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AI에서의 추론

왜 CuTeDSL인가?

개발자 경험

대규모 구성 스페이스

프리필 및 디코딩 특화

그리드 동기화

CuTeDSL 커널

QK 정규화

전문가 혼합(MoE) 디스패치/결합

로컬

단일 노드

데이터 병렬

피어 투 피어

Blackwell 지원

AI 지원 커널 생성

결론 및 향후 과제

참고 문헌

인용