본문 바로가기

AI Compiler framework

AICF v2 통합 역전파 설계 (Integrated Autodiff Design)

1. 설계 철학 (Design Philosophy)

사용자가 별도의 Backward Layer 를 수도응로 쌓지 않아도, Forward 정의만으로 미분 그래프를 자동 생성하는 Tape-based Symbolic Autodiff 방식을 채택

  • 관심사의 분리 : 사용자는 모델 아키텍처에 집중, 미분 로직은 레이어 내부에 캡슐화된다.
  • 컴파일 최적화 : 생성된 전 / 역방향 그래프는 단일 CUDA Graph 로 캡처되어 Python 오버헤드 없이 실행된다.

 

2. 핵심 메커니즘 : Tape & Vid Binding

Forward Trace (Recording)

Model.add() 호출 시, 다음 정보가 _tape 에 기록된다.

  • inputs : 입력 텐서들의 Vid 리스트
  • outputs : 출력 텐서들의 Vid 리스트
  • params : 해당 레이어가 생성한 학습 가능한 파라미터 Vid 리스트

 

Backward Dispatch (Resolution)

build_backward() 실행 시 테이프를 역순으로 순회하며 레이어의 emit_backward 를 호출한다.

데이터 주입 : 레이어는 inputs, outputs, grad_y 를 전달받는다

의존성 해결

  • ReLU : outputs[0] 을 참조하여 마스킹
  • Softmax : outputs[0] 을 참조하여 자코비안 행렬 연산 수행
  • Linear : inputs 과 params 를 참조하여 그라디언트 전파

 

레이어 인터페이스 표준 (Standard Interface)

해당 표준 메서드 시그니처 준수해야 한다.

def emit_backward(self, b, ctx, inputs, outputs, grad_y, params, **kwargs):
    """
    Args:
        b: Builder (그래프 생성기)
        ctx: CudaEmitContext (KID 관리)
        inputs: 순방향 입력 Vid 리스트
        outputs: 순방향 출력 Vid 리스트
        grad_y: 상위에서 전파된 출력 미분 Vid
        params: 가중치 Vid 리스트
    Returns:
        Dict[str, int]: {"input": dx_vid, "weight": dw_vid, "bias": db_vid}
    """