kernbench2/docs/adr/ADR-0024-par-sip-tp-launcher.md

ADR-0024: SIP-level Launcher — rank = SIP

Status

Accepted

Context

목표

torch.distributed collective 호출의 참여 단위(rank)를 SIP(device) 경계에 맞춘다. 실제 PyTorch DDP/TP 스크립트와 호스트 레벨에서 구분 없이 읽히는 bench 코드를 목표로 한다.

real PyTorch와 비교:

차원	real PyTorch	KernBench
프로세스 모델	N개 프로세스, 각 1 GPU	1 프로세스, N greenlet, 각 1 SIP
get_rank()	RANK env var	greenlet-local 레지스트리
get_world_size()	WORLD_SIZE env var	topology의 SIP 수
torch.cuda.set_device(r) (real) / torch.ahbm.set_device(r) (KernBench)	rank → GPU	rank → SIP
mp.spawn	OS 프로세스 fork	greenlet fan-out

풀어야 할 문제

1. 공개 API에서 rank = SIP — bench worker가 PE 개념을 알지 않도록.
2. Greenlet-local rank/device tracking — 1-프로세스 모델 안에서 각 worker greenlet이 자기 rank / 자기 SIP를 정확히 식별.
3. Tensor placement = structural (sip, cube, pe) — rank가 SIP이면 기본 텐서 배치도 구조적 좌표로 표현되어야 함.

Non-problem (이 ADR 밖)

• IPCQ direction addressing → ADR-0025
• DPPolicy.sip/num_sips 제거 → ADR-0026
• Megatron-style TP → ADR-0027
• DTensor → ADR-0028 (future)
• Worker scheduling / mp.spawn / collective drain / exception cleanup → ADR-0027 D0/D1
• Collective algorithm 구현 (intercube_allreduce, SFR config) → ADR-0032

Decision

D1. rank = SIP (world_size 해석)

def _resolve_world_size(self) -> int:
    if "world_size" in self._merged:
        return int(self._merged["world_size"])
    defaults = self._cfg_all.get("defaults", {})
    if "world_size" in defaults:
        return int(defaults["world_size"])
    spec = self.ctx.spec or {}
    return int(spec.get("system", {}).get("sips", {}).get("count", 1))

우선순위: 알고리즘 override > defaults override > SIP count. ccl.yaml override는 legacy "rank = PE" 테스트 경로로 유지.

D2. Greenlet-local rank registry (+ debug warning)

class DistributedContext:
    def __init__(self):
        self._backend = None
        self._rank_by_greenlet: dict = {}

    def _bind_rank(self, g, rank: int) -> None:
        self._rank_by_greenlet[g] = int(rank)

    def get_rank(self) -> int:
        self._ensure_initialized()
        from greenlet import getcurrent
        g = getcurrent()
        if g not in self._rank_by_greenlet:
            if os.environ.get("KERNBENCH_DEBUG"):
                warnings.warn(
                    "get_rank() called outside a bound greenlet — returning 0. "
                    "Likely a bug unless running single-driver."
                )
            return 0
        return int(self._rank_by_greenlet[g])

D3. torch.ahbm.set_device(rank) — SIP 바인딩

KernBench 백엔드 이름은 ahbm (ADR-0023). Real PyTorch는 torch.cuda.set_device(r)이지만 우리는 CUDA가 아니므로 honestly-named namespace를 사용한다.

class _AhbmNamespace:
    """torch.ahbm — per-greenlet SIP device binding.

    Real-PyTorch parity idiom: ``torch.cuda.set_device(rank)``. Since
    KernBench's backend is 'ahbm' (not CUDA), we expose the equivalent
    API under ``torch.ahbm`` to avoid pretending to be a CUDA runtime.
    """

    def __init__(self):
        self._device_by_greenlet: dict = {}

    def set_device(self, device: int) -> None:
        from greenlet import getcurrent
        self._device_by_greenlet[getcurrent()] = int(device)

    def current_device(self) -> int | None:
        from greenlet import getcurrent
        return self._device_by_greenlet.get(getcurrent())

# Attached to RuntimeContext as `self.ahbm = _AhbmNamespace()`.
# Bench code: `torch.ahbm.set_device(rank)` mirrors `torch.cuda.set_device`.

PyTorch 2.x style 병행 지원: 최신 PyTorch는 device-agnostic한 torch.accelerator 네임스페이스를 지향 (torch.accelerator.set_device_index(r), torch.accelerator.current_device_index()). Device vendor에 종속되지 않는 코드를 쓰려는 사용자를 위해 KernBench도 이 표면을 병행 지원한다.

class _AcceleratorNamespace:
    """torch.accelerator — device-agnostic API (PyTorch 2.x style).

    Aliases torch.ahbm for bench code that prefers device-neutral idiom:
        torch.accelerator.set_device_index(rank)
        torch.accelerator.current_device_index()
    """

    def __init__(self, ahbm: _AhbmNamespace):
        self._ahbm = ahbm

    def set_device_index(self, device: int) -> None:
        self._ahbm.set_device(device)

    def current_device_index(self) -> int | None:
        return self._ahbm.current_device()

# RuntimeContext
self.ahbm = _AhbmNamespace()
self.accelerator = _AcceleratorNamespace(self.ahbm)   # alias

Bench 작성자는 다음 중 하나를 선택 — 둘 다 내부적으로 같은 레지스트리를 보유:

torch.ahbm.set_device(rank)                   # KernBench-native, explicit backend
torch.accelerator.set_device_index(rank)      # PyTorch 2.x device-agnostic

D4. Tensor placement = structural (sip, cube, pe) 좌표

resolve_dp_policy가 target_sip을 직접 받아 구조적 좌표로 placement 생성. 세부는 ADR-0026.

# RuntimeContext._create_tensor
current_sip = self.ahbm.current_device()          # (D3 naming)
if current_sip is None:
    current_sip = 0  # single-driver fallback (D2와 일관)
placement = resolve_dp_policy(
    dp, shape=shape_2d, itemsize=itemsize,
    num_pe=eff_num_pe, num_cubes=eff_num_cubes,
    target_sip=current_sip,
)

Post-hoc pe_index shifting 없음 — ShardSpec이 (sip, cube, pe) 구조적 좌표를 직접 보유. ShardSpec 상세는 ADR-0026.

---

Dependencies

• ADR-0023 (IPCQ): backend ahbm namespace의 기원.
• ADR-0026 (DPPolicy intra-device): D4의 resolve_dp_policy 시그니처와 ShardSpec의 구조적 좌표 표현.
• ADR-0027 (Megatron TP + scheduler): worker scheduling, mp.spawn, collective drain, exception cleanup의 구현 기준.

---

Non-goals

• IPCQ protocol 수정: ADR-0023 유지.
• DPPolicy 필드 정리: ADR-0026.
• Megatron-style TP: ADR-0027.
• Worker scheduling / spawn / drain / exception cleanup: ADR-0027 D0/D1.
• Collective algorithm 구현: ADR-0032.
• Multi-node (프로세스 간): 단일 프로세스.

---

Consequences

Positive

• Bench = real PyTorch DDP (공개 API 관점).
• Greenlet-local rank: 1-프로세스 모델에서 cross-rank correctness 가능.
• Structural placement 좌표: ADR-0026 / ADR-0027 / ADR-0032의 다른 ADR이 (sip, cube, pe) 3튜플 위에서 일관되게 동작.

Neutral

• IPCQ PE-level protocol (ADR-0023) 불변.
• IO_CPU 역할 불변 (기존 transit 그대로).