dev_AI_framework (326) 썸네일형 리스트형 From FlashAttention to AI Framework Ops 문서 GPU Kernel Design & Analysis Knowledge Transfer단순히 FlashAttention 구현 기록이 아닌, 하나의 고난도 op 를 구현하면서 얻은 GPU 커널 설계 원칙을 모든 AI ops 구현에 재사용하기 위한 기준 문서 1. FlashAttention 이 특별한 이유FlashAttention 은 단순한 attention op 가 아니다GEMM + reduction + exp + normalizetile streamingonline statisticsstrict numerical orderextreme memory pressure즉, 대부분의 딥러닝 ops 가 겪는 문제를 한 커널에 압축한 사례이를 통해 얻은 지식은 AI ops 전반에 일반화 가능 2. ops 커널 구현의.. AI Compiler 의 역할 재정의 - 기존 역할과 학습 단계에서의 새로운 역할 구분 1. 기존 AI Compiler 의 주된 역할기존 AI Compiler 는 모델의 의미 semantics 를 보존한 채 실행 형태를 변환하는 시스템입력은 모델의 구조, 출력은 동치 (equivalent) iR 또는 실행 코드, 다음의 목표실행 효율 극대화메모리 트래픽 감소런타임 오버헤드 제거하드웨어 제약 하에서의 최적 실행AI Compiler 는 모델이 무엇을 계산하는지는 바꾸지 않는다. 추론은 AI Compiler 가 가장 잘 작동하는 환경그래프 구조 고정파라미터 불변side-effect 없음결정적 실행 가능이로 인해 aggressive optimization, CUDA Graph cpature, kernel fusion 등이 가능 학습 단계에서의 AI Compiler 의 한계파라미터 및 optimi.. 🧭 GE2 Runtime — 설계 회고 (Design Retrospective) “커널 오버헤드를 줄이려다, 새로운 Runtime 구조가 탄생하기까지”1. 시작점: 단순한 목표내가 처음에 목표로 한 건 그리 거창한 것이 아니었다.✔ 목표CUDA Graph capture로 커널 호출 오버헤드를 줄이자.(PyTorch의 cuda.graph처럼 forward 내 모든 kernel을한 번에 capture-replay 할 수 있다면 훨씬 빠르게 실행될 것이다.)그 당시 내 가정은 매우 단순했다:학습의 loss 계산optimizer.step()backward()이런 것들도 capture 안에서 돌 수 있을 거라고 생각했다.그래서 “학습 전체를 하나의 큰 graph처럼 묶어버릴 수 있지 않을까?”라는 시점에서 출발했다.하지만 이 지점에서 근본적인 오해가 있었다.2. 위기: “생각보다 captur.. 어느 환경에서 학습 시 매 번 새로운 graph 가 생성될까 대부분의 일반 모델 학습에서는실제로 그래프 구조가 iteration 마다 크게 달라지지 않음,대부분 iteration 에서 동일한 형태의 dynamic graph 가 반복 생성되는 것단순히 새로 생성되는 것 뿐이지, 구조적으로는 거의 동일한 형태그래프 caching 에 관심, autograd 특성 때문에 정식 지원 확장 불가능 하지만 특정 모델 / 훈련 패턴에서는구조가 iteration 마다 바뀌는 경우도 존재NLP generationMixture-of-ExpertsDynamic routing networkEarly-exit network조건 기반 loss switchingRL / 시뮬레이션 기반 모델variable - length inputvariable - length decodingdropout /.. 상용 딥러닝 프레임 워크는 Inference Runtime 기반 구조이다...! 딥러닝 프레임워크 (Pytorch, TensorFlow, JAX 등 ) 은 본질적으로 inference 성능을 기준으로 설계된 runtime 구조를 가진다.그 이유는 training 과 inference 이 시스템 관점에서 요구하는 특성이 다르기 때ㅜㅁㄴ 1.Training vs Inference 요구 조건의 차이Trainingdynamic control-flow 필요autograd graph 매 iteration 새로 생성메모리, CPU, GPU 모두 매우 복잡하게 움직임python-level branching, 반복, shape 변화가 자연스럽게 발생Training = dynamic execution + autograd engine 중심 구조Inference동일한 연산 패턴 반복control-flo.. gemm_bias_act_f32_tilled_kernel 만 자세히 봐보자 1. 해당 커널이 구현한 기본 패턴크게 보면, C = alpha A B + beta C + bias -> activation 을 M x N 타일로 쪼개고, K 축은 BK 씩 잘라가면서 A/B 의 K-슬라이스를 shared 에 올려서 여러 번 재사용하는 구조 2. 공간 쪼개기 : block / thread / thread-tile2-1. block -> C 타일const int m0 = blockIdx.y * BM_;const int n0 = blockIdx.x * BN_;blockIdx.y, blockIdx.x 가 각각 M 축, N 축 타일 인덱스(m0 ,n0) = 이 블록이 담당하는 C 타일의 좌상단 global index즉, 이 블록은A 의 [m0 ... m0+BM_-1, :]B 의 [:, n0 .... 현재 구현된 gemm 의 fwd 부분 커널 코드 확인 1. 타일 / 스레드 구조 정의constexpr int BM = REGEMM_TILE_M;constexpr int BN = REGEMM_TILE_N;constexpr int BK = REGEMM_TILE_K;constexpr int TDX = REGEMM_BLOCK_TDX;constexpr int TDY = REGEMM_BLOCK_TDY;constexpr int THR_M = REGEMM_THREAD_TILE_M;constexpr int THR_N = REGEMM_THREAD_TILE_N;BM, BN, BK : block 하나가 맡는 타일 크기TDX, TDY : 블록 내 스레드 배치 ( threadidx.x / y 최대값 )THR_M, THR_N : 스레드 하나가 담당는 C 타일 내 미니 타일 크기.. Shared Memory - Bank Conflict Shaered Memory 의 구조Bank 개수 : 32 개, 각 warp 의 thread 수와 동일Bank 너비 : 4bytes(float 기준), 주소 % 32 - bank 결정Access 단위 : Warp 단위 ( 32 threads), conflict / broadcast 판단은 warp 내부에서만 발생 Bank 매핑 규칙bank = (address_in_bytes / 4) % 32 // float 기준thread 는 thread 단위로 동작하는 것처럼 보이지만, shared memory 접근은 warp 단위로 동시에 발생한다.따라서 bank conflict 여부는 thread 가 아니라 warp 내부 패턴으로 결정된다. Access 패턴 시나리오Broadcast (Conflict 없음)__sha.. 이전 1 2 3 4 ··· 41 다음