NVIDIA DGX Spark: MoE 커널 최적화, 추측 디코딩

1. 전문가 계산과 출력층의 읽기 부담을 줄인다

평가 대상은 Inclusion AI의 Ling-3.0-flash-fp4다. 여러 전문가 가운데 일부만 골라 계산하는 MoE 모델로, 전체 1,240억 파라미터 중 토큰마다 51억 개를 활성화한다. DGX Spark에서 요청을 하나씩 처리하는 배치 크기 1로 분석했더니, 토큰 생성은 가중치를 읽는 속도에 묶여 있었다.

주요 비용은 전문가 계산과 어휘별 점수를 구하는 출력층이었다. NVIDIA는 계산을 실행하는 커널, GPU에서 연산을 수행하는 프로그램을 교체하고 출력층의 가중치를 압축해 단계당 읽기 부담을 낮췄다. 여기에 여러 토큰을 미리 제안하고 한꺼번에 검증하는 추측 디코딩을 더해 생성에 필요한 단계 수도 줄였다.

2. 커널 교체와 FP8 출력층을 함께 적용한다

첫 방법은 기본 FlashInfer MXFP4, 블록별 4비트 부동소수점 형식인 전문가 계산 경로를 Humming 커널로 바꾸는 것이다. Humming은 행렬 크기·데이터 형식·GPU에 맞춰 커널을 실행 중 컴파일하고 저장해 재사용한다. 회사 설명에 따르면 개별 커널 컴파일은 수 초에 끝나며, 기본 경로는 97개 대상을 컴파일하는 데 수십 분이 걸린다.

두 번째 방법은 저장된 모델에서 BF16 형식, 16비트 부동소수점 데이터 표현인 출력층만 구성 시점에 더 작은 FP8 형식, 8비트 부동소수점 데이터 표현으로 바꾸는 것이다. 출력층은 생성 토큰마다 어휘 157,184개에 대응하는 가중치 전체를 읽는데, FP8로 바꾸면 이 층의 읽기량이 절반이 된다. 나머지 모델의 양자화 설정은 유지한다.

기준 구성은 가중치 4비트·활성값 8비트 설정에 출력 점수를 FP32로 내보내는 경로를 사용했다. FP8 출력층은 이 경로 대신 적용한다.

이 조합은 단계당 비용과 생성 단계 수를 줄이는 대신, 처음 쓰는 커널을 컴파일해 저장하는 절차와 초안 검증 계산을 추가한다. 도입 시에는 토큰 생성의 이득과 이 준비·검증 비용을 함께 살펴야 한다.

3. 내장 예측 층이 초안을 만들고 본 모델이 검사한다

Ling-3.0-flash의 NEXTN 예측 층, 다음 토큰 후보를 예측하는 층은 추가 학습이나 다운로드 없이 초안 생성에 쓸 수 있다. 초안이 토큰 세 개를 차례로 제안하면 본 모델이 네 위치를 함께 검사하고, 앞에서부터 맞는 부분을 받아들인다. 한 번의 검사로 여러 토큰을 확정하는 방식이다.

회사 발표 기준 검증 한 번의 비용은 일반 생성 한 단계의 약 2.26배였다. 배치 크기 1에서는 이전 토큰의 저장 정보를 읽는 비용보다 전문가 활성화가 더 큰 부담이었다. 일반 생성은 토큰 하나에 전문가 8개를 쓰지만 네 위치를 검증할 때는 최대 32개를 활성화한다.

FP8 출력층과 추측 디코딩을 함께 쓰려면 초안 모델이 본 모델의 출력층 가중치뿐 아니라 모듈 전체를 공유하도록 구성해야 한다. 공유되지 않으면 FP8 가중치를 BF16으로 해석해 한 번에 받아들이는 길이가 1.00토큰으로 떨어지고, 최적화 전보다 느려질 수 있다.

4. 합성 시험에서는 단계별 누적 효과가 컸다

NVIDIA는 DGX Spark에서 입력 32K 토큰, 출력 2K 토큰의 합성 작업에 각 방법을 차례로 추가했다. 회사 발표 기준 세 방법을 모두 적용한 처리량은 초당 102.88토큰으로, 기준 구성보다 185.1% 늘었다.

표: 합성 작업의 누적 최적화 결과

조건: NVIDIA DGX Spark, Ling-3.0-flash-fp4, 입력 32K·출력 2K 토큰, 배치 크기 1, 탐욕적 디코딩, 추측 구성 steps=3 / topk=1 / draft_tokens=4

단위: 토큰당 시간 ms, 처리량 토큰/초, 증가율 %

이 시험은 가장 가능성 높은 토큰을 골라 반복적인 출력을 생성했다. 추측 디코딩은 한 번에 평균 3.90토큰을 받아들여 여러 토큰을 확정했다.

5. 전체 문제 평가에서는 처리량이 42.5% 늘었다

GSM8K 전체 1,319문 평가에서는 예시 5개와 풀이 과정 생성을 사용했다. 두 구성의 샘플링, 정답 판독, 무작위 시드, 메모리 할당을 맞췄고, 회사 발표 기준 평균 출력 길이는 기준 567토큰, 최적화 후 607토큰이었다.

표: GSM8K 전체 평가의 기준 구성과 최적화 구성

조건: NVIDIA DGX Spark, Ling-3.0-flash-fp4, GSM8K 1,319문, 배치 크기 1, 예시 5개·풀이 과정 생성, 온도 0.6·누적 확률, 토큰 확률을 더해 정하는 기준 0.95·상위 토큰, 확률이 높은 후보 토큰 29·최대 출력 32,768토큰

회사 발표 기준 정확도에는 통계적으로 검출 가능한 변화가 없었다. 이 평가에서는 확률에 따라 토큰을 선택했고, 출력 길이의 중간값은 약 310토큰으로 합성 시험보다 짧았다. 한 번에 받아들인 토큰도 평균 2.68개여서, 두 시험의 향상 폭 차이를 설명한다.

도입 판단의 핵심은 합성 시험의 최고 향상률보다 배포 환경을 닮은 입력과 출력으로 세 최적화의 결합 효과를 확인하는 데 있다.