AICF v2 실행 흐름 & 개선안 : 훈련 전체 CUDA Graph 캡처 + Replay 설계 문서

0. 이 시스템이 지금 “성립”하는 핵심 한 줄

AICF v2에서 학습 CUDA Graph replay가 가능한 이유는 딱 3개가 동시에 맞물리기 때문임:

1. externals(input/param/state) 전부를 내부 고정 버퍼(ext_bufs)로 재배치해서 CUDA Graph가 잡는 포인터를 고정한다.
2. ExecPlan.alias가 SSA 출력(outP 등)을 입력(P 등) 버퍼로 alias해 커널이 사실상 in-place 업데이트를 하게 만든다.
3. train 모드에서 replay 시 input만 copy하고 param/state는 copy하지 않아, 업데이트가 내부 버퍼에 누적된다.

이 3개 중 하나라도 깨지면 “replay 학습”은 깨진다.

---

1. 전체 구조(파일/모듈 레벨) 맵

1.1 핵심 모듈

• aicf_v2/model.py
  • Model: 파라미터/상태 실텐서 저장소 + tape 기반 autodiff + Builder 조립
  • Sequential: build/compile/run의 사용자 API
• aicf_v2/builder.py
  • Builder: Value/Op 저장 (vid 기반 IR)
• aicf_v2/compile/compile.py
  • compile_cuda(m): optimize_ir → make_exec_plan_cuda → CompiledProgram(plan)
• aicf_v2/compile/plan.py
  • make_exec_plan_cuda(b): alias 결정(현재는 constraints 기반 + adam_step special-case)
• aicf_v2/runtime/cuda_exec.py
  • CudaExecutor: compile_cached / capture_prebuilt / run_compiled
• aicf_v2/runtime/graph_capture.py
  • capture_cuda_graph: 고정 포인터 버퍼 기반 캡처
  • replay_cuda_graph: input copy + graph.replay
• aicf_v2/layers/*, aicf_v2/emitters/cuda/*
  • 레이어는 emit()/emit_backward()로 IR 생성
  • emitter는 emit_resolved()로 kind_id/schema/blob 포함 op를 Builder에 기록

---

2. 기본 데이터 모델: Value/Op/vid/slot

2.1 vid는 “Value 인덱스”이자 “slot 인덱스”

Builder._new_value():

• vid = len(self.values)
• self.values.append(Value(vid=vid,...))

런타임에서는:

• eager path: slots: Dict[int, Tensor]에서 vid 키로 사용
• graph path: slots: List[Tensor]에서 slots[vid]로 사용

즉 vid는 단순 ID가 아니라 메모리 슬롯 주소 지정자다.

불변조건(invariant)

• len(slots) == len(b.values)가 성립해야 한다.
• Value가 추가되면 vid는 증가하며, 기존 vid의 의미는 변하지 않는다.

실패 모드

• “optimize_ir”에서 Builder를 새로 만들면서 vid를 재번호(reindex)하면, 외부에 저장된 vid(예: 사용자가 get_vid_by_name 같은 걸로 보관한 것)가 깨질 수 있다.
• 현재 optimize_ir는 identity라 안전하지만, 추후 패스 추가 시 매우 위험해질 수 있다.

개선안

• (강력 추천) “vid 안정성”을 깨는 최적화(재번호)는 Compile 이후 외부로 vid를 노출하지 않는 방향으로 제한.
• 사용자/테스트가 슬롯 접근할 때는 vid 직접 노출 대신 name→vid 매핑 API를 공식화:
  • Model.vid(name), Program.tensor(name) 같은 안정 API 제공
• optimize_ir가 Builder를 새로 만들 필요가 생기면:
  • old_vid -> new_vid remap 테이블을 같이 제공하거나,
  • 외부에서 vid를 절대 보관하지 않도록 시스템적으로 막는 게 맞다.

---

2.2 Value role: input / param / state / tmp

Builder는 externals를 role로 분리한다:

• input: 순수 입력(x, y_true, grad_initial, bc1/bc2 같은 외부 주입)
• param: 학습 파라미터(W, b)
• state: 옵티마이저/상태(m, v, step 등)
• tmp: 내부 중간 텐서(activation, grad intermediate 등)

Builder는 또한:

• external_vids = input + param + state를 모아둔다.

불변조건

• CUDA Graph 캡처 모드에서는 external_vids 전부가 고정 포인터 버퍼를 가져야 한다.
• graph_capture는 이를 강제한다(아래 6장에서 설명).

---

3. 빌드 단계: Layer.emit으로 IR 생성 + tape 기록

3.1 Model.add(layer, *args)의 역할

Model.add()는 2가지 일을 동시에 한다.

1. IR 생성

• out_vids = layer.emit(self.b, *args, ctx=self.ctx, ...)

1. 자동미분을 위한 forward 기록(tape)

• "inputs": list(args)
• "outputs": [out_vids] 또는 list
• "params": new_params (emit 도중 새로 추가된 param vids)

그리고 추가로 “자동 파라미터 바인딩”도 한다:

• 레이어가 b.param("fc1.W") 같은 IR 등록만 하면
• Model이 실제 torch.Tensor를 만들어 self.parameters[name]=tensor로 저장한다.

불변조건

• 레이어 emit은 Builder에 Value/Op를 “append-only”로 추가해야 한다.
• b.param_vids는 파라미터 Value들을 포함하며, 동일 이름이면 재사용한다.

실패 모드

• 레이어가 실수로 같은 이름의 param을 중복 생성하면 _new_value가 자동으로 suffix를 붙인다.
  • 그러면 name 기반 디버깅에서 “fc1.W”가 아닌 “fc1.W_123”이 생겨서 찾기 어려워짐.

개선안

• param/state/input은 _name2vid에서 강하게 중복 방지하고(현재 param/state는 name 있으면 기존 vid 반환이라 OK),
• tmp value도 이름 중복을 허용하되 “디버깅용 라벨”과 “키”를 분리:
  • key는 내부적으로 unique id
  • name은 display label
• 또는 tmp는 아예 name을 필수로 받지 않고 자동 생성 규칙을 통일(“opname:out0@opidx” 같은 형태)

---

3.2 Sequential.build 흐름

Sequential.build(input_spec, input_name="x")

1. curr_vid = self.input("x", spec)
2. 레이어 순회:
   • curr_vid = self.add(layer, curr_vid)
3. self.output("output", curr_vid)

의미

• build는 “실행 가능한 최종 결과”를 b.outputs["output"]=vid로 고정해둔다.
• 이 output mapping은 replay 시 _extract_outputs에서 그대로 사용된다.

---

4. Loss & Backward: tape 역순 자동미분 생성

4.1 seed gradient: grad_initial은 input role로 추가됨

build_backward(loss_vid)는:

• loss의 spec을 가져와
• self.b.input("grad_initial", loss_spec)를 만든다.
• grad_map = {loss_vid: grad_initial_vid}

즉 grad_initial은 “외부 feed로 들어오는 값”으로 취급된다.

중요한 점

• train 모드에서 copy_roles=("input",)이므로 grad_initial은 replay마다 바뀔 수 있다(복사 대상).
• 지금 테스트에서 grad_initial을 feed로 주는 이유가 정확히 이 설계 때문.

---

4.2 tape 역순 순회 규칙

현재 코드 핵심은:

 

outs = entry["outputs"]
if not outs or outs[0] not in grad_map:
continue
grad_y = grad_map[outs[0]]

 

즉 “첫 output에 대한 grad가 존재할 때만” backward 확장함.

실패 모드

• multi-output op/layer에서 outs[0] 외 output에 대한 grad가 들어오는 케이스를 다루지 못함.
• branching 그래프에서는 동일 input에 대한 grad가 여러 경로로 모여야 하는데, 지금은 overwrite만 한다.

개선안(필수급)

• grad_map을 vid -> grad_vid로 유지하되, 이미 grad가 있으면 Add로 누적:
  • grad_map[x] = add(grad_map[x], new_grad)
• entry의 outputs 전체에 대해 grad를 처리하도록 확장:
  • 최소: “grad가 존재하는 output들”을 모아 layer.emit_backward로 넘김
• Layer.emit_backward 표준 계약을 강화:
  • 입력 grads를 dict로 반환: { "inputs": [dx0, dx1...], "params": {...} }
  • 혹은 {"input_grads": {vid: grad_vid}} 방식

---

4.3 parameter_grads 생성

layer.emit_backward()가 반환한 dict에서:

• "weight"/"bias" 키가 있으면 parameter_grads[param_vid] = grad_vid

이건 optimizer가 “파라미터에 대한 grad vid를 찾아 step op를 만든다”는 구조의 기반이다.

개선안

• 지금은 param이 2개(weight,bias)라는 전제를 깔고 있음.
• 일반화하려면 “params list와 grads dict를 name/vid 기반으로 매핑”해야 한다.
  • 예: layer_grads["params"] = {param_vid: grad_vid, ...}

---

5. Optimizer: AdamStep는 SSA out을 만든다 → alias로 접는다

5.1 AdamStep 레이어

AdamStep.emit:

• Pout/Mout/Vout를 b.value()로 새 Value로 만든다(SSA 스타일)
• emitter adam_step를 호출하면서:
  • inputs: [P, G, M, V, bc1, bc2]
  • outputs: [Pout, Mout, Vout]
  • constraints: {"inplace_ok": True}
  • hints: view_rank0_inputs=[4,5] (bc1/bc2 스칼라 ABI fix)

핵심 의미

• IR 상으로는 “새 텐서가 나오는 함수형 업데이트”처럼 모델링한다.
• 하지만 실행에서는 in-place를 허용한다는 힌트를 constraints로 준다.

---

5.2 ABI fixups: bc1/bc2 (1,) → rank0 view

runtime의 _apply_abi_fixups()가:

• hints에 view_rank0_inputs가 있으면,
• dim==1 & numel==1 텐서를 .view(())로 바꿔 op_call에 전달한다.

장점

• v2에서 “0-dim spec을 금지”해도 ABI는 만족 가능.

개선안

• 장기적으로는 TensorSpec에 “scalar” 표현을 공식화하는 게 깔끔하다.
  • spec.shape=()를 허용하든,
  • 또는 dtype/shape 외에 “rank0 semantic scalar” 플래그를 두든

---

6. compile_cuda: optimize_ir → plan(alias) → CompiledProgram

6.1 compile_cuda 코드 그대로 의미 해설

compile_cuda(m):

1. b0 = m.b
2. b1 = optimize_ir(b0) (현재 identity)
3. plan = make_exec_plan_cuda(b1)
4. return CompiledProgram(plan=plan)

현실

• 지금은 “IR 최적화도 lowering도 없다”.
• 따라서 plan은 b.ops 그대로를 실행 순서로 삼는다.
• alias만 결정한다.

개선안(로드맵)

• optimize_ir에 들어갈 후보:
  • fuse: gemm + bias_add → gemm_epilogue
  • DCE: 쓰이지 않는 tmp 제거
  • CSE: 동일 계산 재사용
  • layout canonicalization (contiguous 강제 등)
• lower 단계(중요):
  • op.kind string을 backend kind_id/schema/blob로 채우는 작업은 지금 emitter가 하고 있음.
  • 하지만 장기적으로는 high-level kind → backend variant 선택이 필요해짐.
  • 그때 registry가 compile 단계에서 의미를 갖게 된다.

---

7. plan(make_exec_plan_cuda): SSA를 in-place로 접는 핵심 단계

7.1 현재 alias 정책(코드 그대로)

• constraints["inplace_ok"] == True인 op만 고려
• 특수 케이스:
  • adam_step: outputs 3개를 각각 inputs P/M/V에 alias
  • 단, Pout == G 금지
• 일반 케이스:
  • _ALIAS_OUT0_IN0 whitelist(kind 기반)에 속하면 out0 -> in0 alias

이것이 의미하는 것

• AdamStep는 Pout/Mout/Vout를 만들지만 실제 실행 슬롯에서는:
  • slots[Pout] = slots[P]
  • slots[Mout] = slots[M]
  • slots[Vout] = slots[V]
    로 바뀌고, 커널이 out에 써도 결국 in이 업데이트된다.

장점

• 레이어 구현을 SSA로 단순화하면서도 실행은 고성능 in-place 가능.

---

7.2 지금 alias 설계의 위험한 점

1. kind 문자열 의존
   • op.kind가 변경되면 alias 정책이 무력화.
2. 검증 조건 부족
   • adam_step에서 Pout != G만 체크.
   • 실제 커널이 “out이 grad와 alias되면 안 됨”을 더 강하게 요구할 수도 있음(Mout/Vout 포함).
3. 디버깅 난이도
   • alias를 걸어버리면 “Pout이 별도 텐서”라는 SSA 의미가 사라짐.

개선안(바로 적용 가능)

• constraints에 alias policy를 명시:
  • constraints={"inplace_ok": True, "alias_policy": "out0_in0"}
  • adam_step은 "alias_policy": "adam_triplet"
• plan 단계에서 더 강한 검증:
  • out_vid가 “읽기 전용으로 필요”한 다른 op에 사용되는지 분석(간단한 use-def로도 가능)
  • grad input과 alias 금지를 outputs 전체에 대해 적용
• debug 모드:
  • inplace_ok 무시하고 alias 비활성화 옵션 제공(정확성/디버깅 우선)

---

8. CudaExecutor: compile_cached / capture_prebuilt / run_compiled

8.1 compile_cached

• key = id(m.b)
• 동일 builder면 compile 결과 재사용

개선안

• optimize_ir가 Builder를 새로 만들게 되면 compile_cache key 정책이 흔들릴 수 있음.
• plan_id / hash 기반으로 바꾸는 게 더 장기적으로 안정적임.

---

8.2 capture_prebuilt와 모드 정책

_get_roles(mode):

• train:
  • static_roles=("input","param","state")
  • copy_roles=("input",)
• inference:
  • static_roles=("input","param")
  • copy_roles=("input",)

의미

• 캡처 시 externals는 “고정 포인터 버퍼”가 필요하므로 static_roles에 포함되어야 함.
• replay 시 매 step 바뀔 수 있는 것만 copy_roles에 포함.

train에서 param/state가 copy_roles에 없다는 게 핵심:

• 내부 고정 버퍼가 업데이트되며 그대로 유지
• 외부에서 param/state feed를 줘도 replay마다 복사하지 않음 → 사실상 무시됨

개선안(현실적으로 매우 유용)

• “reset/reload” 모드 추가:
  • mode="train_reset": copy_roles=("input","param","state")
• 혹은 API 제공:
  • executor.load_externals(gprog, feed, roles=("param","state"))로 ext_bufs overwrite
• 이거 없으면 “체크포인트 로딩/재시작”에서 사용자 경험이 크게 망가진다.

---

8.3 cache_key 설계

_cache_key는 copy_roles에 해당하는 이름들의 shape/dtype/device만 signature로 사용한다.

train에서 copy_roles=("input",)이면:

• param/state의 shape/dtype이 달라도 key가 같을 수 있음.
• 하지만 어차피 replay에서 param/state를 복사하지 않으니 “외부 주입을 허용하지 않는 정책”이면 괜찮다.

개선안

• 정책을 명시적으로 문서화해야 함:
  • “train 캡처 이후 param/state는 내부에서만 유지되며, 외부에서 주입해도 반영되지 않는다.”
• 또는 strict_static_signature=True 옵션으로 static_roles도 signature에 넣을 수 있게 해두면 디버깅/안정성에 좋다.

---

9. graph_capture: 고정 포인터 버퍼 기반 캡처/리플레이

9.1 capture_cuda_graph 단계별 동작

(1) externals must be static

• external_vids 전부가 static_roles에 포함되어야 함
• 누락 시 에러

이건 CUDA Graph의 가장 중요한 제약(캡처 포인터 고정)을 코드로 강제한 것.

(2) externals에 대해 내부 버퍼 할당

• feed 텐서를 그대로 쓰지 않고
• alloc_from_spec으로 새 GPU 텐서를 할당한 뒤 copy

즉 캡처되는 포인터는 “사용자 feed”가 아니라 “내부 고정 버퍼(ext_bufs)”다.

(3) tmp 포함 전체 slots 채우기

• externals 외 Value들도 spec 기반으로 모두 alloc

(4) alias 적용

• plan.alias를 slot 레벨에서 반영
• 중요한 점: out_vid가 external이면 ext_bufs도 같이 갱신

(5) warmup

• 동일 버퍼에서 미리 실행해 안정화

(6) capture

• copy_roles(예: input)만 ext_buf로 copy
• with torch.cuda.graph(graph): _run_ops(...)

---

9.2 replay_cuda_graph

• copy_roles(보통 input)만 ext_buf로 copy
• graph.replay()
• outputs는 gprog.slots[vid]에서 반환

이것이 의미하는 것

• replay는 “입력 갱신 + 커널 실행”만 수행
• 메모리 주소/할당은 없다 (그래서 빠름)
• param/state는 내부 버퍼이므로 업데이트가 누적된다

---

10. 지금 구조가 “정확하게 작동하는지”를 검증하는 체크리스트

10.1 alias 실제 적용 여부

아래 조건이 True면 alias가 제대로 적용된 것:

• 캡처 후
  • data_ptr(gprog.slots[Pout]) == data_ptr(gprog.slots[P])
  • data_ptr(gprog.slots[Mout]) == data_ptr(gprog.slots[M])
  • data_ptr(gprog.slots[Vout]) == data_ptr(gprog.slots[V])

이걸 한 번 찍으면 “weight 업데이트가 왜 slot에서 보이는지”가 완전 확정된다.

10.2 train 모드에서 param/state가 copy되지 않는지

• replay 직전에 외부 feed의 param 텐서를 바꿔도,
• gprog.slots[param_vid]가 외부 feed와 동기화되지 않아야 정상(현재 정책상)

10.3 bc1/bc2 ABI fixup 동작 여부

• bc1/bc2가 (1,)로 들어갔을 때 _apply_abi_fixups가 view(())로 바꾸는지
• 커널이 rank0 scalar로 정상 인식하는지

---

11. 주요 실패 모드(현실에서 터지는 것들)와 개선안

아래는 “곧바로 터질 가능성 높은 것” 위주.

11.1 Autodiff의 근본적 한계 (지금은 MLP만 되니까 안 보임)

• multi-output 미지원(outs[0]만)
• gradient accumulation 없음(덮어씀)
• elementwise/add/sub의 올바른 미분 규칙 부족
• reduce_sum/softmax/ce 같은 op의 backward 일반화 부족

개선안(우선순위 1)

• grad_map 누적(Add) 도입
• entry outputs 전체 처리
• op별 backward rule 테이블 분리(현재처럼 “일부 kind만 대충”은 확장 불가)

---

11.2 alias 정책의 위험

• kind 문자열에 의존
• 검증 부족
• 디버깅 어렵

개선안(우선순위 1~2)

• constraints에 alias policy 명시 + plan에서 policy 기반 처리
• debug 모드 alias off
• alias 안전성 검사(use-def 기반 최소 검사)

---

11.3 체크포인트/재시작 문제

train에서 param/state copy 안 함 → 외부 주입이 안 먹힘

개선안(우선순위 1)

• mode="train_reset" 또는 executor.reset_externals(...)
• 최소한 “param/state를 ext_buf로 덮어쓰기” API 제공

---

11.4 dtype/device 정확성

• Model.param/state는 torch.float32로 고정 생성(현재 model.py)
• TensorSpec dtype이 f16/bf16이면 mismatch 가능

개선안(우선순위 2)

• TensorSpec dtype → torch dtype 매핑해서 그대로 생성
• fp16/bf16 mixed precision을 하려면:
  • param dtype과 accum dtype 분리(예: master fp32 + model fp16) 정책 필요

---

11.5 optimize_ir/compile_cache의 미래 충돌

optimize_ir가 Builder를 새로 만들면:

• vid가 재번호될 수 있고
• compile_cached key가 id(m.b)라 캐시가 꼬일 수 있음

개선안(우선순위 2~3)

• Builder에 graph_hash(values+ops 기반) 부여
• compiled_cache key를 graph_hash 기반으로
• 패스가 새 Builder를 만들면 remap도 함께 생성

---

11.6 CUDA Graph 캡처 안정성(멀티 스트림/비동기)

현재는:

• stream 인자 지원 안 함
• current_stream_u64로 op_call

개선안(우선순위 3)

• stream 정책 명확화:
  • capture는 특정 stream에서만, replay도 동일 stream에서만
  • 또는 per-mode per-stream graph cache 분리
• 멀티스트림/통신(op) 들어가면 반드시 동기/이벤트 정책이 필요

---

12. 권장 로드맵: “지금 코드 기반으로” 가장 싸게 큰 안정성을 얻는 순서

Phase 1 (바로 체감되는 안정화)

1. reset/reload API 추가 (param/state 덮어쓰기)
2. debug alias off 옵션
3. grad accumulation + multi-output 처리
4. dtype 생성 정확화(spec 기반 tensor 생성)

Phase 2 (컴파일러다운 확장)

1. optimize_ir에 fuse 패스(특히 gemm+bias_add → gemm_epilogue)
2. lowering/variant selection 구조 도입(여기서 registry가 의미를 갖게 됨)
3. alias 안전성 검사 강화 + policy 메타데이터화

Phase 3 (실전 운영)

1. checkpoint/restore 정식화(외부 저장 ↔ ext_bufs 로딩)
2. multi-stream / event / overlap 실행 정책
3. 프로파일/검증 모드(정확성 vs 성능) 스위치

---

13. 한 장짜리 “현재 AICF v2 학습 CUDA Graph 실행 요약” (정리)

• 빌드
  • Layer.emit이 Builder(values/ops)에 SSA IR을 append
  • Model.add가 tape 기록 + param/state 실텐서 자동 생성
• Backward
  • tape 역순으로 emit_backward 호출
  • grad_initial input을 seed로 주입
  • parameter_grads 생성
• Optimizer
  • AdamStep가 (Pout,Mout,Vout) 새 Value 생성 + inplace_ok constraints로 op 추가
• 컴파일
  • optimize_ir(현재 identity)
  • make_exec_plan_cuda가 alias 생성:
    • adam_step: Pout→P, Mout→M, Vout→V
• 캡처
  • externals(input/param/state) 전부 내부 고정 버퍼(ext_bufs)로 재할당
  • alias 적용 후 warmup + torch.cuda.graph 캡처
• 리플레이
  • input만 ext_buf로 copy
  • graph.replay
  • param/state는 내부 버퍼에서 누적 업데이트 → 학습이 진행됨