기술 요약: Xe-Forge

문제 정의

새로운 하드웨어 가속기, 특히 인텔 GPU 로 딥러닝 알고리즘을 포팅할 때, 개발자들은 모든 Triton 커널에 대해 반복적인 저수준 최적화 (정량화, 메모리 병합, 타일 튜닝, 아키텍처별 우회 작업) 를 수동으로 적용해야 합니다. 이러한 수동 작업은 상당한 병목 현상을 초래합니다. 최적화 패턴은 일관되지만, 하드웨어 제약 조건 (와프 수, GRF 모드, SLM 크기) 은 장치마다 다르며 종종 대규모 언어 모델 (LLM) 학습 데이터에 포함되어 있지 않기 때문입니다. 기존 LLM 기반 시스템은 주로 NVIDIA 나 AMD 하드웨어를 위해 커널을 처음부터 생성하는 데 중점을 두고 있어, 인텔의 특정 아키텍처 (Arc/Arc Pro 시리즈) 에 대해 기존에 기능적으로 올바른 Triton 커널을 최적화하는 데 있어 중요한 공백이 존재합니다.

방법론

저자들은 기존 Triton 커널의 최적화를 자동화하도록 설계된 다단계 LLM 기반 파이프라인인 Xe-Forge를 제시합니다. 이 시스템은 다음과 같은 핵심 구성 요소를 통해 작동합니다.

1. 다단계 파이프라인 아키텍처

Xe-Forge 는 기능적으로 올바른 Triton 커널을 최대 9 단계의 순차적 최적화 단계를 거쳐 처리합니다. LLM 기반 플래너는 감지된 문제를 기반으로 하드 의존성 제약 조건 (예: 알고리즘 재구성은 데이터 타입 변환에 선행해야 함) 을 준수하면서 실행 순서를 결정합니다. 단계는 다음과 같습니다:

1. 분석: 최적화 기회를 감지합니다.
2. 알고리즘 최적화: FLOP 및 메모리 액세스를 줄이기 위해 계산을 재구성합니다 (예: 중복된 GEMM 제거).
3. 발견: 명시된 단계에서 다루지 않는 개방형 새로운 최적화를 적용합니다.
4. 데이터 타입 수정: 불필요한 정밀도를 변환합니다 (예: float64 를 float32/16 으로).
5. 퓨전: 별도의 커널 또는 요소별 체인을 병합합니다.
6. 메모리 액세스: 병합, 프리패칭, 데이터 레이아웃을 최적화합니다.
7. 블록 포인터: 수동 포인터 산술을 tl.make_block_ptr 로 대체합니다.
8. 지속적 커널: 반복 간 데이터 재사용을 위해 루프를 재구성합니다.
9. GPU 특화 튜닝: 인텔 특화 제약 조건 (와프 수, GRF 모드, 타일 크기, GROUP_SIZE_M 스와즐링) 을 적용합니다.
10. 자동 튜닝: @triton.autotune 구성 그리드를 생성합니다.

관련 문제가 없는 단계는 비용과 코드 변경을 줄이기 위해 건너뜁니다.

2. 검증 및 정제 체인 (CoVeR)

각 단계는 DSPy 로 구현된 CoVeR 에이전트에 의해 실행되며, 이는 후보를 반복적으로 생성하고 검증한 후 피드백을 기반으로 정제합니다. 에이전트는 다음 루프에서 작동합니다:

• 생성: 코드 변경을 제안합니다.
• 검증: 4 단계 캐스케이드를 실행합니다:
  1. 구문: Python AST 파싱.
  2. 구조: 하드웨어 제약 조건 확인 (예: 2 의 거듭제곱 와프, 유효한 GRF 모드).
  3. 정확성: torch.allclose 를 사용하여 PyTorch 참조와 출력을 비교합니다.
  4. 성능: 실제 하드웨어에서 벤치마크를 수행하며, 기준선보다 느리면 거부합니다.
• 정제: 검증이 실패하면 오류 진단이 다음 반복을 위한 컨텍스트에 피드백됩니다.

3. 큐레이션 지식 베이스

LLM 학습 코퍼스에 인텔 특화 데이터가 부족하다는 문제를 해결하기 위해 Xe-Forge 는 기계 가독성 YAML 지식 베이스를 활용합니다. 이는 인텔 Xe 아키텍처에 특화된 하드 제약 조건 (예: 유효한 와프 수, SLM 용량 제한) 과 최적화 패턴 (전후 코드 변환) 을 인코딩합니다. 이를 통해 모델이 아키텍처적으로 유효한 범위 내에서 작동하도록 보장합니다.

4. 검증 인프라 (AI Bench)

이 시스템은 최적화 로직과 측정을 분리하는 독립형 프레임워크인 AI Bench에 의존합니다. 이는 L2 캐시 플러시를 사용하는 하드웨어 이벤트 타이밍을 사용하여 재현 가능한 타이밍과 정확성 검사를 보장합니다. 중요하게도, 파이프라인은 커널 하네스 분리를 강제합니다. LLM 은 커널 본체만 수정할 수 있으며, Model 래퍼, 입력 생성, 실행 경로는 고정된 상태로 유지됩니다. 이는 모델이 벤치마크 논리를 우회하는 적대적 최적화를 방지합니다.

주요 기여

1. Triton-to-Triton 최적화 공식화: 기존 PyTorch-to-Triton 생성 접근법과 달리, Xe-Forge 는 시맨틱을 보존하면서 기존 커널을 최적화합니다.
2. 최초 인텔 GPU 타겟팅 시스템: 학습 데이터에 없는 제약 조건을 위한 하드웨어 인식 지식 베이스를 갖춘, 인텔 GPU 를 대상으로 하는 최초의 LLM 기반 시스템입니다.
3. 다단계 CoVeR 파이프라인: 각 단계가 전용 패턴과 검증 기준을 갖는 분해된 접근 방식으로, 유연성과 재현성 사이의 균형을 맞춥니다.
4. AI Bench 프레임워크: 실행/타이밍을 최적화 로직에서 분리하여 백엔드 간 공정한 비교를 보장하는 벤치마킹 도구입니다.
5. 적대적 완화: 모델이 벤치마크를 우회하여 정확성 검사를 통과하는 것을 방지하기 위해 실행 하네스를 고정하는 설계입니다.

실험 결과

이 시스템은 Intel Arc Pro B70 GPU 에서 97 개의 Level-2 KernelBench 커널과 Flash Attention 구성에 대해 평가되었습니다.

• 전체 성능: PyTorch eager 대비 기하평균 1.17 배 속도 향상을 달성했으며, **67%**의 커널에서 개선이 나타났습니다.
• 극단적인 속도 향상: 9 개의 커널이 5 배 이상의 속도 향상을 보였으며, 가장 높은 경우 82 배 (Matmul_Min_Subtract) 에 달했습니다. 이러한 이득은 중복 계산을 제거하는 알고리즘 재구성에 기인하여, 원래 계정보다 FLOP 수를 효과적으로 줄였습니다.
• Flash Attention: 모든 16 개 테스트 구성 (다양한 모델 크기, 시퀀스 길이, 배치 크기) 에서 2 배에서 13.3 배의 속도 향상을 달성했으며 회귀는 발생하지 않았습니다. 특히, 시스템은 지능적 커널로는 불가능했던 지속적 커널과 SLM 타일링을 활용하여 32GB GPU 에서 32k 시퀀스 길이 추론을 가능하게 했습니다.
• 정확성: 최적화된 커널의 100% 가 PyTorch 참조에 대한 정확성 검사를 통과했습니다.
• 루프라인 분석: 계산 병목 커널 (GEMM/MatMul) 의 경우 파이프라인이 성능을 최대 효율에 근접하게 끌어올렸습니다. 대역폭 병목 커널의 경우 성능은 일반적으로 경쟁력이 있었으나, 메모리 처리량 제한으로 인해 일부는 경미한 회귀 (0.5~0.8 배) 를 보였습니다.

중요성과 주장

이 논문은 하드웨어 특화 커널 최적화에 있어 모델 규모 단독보다 구조화된 도메인 지식과 하드웨어-인-루프 검증의 결합이 더 효과적이라고 주장합니다.

• 지식 격차 해소: 큐레이션 지식 베이스에 인텔 특화 제약 조건 (와프 수, GRF 모드) 을 인코딩함으로써, LLM 이 이러한 특정 아키텍처에 대한 학습 데이터가 부족하다는 점을 보완합니다.
• 포팅 노력의 체계적 제거: Xe-Forge 는 반복적인 포팅 작업을 자동화하여 새로운 가속기에서 알고리즘 배포를 가로막는 현재 병목 현상을 체계적으로 제거할 수 있음을 보여줍니다.
• 견고성: 변경 가능한 커널 코드와 불변의 하네스 코드를 분리함으로써, 다른 생성 시스템에서 볼 수 있는 모델이 제약 조건을 우회하여 정확성 검사를 통과하는 '회피' 실패 유형을 효과적으로 방지합니다.

저자들은 파이프라인이 완벽하지는 않다는 점 (소수의 커널에서 경미한 회귀가 발생하며 새로운 알고리즘 패턴은 여전히 어렵다는 점) 을 인정하면서도, 생산 워크로드의 대부분을 구성하는 표준 연산자 및 퓨전 구성의 최적화를 성공적으로 자동화했다고 결론지었습니다.