kernbench2/docs/adr-ko/ADR-0062-prog-tl-async-load.md

ADR-0062: Lazy tl.load — 첫 사용 시점 auto-wait를 갖는 non-blocking HBM load

Status

Accepted

ADR-0060(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context 어텐션은 KV-load-bound이며, load/compute 오버랩이 지배적 레버다. 본 ADR은 별도의 load_async op을 추가하는 대신 tl.load 자체를 lazy (non-blocking, wait를 첫 사용 시점에 자동 삽입)로 만든다. 오늘날 유일한 async primitive는 IPCQ 통신용이지 HBM load용이 아니다.

Context

오버랩이 요구하는 것

FlashAttention은 operand를 스트리밍한다: GEMM/MATH 엔진이 현재 타일을 처리하는 동안 DMA 엔진은 이미 다음 operand를 당기고 있어야 한다. 그 오버랩이 있으면 대역폭-bound 커널은 타일당 compute + dma 대신 대략 max(compute, dma)로 돈다.

ADR-0060의 하이브리드 설계에서 두 GEMM은 tl.composite(op="gemm")로 발행되고, 그 K/V operand는 tl.ref(HBM 상주, PE_SCHEDULER가 타일당 스트리밍)이므로 타일당 K/V prefetch는 composite scheduler가 처리한다. lazy tl.load는 나머지 명시적 load(Q group, 그리고 비-composite 커널)를 담당하여 그것들도 greenlet을 멈추는 대신 뒤따르는 compute와 오버랩되게 한다.

현재 존재하는 것

• tl.load(ptr, shape, dtype)은 blocking이다: DmaReadCmd를 emit하고 greenlet 커널은 PE_DMA가 완료를 알릴 때까지 suspend된다 (tl_context.py:177-203; greenlet 구동 kernel_runner.py:146-153). 한 커널에서 두 tl.load가 in-flight일 수 없다.
• async 패턴이 존재하긴 하나 IPCQ 전용이다: tl.recv_async(dir, ...) -> RecvFuture + 지연 wait (kernel_runner.py:248-285). 이는 메커니즘을 입증한다 — future를 반환하고 나중에 wait 체크(if not future.event.triggered: yield future.event)로 resolve되는 non-blocking 커맨드가 greenlet 모델에서 동작한다.
• DMA 엔진은 read 채널을 write 채널과 분리된 SimPy resource(capacity 1, pe_dma.py:45)로 모델한다. 따라서 in-flight read는 표현 가능하며, 단일 read 채널에서 직렬화되면서 PE_GEMM/PE_MATH의 compute와 오버랩된다. 오직 커널-대면 API만이 현재 load-vs-compute를 직렬화한다.

tl.load는 오늘날 뒤따르는 compute와 오버랩될 수 없다.

Decision

tl.load를 lazy로 만든다: DmaReadCmd를 발행하고 핸들을 즉시 반환 (non-blocking); 런타임은 load된 데이터가 실제로 처음 소비되는 지점에 wait를 자동 삽입한다. 커널-대면 API는 불변이다 — 작성자는 계속 tl.load를 쓴다 — 따라서 이는 새 op이 아니라 시맨틱 변경이다. 기존 recv_async/wait 기계장치를 (1) HBM-load 경로로, (2) 명시적 tl.wait 호출에서 암묵적·의존성 기반 first-use wait로 일반화한다.

D1. tl 표면 — 불변

tl.load(ptr, shape, dtype)은 시그니처와 TensorHandle 반환을 유지한다. 바뀌는 것은 언제 blocking하느냐다: 발행 시점에는 결코 아니고, 결과가 소비 op에 의해 처음 읽힐 때만 암묵적으로.

D2. 메커니즘 — non-blocking 발행 + 사용 시점 auto-wait

• tl.load는 완료 event를 yield하지 않고 DmaReadCmd를 PE_DMA에 post하며, 반환 핸들에 pending event를 기록한다(recv_async 패턴을 load에 적용).
• 소비 op(tl.dot, MATH op, tl.store, tl.composite operand, …)이 dispatch될 때, 런타임은 각 입력 핸들의 pending load event를 확인하고 아직 triggered가 아니면 먼저 yield한다 — 즉 wait가 가장 늦은 올바른 지점(첫 사용)에 자동 삽입된다.
• op_log 엔트리는 불변(memory/dma_read): asynchrony는 스케줄링 속성이지 새 op 종류가 아니므로 dma_read_count 및 기존 op_log 소비자가 계속 동작한다.

D3. 범위 — 전역

tl.load는 opt-in 플래그 뒤가 아니라 모든 곳에서 lazy다. 이것이 충실한 모델이다: 매 load마다 blocking하는 것은 naive 커널의 속성이고, 효율 커널(그리고 실제 컴파일러)은 load를 hoist하고 사용 시점에만 wait한다. 결과: tl.load와 첫 사용 사이에 독립 작업이 있는 기존 커널은 더 낮은(빠른) latency를 본다 — 동작의 회귀가 아니라 모델의 정확도 개선. 바뀌는 골든 latency는 재생성해야 하며; load 직후 곧바로 사용하는 커널은 변화가 없다(auto-wait이 즉시 발동하여 blocking과 동일).

D4. Latency / 오버랩 시맨틱

• tl.load는 발행(descriptor push)만 청구하고 커널은 진행한다. (per-op 발행 비용은 dispatch_cycles, 현재 0; op 종류별 차등 발행 비용은 별도로 추적 — ADR-0060 §1, §9, ADR-0064.)
• DMA 전송은 read 채널(capacity 1)을 모델 지속시간만큼 점유하며 커널이 다음에 내는 compute와 병렬; 다중 in-flight load는 채널에서 직렬화되나 compute와 오버랩.
• 자동 삽입된 wait는 전송이 끝나지 않았을 때만 blocking.
• 결정성 보존: wait 지점은 프로그램 순서(첫 사용)로 고정되고, 완료는 modeled DMA 채널의 scheduled event다(SPEC §0.1, R8). latency 빼기 없음 — 오버랩은 실제 모델 동시성.

모델링 가정. auto-wait-at-first-use는 잘 스케줄된 커널(컴파일러가 wait를 가장 늦은 올바른 지점에 둠)을 모델한다. 실제 컴파일러는 이를 근사하며, 일부 load는 hoist 불가(register pressure, aliasing). 성능 시뮬레이터(SPEC §0)에서 잘 스케줄된 경우를 모델하는 것이 의도된 동작이다.

Alternatives

A1. 별도 tl.load_async op (이전 제안)

커널이 손으로 호출하는 명시적 load_async/wait 쌍 추가(더블버퍼 댄스). 기각: 커널-대면 API를 키우고 버퍼 수명 부기를 모든 커널 작성자에게 떠넘긴다. lazy tl.load는 API 표면 변경 없이 동일 오버랩과 컴파일러식 auto-wait를 준다.

A2. 커널별 opt-in lazy load

tl.load를 기본 blocking으로 두고 새 커널만 opt-in. 기각: tl.load 시맨틱을 두 변종으로 쪼개고 기존 bench로부터 모델 개선을 숨긴다; 전역 lazy가 더 깔끔하고 골든 재생성 비용은 1회성(D3).

A3. composite 스트리밍에만 의존 (lazy load 없이)

ADR-0060의 composite는 tl.ref K/V operand를 스트리밍하므로 GEMM operand DMA는 이미 오버랩된다. 그러나 명시적 load(Q group, 비-composite 커널)는 lazy tl.load 없이 여전히 stall. 단독으로는 불충분.

Consequences

Positive

• 커널-대면 API 변경 제로로 load/compute 오버랩; 작성자는 계속 tl.load.
• recv_async/wait와 대칭 — 개념적 표면적 작음.
• 일반적: 어떤 대역폭-bound 커널도 자동 prefetch.

Negative

• load와 사용 사이에 독립 작업이 있는 커널의 기존 골든 latency가 이동(빨라짐) → 1회성 재생성(D3).
• auto-wait는 런타임이 핸들별 pending event를 추적하고 소비-op dispatch에서 확인하도록 요구(데이터 의존성 추적).
• scratch 수명(ADR-0063)과 상호작용: in-flight load의 대상 버퍼는 auto-wait이 발동하기 전에 recycle되면 안 된다. recycling scope는 살아있는(미-wait된) load 버퍼를 제외해야 한다.

Test Requirements

1. 오버랩이 실재: tl.load 후 비슷한 지속시간의 독립 GEMM을 내는 커널이 load+gemm이 아니라 ≈max(load, gemm)으로 완료(end-to-end latency가 직렬 합보다 엄격히 작음).
2. auto-wait 정확성: load된 TensorHandle이 처음 소비될 때, 같은 주소에 대해 오늘날 blocking tl.load와 동일한 바이트를 보유(Phase 2).
3. 둘 in-flight: 서로 다른 주소로의 두 tl.load가 나중에 소비되어 둘 다 올바른 독립 텐서로 resolve; 그 DMA는 read 채널에서 직렬화.
4. op_log 호환성: 각 tl.load은 여전히 정확히 하나의 memory/dma_read를 로그; dma_read_count는 blocking 버전 대비 불변.
5. 무-오버랩 무-변화: load 직후 곧바로 사용하는 커널은 blocking 모델과 동일한 latency(auto-wait 즉시 발동).