AWS Trainium2에서 GLM-5.3-Flash를 위한 NKI 레이어 메가커널: 디코드 스텝의 시간 사용처와 우리가 변경한 것

저희는 AWS Trainium2에서 GLM-5.3-Flash (320B MoE, 1M 토큰 컨텍스트)를 하나의 trn2.48xlarge에서 서비스하는 추론 엔진을 작성하여 오픈 소스화했습니다. 작업의 대부분은 디코드(decode) 과정에 집중되었으며, 제가 가장 흥미롭게 발견한 부분은 모델의 실제 수학 연산이 디코드 스텝에서 차지하는 비중이 매우 작았다는 점입니다. 초기 빌드에서는 배치 1의 하나의 KDA (선형 어텐션) 레이어에서 재귀(recurrence)가 전체 어텐션 시간의 2%에 불과했습니다. 나머지 시간은 mHC 사전 투영(pre-projection) (36%), XLA glue (33%), 전치(transposes) (10%), 스트림 믹스(stream mix) (10%), 그리고 기타 작은 연산들 (8%)이 차지하고 있었습니다. 배치 8의 디코드 스텝은 메모리 하한선인 6.1ms에 비해 24.2ms가 소요되었고, 프로파일링 결과는 그 이유를 보여주었습니다: 약 153ns의 고정 비용을 가진 작은 벡터 엔진 명령어들의 긴 체인, 모든 커널 경계마다 HBM으로의 쓰기 백(write back), 그리고 스텝당 125개의 크로스-칩 콜렉티브(cross-chip collectives) 연산이었습니다. 따라서 저희는 각 레이어 절반을 하나의 NKI 커널(희소 어텐션 레이어 전체를 아우르는 경우)로 재작성했습니다. 이 커널들은 세 가지 규칙을 따릅니다: 1. 작업 집합(working set)을 칩 내에 유지합니다. 물리적 NeuronCore는 SBUF가 28 MiB이므로, 이는 레이어 절반의 활성화 값, 재귀 상태, 그리고 비행 중인 가중치를 처리하기에 충분하며, 오직 레이어의 입력, 출력 및 스트리밍되는 가중치만이 HBM을 가로지릅니다. 2. 논리 코어의 두 물리적 코어를 모두 사용합니다. 한 프로그램이 340–400 GB/s의 가중치를 스트리밍하고, 쌍으로 사용하면 672–708 GB/s가 되므로, 가중치 스트림을 이들 사이에 분할해야 합니다. 3. 콜렉티브 연산이 다음 커널 실행을 차단(block)하는 경우, 이를 커널 내부로 이동시킵니다. XLA를 통해 실행되는 모든-축소(all-reduce)는 다음 커널이 모두 완료되기를 기다리게 만듭니다. 그러나 커널 내부에서는 독립적인 작업들이 이 과정을 중첩(overlap)할 수 있습니다. 같은 레이아웃에서 배치 8의 디코드 스텝은 저희가 10월 17일에 만든 빌드에서는 32.2ms였으나, 11월 1일에는 24.2ms로 줄었습니다. 이 변화의 대부분은 두 가지 변경 사항 덕분이었는데, 각각 전체 스텝에서 측정했을 때 배치 8 스텝당 3.2ms를 절약해 준 결합된 희소 어텐션 레이어와 KDA 및 MoE 커널이 함께 적용되면서 또 다른 3.2ms를 절약해주었습니다. 비교할 수 있는 기본 AWS 기준선은 없습니다: 최신 vLLM Neuron 플러그인 릴리스에는 GLM 모델이 등록되어 있지 않으며, NxD Inference는 유지보수 모드에 있습니다.

남은 24.2ms 중 6.1ms는 메모리 바닥(memory floor)이며, 10.0ms는 그 이상의 엔진 사용 시간(주로 벡터 엔진 명령어 체인), 4.6ms는 커널 내부 대기 시간(waits inside kernels), 2.8ms는 노출된 집합 연산(exposed collectives), 그리고 0.7ms는 유휴 시간입니다.

작동하지 않은 부분: 컴파일러가 이미 오버랩(overlapped)한 심(seams) 융합입니다. 저희는 이를 단계당 2~4ms로 책정했지만, 해당 단계가 이미 하나의 컴파일된 그래프였기 때문에 한 배치 크기에서 0.85ms를 절약했습니다.

고립된 벤치마크를 신뢰하는 경우. 밀집 FFN(Dense FFN) 커널은 단독으로 실행될 때 145µs 대신 84µs로 작동했고, 비트가 동일했습니다. 그러나 컴파일러가 새로운 호출 주변의 모든 것을 재조정했기 때문에 배치 크기 2에서 전체 단계는 55µs 느려졌습니다.

시뮬레이터를 신뢰하는 경우. 세 개의 커널이 NKI의 시뮬레이터를 통과했지만 장치(device)에서는 잘못된 답변을 내놓았습니다: PSUM 타일 슬라이스 내 두 개의 matmul 누적 그룹, 펜스(fences) 없는 커널 내부 집합 연산, 그리고 쌍을 이루는 다른 프로그램이 방금 작성한 공유 HBM을 읽는 하나의 프로그램입니다.

정확도(Accuracy): 모든 커널은 비트가 변경되어서는 안 되는 부분에서는 바이트 단위로 확인되며, 그렇지 않은 경우에는 이전 빌드의 1.05배로 오류가 제한된 고정밀 참조 값과 비교됩니다.

근사화(approximation)는 희소 어텐션 프리필(sparse-attention prefill)에서 여전히 남아있으며, 전체 선택 중 1.2%에 해당합니다: 정확한 재실행과 비교했을 때 평균적으로 다음 토큰 로그 확률을 0.012~0.021 nats만큼 이동시켰고, 448개 위치 중 3곳에서 최고 토큰이 변경되었습니다.

저희는 이 빌드에 대해 작업 수준 평가(task-level eval)를 수행하지 않았습니다. 엔드투엔드(End to end)로, SemiAnalysis의 AgentX 벤치마크(자체 실행, 지점당 1시간 소요, 공식 InferenceX 아님)에서 한 박스당 분당 최대 28.9개의 에이전트 요청을 측정했으며, 16개 에이전트에서는 중앙값 요청 디코딩 속도가 35.9 tokens/s였습니다.

동일한 트레이스(traces)로 GPU를 실행하지 않았으므로 여기에는 GPU 비교가 없습니다.

공개 고지: 저는 RunAnywhere에서 근무하며, 이 코드는 저희 소유입니다.

코드: https://github.com/RunanywhereAI/trnrt (모든 숫자를 다시 실행하는 명령어는 docs/RESULTS.md에 있습니다) 프로필 및 모든 숫자가 포함된 작성 자료: https://runanywhere.ai/blog/glm-5-3-flash-trainium2-nki-megakernels 저희의 디자인은 vLLM의 AMD 전 레이어 디코드 커널과 Hazy Research 및 Mirage의 메가커널 작업에서 영감을 받았습니다. 비판을 듣게 되어 기쁩니다. 특히 Neuron 환경에서 NKI나 XLA와 씨름해 본 분들의 의견을 기대합니다. 제출자: /u/WildCard727 [링크] [댓글]