현재 CUDA Graph Capture 과정

CUDA Graph 캡쳐가 하는 일 (핵심 요약)

• 개념: 특정 스트림에서 발생하는 GPU 작업들을 “녹화”해서 cudaGraph_t로 만들고, 이를 인스턴스화(cudaGraphExec_t)한 뒤 한 번의 런치로 재실행.
• 잡히는 노드들: 커널 실행(launch), memcpy(H2D/D2H/D2D), memset, 이벤트 record/wait, (옵션) host callback 등.
• 이점: 수십~수백 개의 작은 커널/복사 호출을 CPU→GPU 런치 오버헤드 없이 한 번에 던질 수 있음 → 특히 “작은 커널 多”일수록 이득이 큼.
• 제약:
  • 캡쳐 중에는 새로운 cudaMalloc 같은 비결정적 할당이 끼면 안 됨(그래프 재현성이 깨짐). 그래서 보통 워밍업으로 메모리 풀/커널 초기화를 끝낸 뒤 캡쳐함.
  • 포인터/shape가 고정되어야 함. (바뀌면 재캡쳐 or GraphExec 업데이트 필요)
  • 캡쳐된 커널은 같은 스트림에서 동작해야 함(내부에서 default stream 쓰거나 다른 스트림으로 새 커널을 쏘면 캡쳐 실패/비호환 가능).

 

2. 캡쳐 순서 (지금 구조 해석)

IR 구성 + Arena 할당: compile()에서 GraphIR 만들고, MemoryArena로 출력/중간 텐서를 고정 주소의 장치 버퍼로 확보.

워밍업

warm_x = cp.empty(input_shape, dtype=cp.float32)
_ = self._run_impl(warm_x)       # 처음 한 번 실행해서 모든 WS/커널 로드/메모리풀 채움
cp.cuda.get_current_stream().synchronize()

 

→ 캡쳐 중 새 할당이 일어나지 않게 선행 초기화.

 

캡쳐 시작:

self._cap_stream = cp.cuda.Stream(non_blocking=True)
with self._cap_stream:
    self._cap_stream.begin_capture()
    _ = self._run_impl(None)     # 입력 복사 없이, Arena 내부 버퍼들만 사용
    graph = self._cap_stream.end_capture()

 

_run_impl(None)이 이미 할당된 Arena 버퍼만 쓰는 forward 경로를 실행 → 커널/복사 호출들이 그래프에 기록.

 

그래프 인스턴스화 & 업로드: CuPy 버전에 따라

 

• graph.upload/graph.launch (graph-object 경로)
• cp.cuda.graph.upload/launch (module-func 경로)
• 또는 런타임 포인터(graphInstantiate → graphUpload → graphLaunch)를 선택적으로 사용(네 코드에 3중 백업 경로 이미 구현).

 

실행:run(x)에서

• 캡쳐가 있으면: Arena 입력 슬롯에 x를 복사한 뒤(inp[...] = x) → 그래프 런치.
• 캡쳐가 없으면: _run_impl(x)로 즉시 실행.

 

3) 캡쳐에 들어가는/안 들어가는 것

• 들어감: _run_impl(None) 동안 실행된 Conv/GEMM 커널, 내부 D2D memcpy, 필요한 memset, 이벤트 동기화 등.
• 안 들어감:
  • run(x)에서 호스트→Arena 입력 버퍼로 복사하는 부분(현재는 파이썬 단계에서 inp[...] = x라 그래프 밖).
  • backward()/optimizer step(네가 캡쳐 구간에 포함시키지 않았으니까 당연히 밖).
  • 파이썬의 if/for 같은 호스트 제어 흐름(그래프엔 커널/복사 호출만 녹음됨).

4) “잘 캡쳐되게” 하는 필수 규칙 (실전 체크리스트)

• ✅ 워밍업 필수: 첫 호출에서 만들어지는 모든 WS/커널 핸들/메모리를 캡쳐 전에 만들어 둬라. (지금 이미 만족)
• ✅ 하나의 스트림: 캡쳐는 self._cap_stream 안에서만 이루어지게 하라. 내부 ops가 반드시 현재 스트림을 받게 하라. (지금 stream_ptr로 넘겨서 OK)
• ✅ 고정 shape/포인터: 캡쳐 후엔 입력/출력 텐서의 주소와 shape이 동일해야 한다. 바뀌면 재캡쳐. (Arena가 주소 고정이라 OK)
• ✅ 기묘한 동기화 금지: 캡쳐 중 cudaDeviceSynchronize() 같은 전역 동기화/블로킹 호출은 피하라(캡쳐 모드에 따라 에러). 이벤트 기반 동기화(같은 스트림)만 쓰는 게 안전.
• ✅ 메모리 할당 금지: 캡쳐 중 cudaMalloc류는 대부분 불가/비권장. CuPy 메모리 풀도 워밍업으로 충분히 채워두면 캡쳐 중 새 할당이 안 생김.
• ✅ 라이브러리 호환: cuBLAS/cuDNN 등 호출 시 capture-safe 버전인지(대부분 최신은 지원)와 스트림 인자 전달을 확인. (너는 커스텀 CUDA 커널이 주류라 스트림만 맞으면 OK)

5) 입력 복사도 그래프에 넣고 싶다면 (두 가지 패턴)

1. Pinned Host 스테이징(H2D memcpy 노드 캡쳐):
   • 캡쳐 전에 고정 주소의 pinned host 버퍼 h_pinned를 만들고, 그래프 안에 h_pinned → arena_input H2D memcpy를 녹인다.
   • 매 반복 때는 그 주소를 유지한 채로 h_pinned에 CPU가 데이터를 덮어쓴 뒤 그래프 런치만 한다.
   • 포인터가 변하지 않으니 캡쳐 재사용 가능.
   • (CuPy: cp.cuda.alloc_pinned(nbytes) 같은 API로 page-locked 메모리 확보 가능)
2. Device Staging(D2D memcpy 노드 캡쳐):
   • 장치에 고정 주소의 x_stage_dev(staging 버퍼)를 하나 잡고, 그래프 안에 x_stage_dev → arena_input D2D memcpy를 녹인다.
   • 매 반복 때는 그래프 밖에서 cp.copyto(x_stage_dev, x)로 데이터를 올려두고 → 그래프 런치.
   • H2D는 그래프 밖이지만, 나머지 노드는 그래프 안. 구현이 간단하고 안정적.

네 현재 구조는 “Arena 입력 슬롯에 복사 → 그래프 런치”. 위 ② 패턴으로 바꾸려면 x_stage_dev 하나만 더 두면 됨.

6) 학습 스텝까지 한 번에 캡쳐하는 법 (확장 방향)

• 트레이닝 스텝 함수를 하나로 묶기:
  forward → loss → backward → optimizer.step() → (opt) zero_grad
• 이 함수에 필요한 모든 임시/WS/그라드/옵티마 상태 버퍼를 사전할당(워밍업 때 한 번 실행).
• 그리고 그 함수를 캡쳐 구간에서 호출. (shape, 포인터, 옵티마 상태 텐서 주소 모두 고정)
• 변하는 스칼라(예: lr, momentum)는 디바이스 상수 버퍼에 넣고, 그래프 밖에서 그 버퍼 값만 바꿔주면 그래프 재사용 가능.
• 주의: loss가 리덕션을 포함하면 내부 임시 버퍼가 생길 수 있으므로 워밍업에서 동일 shape로 한 번 돌아야 안전.

7) 그래프 업데이트/재사용 팁

• GraphExec 업데이트: CUDA에는 cudaGraphExecUpdate가 있어 토폴로지 동일한 범위에서 파라미터(커널 args, grid/block, memcpy 크기 등)를 업데이트할 수 있음.
  CuPy에서 직접 노드 업데이트 API를 쓰기 어렵다면, 장치 상수 버퍼 값만 바꾸는 방식이 현실적.
• shape 바뀜: 그래프 토폴로지가 달라지면 업데이트 불가 → 재캡쳐.

8) 언제 이득이 커지나?

• 커널이 가볍고 개수가 많을수록(예: elementwise/BN/activation/작은 conv chain) → 런치 오버헤드 절감 효과 큼.
• 반대로 커널이 적고 무거우면(예: 큰 GEMM 1~2개) → 그래프 이득이 거의 없음.
• H2D가 큰 워크로드면 입력 복사 최적화(스테이징/파이프라이닝)가 체감 차이를 만듦.

9) 에러/디버그 포인트 (현장에서 자주 맞닥뜨리는 것)

• cudaErrorStreamCaptureUnsupported / IllegalState: 캡쳐 모드에서 허용 안 되는 호출(새 스트림 생성/디바이스 동기화/새 메모리 할당) 시 흔함. → 워밍업/단일 스트림/사전할당 체크.
• “포인터가 바뀜”: 메모리 풀에 의해 다른 주소가 나올 수 있음 → Arena/WS는 compile 때 고정, 외부 입력은 스테이징 버퍼(고정 주소)로 통일.
• “그래프가 너무 작음”: 로그는 멀쩡한데 속도 이득이 1.0x 근처 → 커널 수/작업량 늘리거나 train-step 캡쳐를 시도.