TorchNeuron 퓨전, Trainium에서 RMSNorm 지연 시간 1/6으로 단축

• AWS Trainium 기반 TorchNeuron 시스템 레벨 프로파일링은 torch.profiler와 Neuron Explorer를 활용해 프레임워크 디스패치부터 하드웨어까지 실행을 추적함
• torch.rms_norm 또는 torch.compile을 통한 연산자 퓨전은 RMSNorm 디바이스 시간을 네이티브 구현 대비 약 1/6 수준(375 µs 대 2,355 µs)으로 단축함
• Neuron 백엔드를 적용한 torch.compile은 중간 값 물리 생성을 제거하여 375.5 µs의 지연 시간을 기록했으며, 이는 퓨전된 torch.rms_norm 커널의 403.3 µs와 유사한 수준임
• 네이티브 PyTorch 연산은 6개의 연산자 호출과 2,355 µs의 디바이스 시간을 발생시켰으며, 이 중 aten::mul이 자체 NEURON 시간의 49.5%를 차지함

Profiling Tools and Configuration

• torch.profiler에 전달되는 NeuronConfig는 experimental_config를 통해 데이터 스트림 캡처와 NeuronCore 선택을 제어함
• Neuron Explorer는 .neff와 .ntff 아티팩트를 활용해 시스템 트레이스, 디바이스 타임라인, 엔진 분해 정보를 시각화함
• correlate_device_time=True 설정은 key_averages() 테이블에서 CPU와 Neuron 디바이스 시간 비교를 가능하게 함

Performance Analysis

• 퓨전은 하드웨어 실행 간 공백을 제거하고 x_sq, mean_sq 같은 중간 값의 물리적 생성을 방지함
• 워크플로우는 key_averages()에서 핫스팟을 식별하고 트레이스에서 검증한 뒤 퓨전 연산 또는 컴파일을 적용함
• 시리즈 2부에서는 TensorEngine, VectorEngine, DMA 분해 등을 포함한 디바이스 레벨 프로파일링을 다룰 예정임