AWS Trainium·Inferentia 전용 'Neuron Agentic'

하드웨어 최적화는 칩 설계 전문가만 할 수 있는 영역이라는 믿음이 있다. 기존에는 칩 성능을 최대치로 끌어올리려면 칩이 데이터를 처리하는 세부 경로인 커스텀 커널을 직접 설계해야 했다. 이 과정은 수개월 동안 수동으로 성능을 측정하고 수정하는 반복 작업이었으며, 칩 설계도와 씨름하며 데이터 병목 지점을 일일이 찾아내야 했기에 소수의 전문가만 수행 가능했다. AWS는 이 진입장벽을 허물기 위해 AWS Trainium과 AWS Inferentia 칩에서 작동하는 Neuron Agentic Development 기능을 공개했다. 이제 일반적인 머신러닝 엔지니어도 NKI(Neuron Kernel Interface)를 통해 하드웨어에 최적화된 커널을 직접 작성하고 관리할 수 있다.

이번 기능은 독립 프로그램이 아니라 Kiro나 Claude 같은 코딩 에이전트에 스킬 형태로 통합되어 제공된다. 개발자는 VS Code, Cursor, Kiro와 같은 IDE의 `.kiro/skills` 또는 `.claude/skills` 디렉토리에 해당 스킬을 추가하는 것만으로 기능을 즉시 사용할 수 있다. 칩 레벨의 깊은 지식 없이도 ML 엔지니어가 직접 하드웨어 최적화를 수행함으로써 추론 비용을 낮추고 모델 성능을 극대화할 수 있는 실질적인 도구가 마련된 것이다.

작성부터 분석까지, 5가지 전문 스킬의 자동화 파이프라인

개발자가 특정 연산의 최적화를 요청하면 오케스트레이터인 `neuron-nki-agent`가 전체 흐름을 제어하며 작성, 디버깅, 프로파일링, 쿼리, 문서화라는 다섯 가지 단계의 스킬을 자동으로 연결한다. 먼저 `neuron-nki-writing`이 PyTorch나 NumPy 코드, 혹은 자연어 설명을 NKI 코드로 변환한다. 이때 128 파티션 차원이나 512 또는 4096 PSUM 프리 차원 같은 하드웨어 제약 조건을 반영해 타일링 전략을 짜고, DMA 크기 설정이나 SBUF 재사용 효율성 가이드라인을 적용해 코드를 생성한다.

코드를 실행하다 에러가 나면 `neuron-nki-debugging`이 28개의 NCC 에러 코드 인덱스를 바탕으로 오류 유형을 찾아 수정안을 제시한다. `--target` 플래그를 이용한 환경 설정부터 CPU 계산 결과와 비교하는 수치 검증까지 수행하며 실행 오류를 잡는다. 이 과정에서 `neuron-nki-docs`가 `nisa.*` 또는 `nl.*` API의 정확한 시그니처와 Trainium 1, 2, 3 세대별 아키텍처 가이드를 제공해 개발자가 칩의 물리적 구조를 이해하도록 돕는다.

마지막으로 성능 최적화 단계에서는 `neuron-nki-profiling`과 `neuron-nki-profile-querying`이 협업한다. `neuron-explorer`를 통해 실제 칩의 실행 트레이스를 캡처하고 NEFF(Neuron Execution File Format)와 NTFF 파일을 생성하며, DMA Graph Engine(DGE) 알림을 포함한 세밀한 데이터까지 기록한다. 이후 이 파일들을 DuckDB나 pandas로 읽어 SQL 쿼리를 날려 성능 한계치를 계산하고 병목이 발생하는 엔진을 식별한다. 이를 NKI 소스 코드의 특정 줄과 연결해 정밀한 최적화 지점을 찾아낸다.

하드웨어 전문가의 전유물에서 ML 엔지니어의 도구로

이러한 자동화 파이프라인은 기존의 수동 최적화 방식을 완전히 바꾼다. 과거에는 숙련된 엔지니어조차 칩 아키텍처 특성을 익히기 위해 코드 한 줄을 고치고 다시 실행해 결과를 확인하는 시행착오를 수만 번 반복해야 했다. 하지만 이제는 AI 에이전트가 구현 단계마다 적절한 지침을 제공하므로, 개발자가 칩의 메모리 구조나 데이터 전송 방식 같은 복잡한 제약 사항을 일일이 외울 필요가 없다.

실제로 다른 칩 아키텍처에 익숙한 숙련자가 Trainium 환경에 적응하는 데 걸리는 시간이 수개월에서 며칠 단위로 줄어들었다. 에이전트의 가이드를 통해 기존 아키텍처에서의 경험을 Trainium 방식으로 빠르게 치환할 수 있게 되었기 때문이다. 개발자는 칩의 동작 원리를 밑바닥부터 공부하는 대신 모델의 성능을 높이는 결과물에만 집중한다.

작업 방식의 핵심은 사람이 직접 가설을 세우고 수정하며 병목 지점을 찾던 느린 루프를 에이전트 기반의 자동화된 구현으로 전환한 데 있다. 구현부터 디버깅, 분석까지 유기적으로 연결된 흐름을 통해 칩 성능을 끝까지 뽑아내기 위해 투입되던 물리적인 시간과 비용이 획기적으로 낮아졌다.

Softmax부터 SwiGLU까지, 실제 구현 성능과 인프라 요구사항

실제 Softmax 커널 구현 사례를 보면 bfloat16 정밀도를 유지하면서 하드웨어 제한치인 P_MAX=128과 F_MAX=2048 설정을 정확히 맞췄다. 행의 최댓값을 찾고 지수 합을 구한 뒤 정규화하는 과정을 하드웨어 가속 기능인 `nisa.activation(np.exp, ...)`로 구현하고, 수치적 안정성을 위해 float32 누적 방식을 적용했다. PyTorch 참조 모델과 비교해 수치적 일치 여부를 검증하는 과정에서 매뉴얼을 일일이 찾아가며 수정하던 수동 반복 작업이 사라졌다.

더 복잡한 LLM 핵심 모듈인 SwiGLU MLP 커널에서도 실질적인 성능 향상이 확인된다. 생성된 NEFF와 NTFF 파일을 분석해 칩 내부에서 데이터가 막히는 지점과 실행 시간을 잡아먹는 명령어를 식별하는 바운즈 분석을 수행했다. DuckDB와 pandas를 이용해 실행 기록을 SQL로 쿼리하며 정체 구간을 찾아냄으로써, 과거 칩 설계 전문가가 수행하던 영역을 ML 엔지니어가 직접 처리해 추론 비용을 낮췄다.

이 모든 과정은 Trainium 기반의 Amazon EC2 인스턴스 환경에서 작동하며 `trn2.3xlarge` 같은 인스턴스가 필수적으로 요구된다. 설치는 VS Code, Cursor, Kiro 같은 AI 기반 개발 도구의 `.kiro/skills` 또는 `.claude/skills` 디렉토리에 스킬 파일을 추가하는 것만으로 완료된다. 복잡한 라이브러리 의존성 설정 없이 정해진 폴더에 파일을 넣는 것만으로 칩 제어 능력을 갖추게 된다.

한국 AI 실무자에게 주는 의미: 추론 비용 절감의 현실적 대안

그동안 많은 AI 팀들은 모델의 추론 비용을 낮추기 위해 칩 내부 구조를 꿰뚫고 있는 소수의 전문가에게 모든 최적화 작업을 의존해 왔다. 커널 최적화의 진입 장벽이 너무 높아 일반 개발자가 접근하기 어려웠기 때문이다. 이제는 고가의 전문 인력을 직접 고용하지 않고도 일반 ML 엔지니어가 하드웨어 성능을 직접 최적화해 운영 비용을 절감할 수 있는 길이 열렸다.

특히 칩 레벨의 지식이 없는 엔지니어가 커널 코드를 수정할 때 겪던 설정 실수와 그로 인한 실행 오류 문제를 에이전트가 단계별 가이드로 해결해 준다. 소수 전문가의 머릿속에만 있던 최적화 노하우가 에이전트라는 도구의 형태로 팀 전체에 확산되는 구조다. 이는 곧 서비스의 응답 속도를 높이고 사용자 경험을 개선하는 결과로 이어진다.

AWS 인프라를 사용하는 한국 기업, 특히 자본력이 부족한 스타트업일수록 LLM 서비스의 성패는 토큰 하나를 생성할 때 드는 비용을 얼마나 낮추느냐에 달려 있다. 칩 설계 지식 없이도 추론 비용을 실질적으로 낮출 수 있는 도구를 갖게 됨으로써, 하드웨어 최적화라는 극소수의 전유물이 일반 개발자의 업무 범위로 내려와 인프라 운영비 절감의 현실적인 대안이 마련됐다.

칩 설계도를 외우고 수개월간 수동으로 프로파일링하며 씨름하던 커스텀 커널 작업이 이제는 에이전트와의 대화로 바뀐다. 작성부터 디버깅, 문서화까지 이어지는 neuron-nki-agent의 자동화 흐름은 칩 전문가의 전유물이었던 하드웨어 최적화 도구를 일반 ML 엔지니어의 손에 쥐여주었다.

이제는 칩 내부의 복잡한 아키텍처를 꿰뚫는 지식보다, 내가 다루는 모델의 특성을 얼마나 정확히 이해하고 최적화 지점을 찾아내느냐가 추론 비용과 성능을 결정하는 핵심 기준이 된다. 하드웨어라는 거대한 장벽이 사라진 자리에는 오직 효율적인 구현을 향한 엔지니어의 의지만이 남는다.