graph_executor_v2 — 아키텍처 개요 & 통합 가이드

0) 핵심 요약

• 상위는 그래프/오토그래드/디스패치(Registry), 하위는 백엔드(CUDA 등) 커널로 분리.
• 연산별 Op 스키마(입력/속성/검증) ↔ 디스패치(키 매칭) ↔ 백엔드 런처 ↔ 커널.
• 기존 regemm 커널/런처는 그대로 사용하고, 상위에서 래퍼로 파라미터 매핑.
• (옵션) 레거시 ge2_launch(name, bufs, stream) 경로도 그대로 유지 가능(브리지 사용 시).

 

graph_executor_v2/
├─ include/ai/                 # 공용 C++ API
│  ├─ tensor.hpp               # Tensor/TensorDesc/Device/DType/Layout
│  ├─ op_schema.hpp            # ActKind, GemmAttrs(=leaky_slope 포함)
│  ├─ dispatch.hpp             # OpKey/OpRegistry(디스패치)
│  ├─ graph.hpp, autograd.hpp  # (확장 예정)
├─ src/
│  ├─ ops/gemm.cpp             # GEMM 스키마/검증 + gemm_run()
│  ├─ dispatch/registry.cpp    # Registry 구현
│  ├─ dispatch/selector_rules.cpp
│  ├─ executor/*               # 스케줄러/메모리/오토그래드(토대)
│  └─ bindings/py_api.cpp      # 파이썬 바인딩(데모)
├─ backends/cuda/
│  ├─ register_ops.cpp         # 이 백엔드가 지원하는 조합을 레지스트리에 등록
│  └─ ops/gemm/
│     ├─ launcher.cu           # ai::Tensor → regemm::GemmBiasActParamsEx 매핑(FWD)
│     └─ backward.cu           # ai::Tensor → regemm::GemmBiasActBwdParams 매핑(BWD)
├─ kernels/regemm_epilogue/    # 기존 regemm 모듈(그대로 사용)
│  ├─ include/regemm/*.h       # api.h, activations.h, bias.h, config.h, nvtx_shim.h
│  └─ src/*.cu                 # regemm_gemm_bias_act.cu, regemm_backward.cu, launcher.cu
├─ tests/{unit,backend_cuda,python}  # 테스트 스켈레톤
└─ python/graph_executor_v2/_core.pyd  # 파이썬 확장 빌드 산출물

 

2) 실행 경로 (데이터 플로우)

Forward (현재 구현)

Python → py_api.cpp → ai::ops::gemm_run()
  → OpRegistry.find_best(OpKey{kind=GEMM, dev=CUDA, dtype=F32, layout=RowMajor, act=..., with_bias=...})
    → backends/cuda/ops/gemm/launcher.cu : GemmCudaLaunch(...)
      → regemm::gemm_bias_act_f32_ex(GemmBiasActParamsEx, cudaStream_t)
        → launch_gemm_bias_act_f32_{smoke|tiled}_ex(...) → CUDA kernels

 

3) 디스패치(Registry) 키 & 등록

• OpKey = (OpKind, Device, DType, Layout, ActKind, with_bias)
• ActKind는 None, ReLU, LeakyReLU, GELU, Sigmoid, Tanh까지 지원.
  GemmAttrs에 leaky_slope 포함(기본 0.01).

for (bool wb : {false, true}) {
  for (auto a : {ActKind::None, ActKind::ReLU, ActKind::LeakyReLU,
                 ActKind::GELU, ActKind::Sigmoid, ActKind::Tanh}) {
    R.reg({OpKind::GEMM, Device::CUDA, DType::F32, Layout::RowMajor, a, wb},
          &GemmCudaLaunch);
  }
}

 

4) CUDA 백엔드 래퍼 (핵심)

• launcher.cu : ai::Tensor/GemmAttrs → regemm::GemmBiasActParamsEx 매핑 후
  regemm::gemm_bias_act_f32_ex(...) 호출 (공식 엔트리 직접 호출).
• backward.cu : 동일 컨셉으로 regemm::gemm_bias_act_bwd_f32(...) 호출.

제약(현재 버전)

• F32, RowMajor, 비전치(trans_a=false, trans_b=false)만 연결.
• Bias는 1D 텐서 길이로 Scalar | PerM | PerN 자동 판정.

 

5) 레거시 ge2_launch 호환 (선택)

• 원한다면 ge2_launch(name, bufs, stream) 브리지를 통해 기존 이름/버퍼 레이아웃을 그대로 사용할 수 있습니다.
• 이 경우 상위 래퍼에서 bufs와 params_ex를 구성해 ge2_launch("gemm_bias_act_f32_ex", ...)로 호출 → 기존 launch_table.cpp 분기가 실행됩니다.
• 현재 기본 경로는 브리지 없이 regemm 공식 API 직호출입니다(코드 단순/안전).

 

6) 오토그래드(연결 가이드)

• Forward에서 Z stash가 필요하면 GemmAttrs에 save_preact/출력 Z를 추가하고, launcher.cu에서 p.Z, p.save_preact를 세팅.
• Backward 호출 시 입력으로 A, B, (C), gY, Z와 출력 gA, gB, (gC), (gBias)를 연결.
• 현재 regemm_backward.cu는 내부에서 gZ를 만들고, cuBLAS로 gA/gB를 계산합니다.

 

7) Python 바인딩(데모)

• src/bindings/py_api.cpp에 간단한 gemm_bias_act(a, b, bias=None, act="relu") 예제가 들어있습니다.
• 실제 제품화에서는 디바이스/스트림/메모리 소유권을 명시적으로 처리하세요(현재 코드는 데모).

 

8) 빌드 & 링크

• 루트 CMakeLists.txt에서:
  • ai_core : 디스패치/실행기/ops/파이썬 바인딩
  • ai_backend_cuda : backends/cuda/* + kernels/regemm_epilogue/src/*.cu
  • target_include_directories(ai_backend_cuda PUBLIC kernels/regemm_epilogue/include include)
  • CUDA arch, 컴파일 옵션(TF32, fast-math 등)은 GPU에 맞춰 조정.

 

9) 확장 방법

A. 새 연산 추가(예: LayerNorm)

1. include/ai/op_schema.hpp에 속성 struct 정의(에psilon 등)
2. src/ops/에 스키마/검증/shape infer 추가
3. backends/<dev>/ops/<op>/에 런처/커널 또는 기존 커널 래퍼 구현
4. backends/<dev>/register_ops.cpp에서 조합 등록
5. (옵션) 오토그래드 규칙 정의 → bwd 래퍼 추가

B. FP16/TC 경로 추가

• DType::F16 키로 레지스트리 등록 추가
• launcher.cu에서 FP16 파라미터로 매핑 후 regemm의 FP16 경로(존재 시) 호출
• 정확도 정책(TF32/FP16/bf16)은 GemmAttrs에 플래그로 확장

C. Multiple layouts / Transpose

• Layout::ColMajor/trans_a/b=true 조합에 맞춰 leading dim/shape 계산을 확장
• 레지스트리 키에 해당 조합 등록

 

10) 제약/현재 상태

• ✅ GEMM fwd/bwd(F32, RowMajor, No-Transpose) 동작
• ✅ ActKind: None/ReLU/LeakyReLU/GELU/Sigmoid/Tanh
• ✅ Bias: Scalar/PerM/PerN
• ⏳ alpha/beta 상위 노출 미정(기본 α=1, β=0), 필요 시 GemmAttrs 확장
• ⏳ Z-stash 상위 노출(Autograd 엔진과 연계) — 옵션 추가 예정
• ⏳ 다른 연산(Conv/Norm/Attention 등) — 스키마→디스패치→백엔드 패턴으로 동일 확장

 

11) 빠른 FAQ

• 왜 regemm API를 직접 호출하나요?
  레거시 bufs 파서를 우회해 표면적 API를 단순화/안정화하기 위해서입니다. 내부에서는 기존 launch_gemm_*가 그대로 호출됩니다.
• 레거시 디버깅 경로 유지 가능?
  네, 필요하면 ge2_launch 브리지를 통해 예전 이름/로그/분기 그대로 사용 가능합니다.
• 성능 튜닝은 어디서?
  대부분 regemm 커널/타일/SMEM 정책에서 이뤄집니다. 상위에서는 텐서 레이아웃/스트림/워크스페이스 캐시/메모리 플래너가 관여합니다.

 

12) 용어

• Op 스키마: 연산의 타입/shape/속성 정의와 검증 레이어
• 디스패치(Registry): (op, device, dtype, layout, act, with_bias) → 함수 포인터 선택
• 백엔드: 디바이스별 구현체(CUDA/CPU/ROCm…)
• Epilogue: GEMM 결과에 bias/activation 등 후처리 적용 단계
• Z-stash: activation 이전의 pre-activation 값을 저장해 backward에서 사용