kernbench2/docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md

ADR-0063: tl.scratch_scope — long-context를 위한 tile 단위 scratch 재활용

Status

Accepted

ADR-0060(AHBM GQA Fused Attention)의 보조 ADR. Long-context attention 은 수많은 K/V tile을 sweep하며, 매 tile의 중간 결과 (scores, P, exp, partial O, …) 는 현재 kernel 호출 동안 회수되지 않는 fresh scratch를 새로 할당한다. PE당 1 MiB scratch 예산은 실제로 의미 있는 context length에 도달하기 훨씬 전에 소진된다.

Context

Bump allocator

TLContext 는 모든 math/compute output handle을 PE-local scratch pool 에서 linear bump cursor 로 할당한다 (src/kernbench/triton_emu/tl_context.py; _scratch_alloc):

def _scratch_alloc(self, nbytes):
    aligned = (nbytes + 15) & ~15
    addr = self._scratch_base + self._scratch_cursor
    self._scratch_cursor += aligned
    if self._scratch_cursor > self._scratch_size:      # default 1 << 20 = 1 MiB
        raise RuntimeError("TLContext scratch overflow: ...")
    return addr

Docstring은 cursor가 "매 kernel 호출마다 reset된다" 고 명시한다 — 즉 kernel entry에서만 reset, kernel body 내에서는 절대 reset되지 않는다. 모든 tl.dot, tl.softmax, tl.exp, tl.sum, a - b, a * b, … 는 fresh slice를 가져가고 kernel이 return될 때까지 아무것도 회수되지 않는다.

왜 GQA kernel에서 문제가 되는가

현재 milestone bench는 이 한계 때문에 S_q = S_kv_per_rank = 16 으로 명시적으로 고정해 두었다 (tests/attention/test_milestone_gqa_llama70b.py:123-148):

"S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the simulator's 1 MB per-PE TCM kernel scratch is not exhausted by the bump-allocated handle outputs of softmax/exp/dot/sum chains over n_ranks ring steps."

n_tiles개의 tile sweep에 대해 FlashAttention은 O(n_tiles × tile당 중간값) 만큼 할당한다. 의미 있는 context에서는 overflow된다. Flash attention의 수학 자체는 O(1) 의 live scratch만 필요로 한다 — running (m, l, O) + 현재 tile의 working set — 매 tile의 임시값은 running state에 fold된 직후 dead가 되기 때문이다. 다만 allocator가 그 사실을 모를 뿐이다.

Decision

Kernel이 할당 영역을 재활용 가능(reclaimable) 으로 표시할 수 있는 scratch scope 을 추가한다. 이를 통해 tile당 임시값은 재활용되고, live running state(및 in-flight prefetch buffer)는 보존된다.

D1. tl 표면 — context manager

with tl.scratch_scope():
    s   = tl.dot(q, k_t)        # `with` 내부의 모든 handle은
    p   = tl.softmax(s)         # 종료 시 rewind되는 영역을 공유
    o_j = tl.dot(p, v)
    # ... o_j를 외부 영역에 사는 running (m,l,O) 에 fold ...
# __exit__: cursor가 __enter__ 시점 값으로 rewind

Semantics:

• __enter__ 는 현재 _scratch_cursor 를 save-point로 기록.
• __exit__ 는 cursor를 save-point로 복원, 내부에서 할당된 모든 것을 free.
• Scope 외부에서 할당된 handle (running m,l,O, ADR-0062의 LoadFuture 가 보유한 prefetch buffer) 은 주소를 유지 — save-point 이전에 또는 enclosing scope에서 할당되었기 때문.

D2. 안전 계약

Scope 내부에서 할당된 handle은 scope exit 이후 읽으면 안 된다 — 그 바이트는 다음 scope의 할당으로 덮어쓰여질 수 있다. Flash loop은 이를 자연스럽게 지킨다: tile iteration 사이에 살아남는 값은 running accumulator 뿐이고, 이들은 per-tile scope 외부에서 할당되며 내부의 op들이 외부 주소로 쓰면서 갱신된다(merge가 새 running state 를 쓰는 방식 — D3 참조).

D3. Running accumulator와의 상호작용

Online-softmax merge는 이전 running (m, l, O) 와 현재 tile의 (m_j, l_j, O_j) 를 읽어 새 running 값을 만든다. 새 running 값을 재활용 영역 바깥에 두기 위해, merge는 결과를 loop 시작 전 1회 할당된 stable scratch (per-tile scope과는 별개의 작은 고정 "running-state" arena) 로 쓴다. 구체적으로 kernel은 두 개의 arena를 둔다:

• persistent arena (1회 할당): m, l, O (merge용 double-buffer가 필요한 경우 그것까지).
• scoped arena (매 tile rewind): scores, P, exp, O_j, scale_*.

이는 실제 flash-attention SRAM budgeting의 모습과 같다: 작은 persistent accumulator 영역 + 재활용되는 tile working set.

D3.1 Persistent-arena 쓰기 mechanism: tl.copy_to(dst, src)

Merge op (tl.maximum, tl.exp, binary * / +) 모두 D1의 bump cursor로부터 할당하는 _make_compute_out(...) 을 호출한다. 따라서 scratch_scope 내부에서 그 결과 handle들은 scope 내부에 살고 __exit__ 시점에 사라진다. D3의 two-arena 분리를 실현하려면 kernel은 scoped 결과의 바이트를 persistent 주소로 쓰는 방법이 필요하다.

이 gap을 메우는 primitive:

def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
    """``src`` 의 바이트를 ``dst`` 의 주소로 복사한다 (양쪽 모두 TCM).

    Shape, dtype 일치 필요. ``dst`` 는 통상 어떤 활성 ``scratch_scope`` 도
    바깥에서 할당된 handle (persistent arena); ``src`` 는 scope ``__exit__``
    이후에도 바이트가 살아 있어야 하는 scoped handle.
    """

tl.store (HBM 측 바이트 복사) 와 대칭이지만, running-state writeback이 op_log를 불필요한 DMA로 오염시키지 않도록 TCM 전용으로 유지한다.

Mechanics:

• 신규 CopyCmd(src, dst, nbytes, data_op=True) command.
• op_log: op_kind="math", op_name="copy" — vector engine에서 실행.
• Latency: pe_math._compute_ns(prod(shape)) — HBM 전송이 아니라 on-chip register writeback을 모델링.
• Emit-time 검증: dst.shape == src.shape, dst.dtype == src.dtype, dst.space == "tcm", src.space == "tcm". 작성자 오류는 Phase 2 data execution 깊숙한 곳이 아니라 Phase 1에서 노출.

호출 패턴:

m, l, O = init_running(...)                  # persistent (scope 외부)
for j in range(n_tiles):
    with tl.scratch_scope():
        ...                                  # tile 작업 (재활용)
        m_new = tl.maximum(m, mj)            # scoped scratch
        l_new = l * scale_old + l_step * scale_step
        O_new = O * scale_old + O_step * scale_step

        tl.copy_to(m, m_new)                 # ← 새 running state를 persist
        tl.copy_to(l, l_new)
        tl.copy_to(O, O_new)
    # exit: scoped m_new/l_new/O_new 는 사라지고, 그 바이트는
    #       persistent m/l/O 에 살아있다

copy_to 는 __exit__ 이전 에 실행되므로 src (scoped) 의 read는 유효하며, exit 이후에는 dst (persistent) 만 read되어 D2의 "exit 이후에는 읽지 말 것" 안전 계약을 만족한다.

왜 모든 math op에 dst= kwarg를 추가하지 않고 전용 primitive 인가 (검토 후 기각): tl.maximum, tl.exp, _binary_math, _unary_math, _reduction 에 dst= 를 추가하는 것은 5개 op-family 에 걸쳐 ~25 LOC 이며, "호출은 fresh handle을 return한다" 라는 일관된 패턴을 깬다. tl.copy_to 는 단일 primitive, 단일 command, 단일 executor handler — 같은 효과에 비해 최소한의 surface area.

D4. Nesting

Scope는 nest 가능 (save-point의 stack). Inner scope exit은 inner save-point로 rewind, outer exit은 더 멀리 rewind한다. Prefill의 "per-query-block" 외부 scope을 둘러싼 "per-KV-tile" 내부 scope을 지원한다.

Alternatives

A1. Temporary를 HBM에 round-trip

중간값을 HBM에 store하고 reload하여 TCM scratch를 "free" 한다. 기각: TCM-resident streaming kernel을 HBM-bandwidth-bound로 만들어 목적에 정반대이며, op_log를 불필요한 DMA로 오염시킨다.

A2. Tiled tl.composite (scheduler-관리 scratch)

tl.composite 는 PE_SCHEDULER 내부에서 tile당 scratch를 자동 재활용 한다. ADR-0062 A1과 마찬가지로 매력적이지만 flash-capable composite kind (두 GEMM + carry되는 (m,l,O)) 가 필요하므로 훨씬 큰 변경에 종속된다. scratch_scope 는 작고 일반적인 primitive로 greenlet kernel 에 동일한 memory 거동을 제공한다. 두 방식은 공존 가능.

A3. Scratch 예산 증가

pe_tcm.kernel_scratch_mb 를 증가. 해결책으로는 기각: O(n_tiles) scratch leak이 여전히 존재하면서 context-length ceiling만 선형적으로 밀어내며, hardware 실제 모습 (실 TCM은 작고, flash attention의 핵심은 O(1) working set) 을 잘못 표현한다. 거친 조절 knob으로만 유용, 재활용의 대체재는 아니다.

Consequences

Positive

• 인위적인 S = 16 validation-scale ceiling 제거; 시간/데이터 mode 양쪽에서 의미 있는 context length 가능.
• Flash attention의 O(1) working-set 특성을 충실히 모델링.
• 작고 일반적인 primitive (모든 tiled kernel이 이점).

Negative

• Use-after-scope 버그는 stale 바이트를 조용히 읽는다. D2 계약, attention helper 안에서 scope discipline 공유, (선택적으로) rewound 영역을 poison시키는 debug build로 완화.
• ADR-0062와 조정 필요: prefetch buffer는 tile iteration 사이에 살아 있어야 하므로 scoped arena가 아니라 persistent arena에 속한다.

Test Requirements

1. Recycling: tl.scratch_scope() 내부의 N tile loop에서 peak _scratch_cursor 가 N과 무관하게 한 tile의 footprint로 bound됨 (오늘날은 선형 증가 후 overflow).
2. Correctness: scope이 있는 flash sweep이 scope이 없는 동일 sweep 과 (재활용 없이도 수렴하는 작은 N 에서) 동일한 O 를 산출 (Phase 2).
3. Long context: scope 없이 1 MiB를 초과하는 N 의 sweep이 완료 (오늘날 S=16 cap이 회피하는 그 실패).
4. Persistent-vs-scoped isolation: scope 외부에서 할당된 running (m,l,O) 가 __exit__ 이후에도 올바른 값을 유지.
5. Nesting: 중첩된 scope이 올바른 save-point로 rewind.