ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
222815d3742187880399a5a9c8746fcfff6fb462
kernbench2/docs/adr-proposed/ADR-0028-par-dtensor-support.md
T
ywkang 687c98086d ADR housekeeping: category prefixes, lifecycle folders, retroactive 0034-0037 
Filename + lifecycle:
- ADR rename to ADR-NNNN-<cat>-title.md with 8 3-letter category prefixes
  (dev / mem / lat / prog / algo / par / api / ver). Numbers stay immutable.
- ADR Lifecycle split into 3 folders, documented in CLAUDE.md Part 2:
  docs/adr/ (Accepted), docs/adr-proposed/ (Proposed/Stub/Draft),
  docs/adr-history/ (Superseded/Merged). Status field gains "Draft" for
  retroactive docs pending verification.

Merges (one ADR per topic, no change-history annotations):
- ADR-0017 absorbs ADR-0019 (Cube NOC + per-PE HBM connectivity, 10 D-items)
- ADR-0014 absorbs ADR-0021 (PE pipeline execution model, 8 D-items incl.
  TileToken self-routing and multi-op composite epilogue scope)
- ADR-0023 absorbs docs/ipcq-dma-codesign-hw.md as new "HW Realization
  Notes (Informative)" section (D16-D23 + Open HW Questions). codesign-hw.md
  deleted; ADR-0019/0021 moved to adr-history with one-line stub status

Retroactive documentation (G4 closures, code-verified):
- ADR-0037 forwarding component (TransitComponent: first-flit overhead,
  serial worker, path-based routing, single impl/multiple names)
- ADR-0036 IO_CPU component (target_start_ns global barrier stamping,
  per-cube fan-out, response aggregation)
- ADR-0035 M_CPU & M_CPU.DMA component (3 fan-out paths, DMA Resources,
  target_start_ns passthrough)
- ADR-0034 HBM controller internal design (per-PC state, address-based
  selection, flit-aware per-flit commit, async finalize, command-only
  fallback path)

Content updates:
- ADR-0010 expanded to full CLI surface (run/probe/web), retitled
  "Command Line Interface and Execution Semantics"
- ADR-0007 D2 rewritten to current state; ADR-0015 supersession notes pruned
- ADR-0005 wrapped in Decision header with D1-D5; ADR-0022 metadata
  block replaced with standard Status header
- ADR-0024 trimmed to rank=SIP launcher essentials (D1-D4);
  ADR-0027 cleaned of supersession history
- ADR-0033 D6 cleanup: address-based PC selection moved out of future-work
  (now documented in ADR-0034 D3); related D1/D3 wording realigned
- Cross-references back-filled in 5 ADRs (G3 gaps closed)

Onboarding docs split:
- docs/onboarding/ created
- moved: hw-architecture-overview.md, latency-model.md, di-presentation.md,
  ccl-author-guide{,.en}.md
- references updated in README, ADR-0023{,.en}, src/kernbench/ccl/__init__.py

Source / test / yaml: ADR-NNNN cross-references in docstrings and YAML
comments updated after the merges (ADR-0021->0014 D6, ADR-0019->0017 D8).
No behavior change.

Tooling:
- tools/verify_adr_lang_pairs.py + tests/test_verify_adr_lang_pairs.py
  (ADR EN/KO pair invariant checker)
- .claude/commands/report.md tracked (/report slash command)
- .gitignore: allow .claude/commands/*.md while keeping settings files ignored

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-20 01:15:55 -07:00

6.4 KiB
Raw Blame History

ADR-0028: DTensor Support — 선언적 분산 텐서 (Stub / Future)

Status

Stub (Future Work)

Context

목표

선언적 분산 텐서 추상화(PyTorch 2.x DTensor 스타일)를 KernBench에 도입하기 위한 디자인 공간 preliminary exploration. 본 ADR은 구현 계획이 아닌 future 작업의 파일 플레이스홀더 + 초기 질문 목록이다.

Megatron-style TP와의 차이 (Why DTensor)

관점 Megatron (ADR-0027) DTensor (이 ADR)
표현 명시적 parallel layer 텐서 + placement spec
호출 형태 ColumnParallelLinear(...) distribute_tensor(x, mesh, [Shard(1)])
Collective 삽입 레이어 내부 명시 연산 dispatch가 자동
Learning curve 낮음 (명시적) 중~높음 (선언적 의미 이해)
유연성 레이어 단위로 고정 레이어 경계 무관, 어디서나
KernBench에 선행 필요한 것 launcher (ADR-0024) + TP (0027) 그 + operator dispatch overhaul

DTensor는 operator-level에서 "텐서의 placement를 보고 자동으로 collective 삽입". KernBench가 이를 지원하려면 operator dispatch layer에 placement-aware rewriting이 들어가야 한다. 이는 비-trivial.

현재 상태

  • KernBench는 operator dispatch 레이어가 없음 (torch.matmul은 없음; kernel launch로 대체).
  • DPPolicy는 정적 placement metadata를 보유 (ADR-0026 후: intra-device only).
  • ADR-0024 launcher가 rank / device 개념 제공.
  • Megatron-style TP (ADR-0027)가 명시적 대안으로 기능할 것.

Preliminary decision space

DQ1. PyTorch DTensor API 수용 범위

  • DeviceMesh: rank들의 논리적 grid.
  • Placements: Shard(dim), Replicate(), Partial(reduce_op).
  • distribute_tensor(tensor, device_mesh, placements): local tensor → DTensor.
  • Redistribute: dt.redistribute(new_placements)로 collective 자동 삽입.
  • Operator forward: dt @ dt, dt + dt 등 → 적절한 collective 자동 dispatch.

KernBench가 어느 수준까지 지원할지 결정 필요. 최소: distribute_tensor + redistribute. 최대: 모든 operator overloading.

DQ2. Operator dispatch 레이어

KernBench에서 dt @ dt를 정의하려면 Tensor의 __matmul__이 placement를 보고 적절한 action 수행:

  • 둘 다 replicated → local matmul
  • A column-sharded, B row-sharded → local matmul + all-reduce (RowParallel)
  • A replicated, B column-sharded → local matmul (ColumnParallel)
  • etc.

이는 Megatron-style의 자동화된 버전. Kernel은 기존 matmul kernel 사용.

DQ3. DeviceMesh와 기존 topology

KernBench topology는 이미 SIP/cube/PE 계층. DTensor의 DeviceMesh는 추상 (tp_size, dp_size, ...) grid. 매핑:

  • 1D mesh of size = SIP count → rank = SIP
  • 2D mesh (tp × dp) → SIP을 그룹 분할 (pure TP 대신 mixed parallelism)

초기엔 1D mesh만, DP × TP 2D는 future.

DQ4. Placement의 intra-device (DP) 통합

KernBench 특이점: 한 rank 내부에서 DPPolicy로 cube/PE에 분산. DTensor는 device 내부를 보지 않음. 통합:

  • DTensor placement = rank (SIP) 간 분산
  • 각 rank의 local tensor는 여전히 DPPolicy로 cube/PE 배치
  • → DTensor wrapper가 local tensor의 DPPolicy도 보관

DQ5. Collective 자동 삽입 지점

redistribute 또는 operator forward 시. ADR-0024의 submit+yield+wait 패턴을 자동으로 호출하는 형태. _launch_submit 내부화.

DQ6. Autograd

DTensor는 autograd와 상호작용 (backward에서 reverse collective). KernBench가 backward 지원하기 전까지는 forward-only DTensor.

Open questions (to resolve before real design)

  1. 우선순위: Megatron-style(ADR-0027)이 먼저 안착한 후 DTensor를 위에 얹는가, 아니면 공통 lower-layer를 먼저 설계하는가?
  2. 호환성 목표: PyTorch DTensor API와 몇 %까지 일치시키는가? 독자 API vs 거의 동일?
  3. Operator dispatch: KernBench Tensor 클래스에 __matmul__ 등 연산자 overloading을 도입하는가? (현재는 kernel launch만)
  4. Redistribute 정책: Shard(0) → Replicate() 변환 시 어떤 collective 사용? all_gather가 없으면 구현 전까지 제약.
  5. Mesh × DPPolicy interaction: 하나의 DTensor가 2개 layer 분산을 갖는 경우의 metadata 표현.
  6. Partial placement의 reduce 시점: 자동 vs 명시 redistribute 호출.
  7. Bench authoring impact: 기존 Megatron-style bench가 DTensor 기반으로 얼마나 쉽게 포팅되는가?

Non-goals (for future real ADR)

  • 이번 stub에서 API 확정. Future ADR에서 구체화.
  • Implementation timeline. 이번 round에서는 설계 공간 매핑만.

Dependencies (potential)

  • ADR-0024 (launcher): rank / device 기반
  • ADR-0026 (DPPolicy cleanup): DTensor placement와의 분리 명확화
  • ADR-0027 (Megatron TP): 실용 TP 패턴 경험을 DTensor 설계로 환류
  • Future ADR (operator dispatch layer): KernBench Tensor에 operator overloading 도입

Expected consequences (hypothetical)

Positive

  • PyTorch training code 이식이 매우 쉬워짐 (DTensor 코드 그대로).
  • TP + DP + 더 복잡한 parallelism을 하나의 추상화로 표현.
  • Collective 삽입이 자동 → bench 작성자 부담 감소.

Negative

  • Operator dispatch layer 신규 구축 → 상당한 엔지니어링.
  • Implicit behavior 증가 → 디버깅 / 성능 분석 복잡.
  • KernBench의 "명시적 kernel launch" 철학과 tension.

Action

  • Phase 1 (현재): 본 stub 유지. Megatron-style (ADR-0027) 먼저 구현 + 사용 경험 축적.
  • Phase 2 (future): 사용 경험을 바탕으로 본 ADR을 real design으로 승격. 위 Open questions에 대한 답을 제시.
  • Phase 3 (future): Implementation.

현재 구현 작업은 없음. 디자인 공간 매핑만.

Affected files

본 ADR은 stub이므로 production 변경 없음. Future real ADR에서 갱신될 파일 후보:

File 예상 변경 (future)
src/kernbench/dtensor/__init__.py 신규 패키지
src/kernbench/dtensor/device_mesh.py DeviceMesh
src/kernbench/dtensor/placements.py Shard/Replicate/Partial
src/kernbench/dtensor/api.py distribute_tensor, redistribute
src/kernbench/dtensor/ops/*.py Operator dispatch (matmul 등)
src/kernbench/runtime_api/tensor.py Tensor에 __matmul__ 등 추가
Reference in New Issue View Git Blame Copy Permalink