ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ADR-0026: DPPolicy intra-device only + ShardSpec structural coords

Browse Source 
DPPolicy no longer carries a cross-SIP axis. SIP-level placement is
solely controlled by torch.ahbm.set_device(rank) (ADR-0024); DPPolicy
itself describes only the cube × PE layout within one SIP. ShardSpec
switches to structural (sip, cube, pe) coordinates; the flat pe_index
field/property is fully removed — silent drift between global-flat and
SIP-local interpretations was a foot-gun flagged by ADR-0024 D11.

Breaking API (explicit TypeError / AttributeError):
- DPPolicy(sip=...) / DPPolicy(num_sips=...) -> TypeError
- ShardSpec.pe_index -> AttributeError
- ShardSpec(pe_index=...) -> TypeError
- resolve_dp_policy now takes target_sip= (required), no num_sips.

Downstream migration:
- PE allocator dict keyed by (sip, cube, pe) tuples, in both
  _ensure_allocators and _free_tensor. deploy_tensor uses tuple lookup.
- _create_tensor passes target_sip=current_sip; post-hoc pe_index
  shifting removed entirely.
- launch._compute_local_shape drops the dp.sip branch.
- Internal resolvers (column_wise / row_wise / replicate / tiled_*)
  return _LocalPeShard (cube-local identifier) instead of ShardSpec —
  resolve_dp_policy lifts them to full structural coords.

Tests:
- New tests/test_adr0026_dppolicy_intra_device.py (12 tests) pins the
  contract end-to-end.
- test_sip_parallel.py rewritten: SIP composition now modeled as two
  resolve_dp_policy(target_sip=...) calls (ADR-0024 launcher style).
- Call-site migration: test_tensor, test_va_integration, test_va_offset,
  test_runtime_api_tensor, test_tl_recv_async, test_ccl_* and benches
  gemm_single_pe, gpt3_qkv, va_offset_verify, ccl_allreduce (legacy
  branch) all use intra-device DPPolicy and structural ShardSpec.

Result: 523 passed, 1 strict xfail (ring_default_ws — unchanged
ADR-0024 Phase B blocker; architectural fix deferred to ADR-0027).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-04-14 13:02:19 -07:00
parent 787409ced1
commit 357cab525b
20 changed files with 549 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/policy/placement/dp.py
+90 -70
View File
@@ -1,3 +1,14 @@
"""Data-parallel placement policy (ADR-0026: intra-device only).
``DPPolicy`` describes how a tensor is sharded *within a single SIP* across
that SIP's cubes and PEs. Crossing the SIP boundary is not a DPPolicy
concern: ADR-0024's ``torch.ahbm.set_device(rank)`` picks the SIP, and
Megatron-style TP (ADR-0027) expresses multi-SIP tensors when needed.
``ShardSpec`` is expressed in structural ``(sip, cube, pe)`` coordinates.
The former flat ``pe_index`` field/property is fully removed — callers
needing a flat integer key compute it explicitly at the call site.
"""
from __future__ import annotations
from dataclasses import dataclass
@@ -7,25 +18,58 @@ from typing import Literal
@dataclass(frozen=True)
class DPPolicy:
"""Three-level data-parallel policy: sip-level + cube-level + pe-level.
"""Intra-device (cube × PE) data-parallel policy.
Policies:
SIP-level placement is controlled by ``torch.ahbm.set_device(rank)``
(ADR-0024). For tensors that must cross SIP boundaries, use
Megatron-style parallel layers (ADR-0027). DPPolicy itself never
crosses a SIP boundary.
Policies (per axis):
- "replicate": full copy at each unit
- "column_wise": split K (column) axis across units
- "row_wise": split M (row) axis across units
Optional overrides (default None = use topology dimensions):
- num_pes: override PEs per cube (e.g., 1 for single-PE test)
- num_cubes: override cubes per SIP (e.g., 1 for single-cube test)
- num_sips: override SIP count
Optional overrides (``None`` = use topology dimensions):
- num_pes: override PEs per cube
- num_cubes: override cubes per SIP
"""
sip: Literal["replicate", "column_wise", "row_wise"] = "replicate"
cube: Literal["replicate", "column_wise", "row_wise"] = "replicate"
pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
num_pes: int | None = None
num_cubes: int | None = None
num_sips: int | None = None
@dataclass(frozen=True)
class ShardSpec:
"""Structural shard placement — ``(sip, cube, pe)`` coord (ADR-0026).
Global-flat ``pe_index`` was removed: callers must use structural
coords directly. If a flat integer key is needed in a local context
(e.g. internal dict lookup), compute it explicitly at the call site
and do not expose it in any public API.
"""
sip: int
cube: int
pe: int
offset_bytes: int
nbytes: int
@dataclass(frozen=True)
class _LocalPeShard:
"""Internal — PE resolver's return type (ADR-0026 D3).
Holds a cube-local PE identifier (``local_pe``) plus the shard's
byte payload. Lifted into ``ShardSpec`` with full ``(sip, cube, pe)``
coordinates inside ``resolve_dp_policy``.
"""
local_pe: int
offset_bytes: int
nbytes: int
def _split_shape(
@@ -52,14 +96,13 @@ def resolve_dp_policy(
itemsize: int,
num_pe: int,
num_cubes: int = 1,
num_sips: int = 1,
target_sip: int,
) -> list[ShardSpec]:
"""Resolve a DPPolicy into a list[ShardSpec] with three-level resolution.
"""Resolve a DPPolicy into a list[ShardSpec] on a single SIP.
SIP-level → cube-level → pe-level.
num_cubes is cubes per SIP (not total).
ShardSpec.pe_index uses flat indexing:
sip_id * num_cubes * num_pe + cube_id * num_pe + pe_id
Two-level resolution (cube × PE) within ``target_sip``. Each returned
``ShardSpec`` carries ``sip=target_sip`` and cube/pe local to the SIP.
No SIP-level split — DPPolicy is intra-device only (ADR-0026).
"""
_PE_RESOLVERS = {
"replicate": replicate,
@@ -70,84 +113,61 @@ def resolve_dp_policy(
if resolver is None:
raise ValueError(f"Unknown pe-level policy: {policy.pe}")
cubes_per_sip = num_cubes
all_shards: list[ShardSpec] = []
# Level 1: SIP
sip_splits = _split_shape(policy.sip, shape, num_sips, itemsize)
# Level 1: cube within SIP
cube_splits = _split_shape(policy.cube, shape, num_cubes, itemsize)
for sip_id, (sip_shape, sip_offset) in enumerate(sip_splits):
# Level 2: Cube within SIP
cube_splits = _split_shape(policy.cube, sip_shape, cubes_per_sip, itemsize)
for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
# Level 2: PE within cube — resolver returns _LocalPeShard
local_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
# Level 3: PE within cube
pe_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
for ps in pe_shards:
flat_idx = (
sip_id * cubes_per_sip * num_pe
+ cube_id * num_pe
+ ps.pe_index
)
all_shards.append(ShardSpec(
pe_index=flat_idx,
offset_bytes=sip_offset + cube_offset + ps.offset_bytes,
nbytes=ps.nbytes,
))
for ls in local_shards:
all_shards.append(ShardSpec(
sip=target_sip,
cube=cube_id,
pe=ls.local_pe,
offset_bytes=cube_offset + ls.offset_bytes,
nbytes=ls.nbytes,
))
return all_shards
@dataclass(frozen=True)
class ShardSpec:
pe_index: int
offset_bytes: int
nbytes: int
def column_wise(
*, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]:
) -> list[_LocalPeShard]:
"""Split K axis into num_pe equal parts. Each PE gets (M, K/P)."""
M, K = shape
chunk_k = K // num_pe
chunk_bytes = M * chunk_k * itemsize
shards = []
for i in range(num_pe):
shards.append(ShardSpec(
pe_index=i,
offset_bytes=i * chunk_bytes,
nbytes=chunk_bytes,
))
return shards
return [
_LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
for i in range(num_pe)
]
def row_wise(
*, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]:
) -> list[_LocalPeShard]:
"""Split M axis into num_pe equal parts. Each PE gets (M/P, K)."""
M, K = shape
chunk_m = M // num_pe
chunk_bytes = chunk_m * K * itemsize
shards = []
for i in range(num_pe):
shards.append(ShardSpec(
pe_index=i,
offset_bytes=i * chunk_bytes,
nbytes=chunk_bytes,
))
return shards
return [
_LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
for i in range(num_pe)
]
def replicate(
*, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]:
) -> list[_LocalPeShard]:
"""Full copy per PE. Each PE gets (M, K)."""
M, K = shape
full_bytes = M * K * itemsize
return [
ShardSpec(pe_index=i, offset_bytes=0, nbytes=full_bytes)
_LocalPeShard(local_pe=i, offset_bytes=0, nbytes=full_bytes)
for i in range(num_pe)
]
@@ -155,20 +175,20 @@ def replicate(
def tiled_column_major(
*, shape: tuple[int, int], itemsize: int, num_pe: int,
tile_m: int, tile_k: int,
) -> list[ShardSpec]:
) -> list[_LocalPeShard]:
"""2D tiling, column-major order (K axis first), round-robin across PEs."""
M, K = shape
tiles_m = ceil(M / tile_m)
tiles_k = ceil(K / tile_k)
tile_bytes = tile_m * tile_k * itemsize
row_bytes = K * itemsize
shards = []
shards: list[_LocalPeShard] = []
idx = 0
for mi in range(tiles_m):
for ki in range(tiles_k):
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
shards.append(ShardSpec(
pe_index=idx % num_pe,
shards.append(_LocalPeShard(
local_pe=idx % num_pe,
offset_bytes=offset,
nbytes=tile_bytes,
))
@@ -179,20 +199,20 @@ def tiled_column_major(
def tiled_row_major(
*, shape: tuple[int, int], itemsize: int, num_pe: int,
tile_m: int, tile_k: int,
) -> list[ShardSpec]:
) -> list[_LocalPeShard]:
"""2D tiling, row-major order (M axis first), round-robin across PEs."""
M, K = shape
tiles_m = ceil(M / tile_m)
tiles_k = ceil(K / tile_k)
tile_bytes = tile_m * tile_k * itemsize
row_bytes = K * itemsize
shards = []
shards: list[_LocalPeShard] = []
idx = 0
for ki in range(tiles_k):
for mi in range(tiles_m):
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
shards.append(ShardSpec(
pe_index=idx % num_pe,
shards.append(_LocalPeShard(
local_pe=idx % num_pe,
offset_bytes=offset,
nbytes=tile_bytes,
))
src/kernbench/runtime_api/context.py
+14 -45
View File
@@ -89,7 +89,7 @@ class RuntimeContext:
_handles: list[RequestHandle] = field(default_factory=list, init=False)
_completed: set[RequestHandle] = field(default_factory=set, init=False)
_allocators: dict[int, Any] = field(default_factory=dict, init=False)
_allocators: dict[tuple[int, int, int], Any] = field(default_factory=dict, init=False)
_va_allocator: Any = field(default=None, init=False)
_tensor_counter: int = field(default=0, init=False)
_traces: list[dict] = field(default_factory=list, init=False)
@@ -270,12 +270,7 @@ class RuntimeContext:
# Return PA space
if self._allocators:
for shard in handle.shards:
flat_idx = (
shard.sip * self._num_cubes * self._pes_per_cube
+ shard.cube * self._pes_per_cube
+ shard.pe
)
alloc = self._allocators.get(flat_idx)
alloc = self._allocators.get((shard.sip, shard.cube, shard.pe))
if alloc is not None:
from kernbench.policy.address.phyaddr import PhysAddr
alloc.free_hbm(PhysAddr.decode(shard.pa), shard.nbytes)
@@ -339,17 +334,15 @@ class RuntimeContext:
tcm_scheduler_reserved_bytes=4 * (1 << 20),
sram_bytes_per_cube=32 * (1 << 20),
)
# Create allocators scoped to target SIP(s) only
# Flat index: sip_id * cubes_per_sip * pes_per_cube + cube_id * pes_per_cube + pe_id
# Create allocators scoped to target SIP(s) only.
# ADR-0026 D5: dict key is the structural (sip, cube, pe) tuple.
self._pes_per_cube = pes_per_cube
self._num_cubes = cubes_per_sip
self._num_sips = sip_count
cubes_x_pes = cubes_per_sip * pes_per_cube
for sip_id in sip_range:
for cube_id in range(cubes_per_sip):
for pe_id in range(pes_per_cube):
flat_idx = sip_id * cubes_x_pes + cube_id * pes_per_cube + pe_id
self._allocators[flat_idx] = PEMemAllocator(
self._allocators[(sip_id, cube_id, pe_id)] = PEMemAllocator(
rack_id=0, sip_id=sip_id, cube_id=cube_id, pe_id=pe_id, cfg=cfg,
)
@@ -436,44 +429,23 @@ class RuntimeContext:
# DPPolicy overrides take precedence over topology dimensions
eff_num_pe = dp.num_pes if dp.num_pes is not None else self._pes_per_cube
eff_num_cubes = dp.num_cubes if dp.num_cubes is not None else self._num_cubes
# ADR-0024 D11: if torch.ahbm.set_device(r) is active AND DPPolicy
# leaves the SIP dimension at its default (replicate + no num_sips
# override), scope the tensor to SIP r only.
# NOTE: this path uses post-hoc pe_index shifting as a temporary
# measure; ADR-0026 replaces it with structural (sip, cube, pe)
# coords in ShardSpec.
# ADR-0026 D4: resolve structural coords directly at resolve time.
# ``torch.ahbm.set_device(rank)`` (ADR-0024 D10) selects the target
# SIP; if unset, fall back to SIP 0 for single-driver compatibility.
current_sip = (
self.ahbm.current_device() if hasattr(self, "ahbm") else None
)
scope_to_current_sip = (
current_sip is not None
and dp.sip == "replicate"
and dp.num_sips is None
)
if scope_to_current_sip:
eff_num_sips = 1
else:
eff_num_sips = (
dp.num_sips if dp.num_sips is not None else self._num_sips
)
if current_sip is None:
current_sip = 0
placement = resolve_dp_policy(
dp, shape=shape_2d, itemsize=itemsize,
num_pe=eff_num_pe, num_cubes=eff_num_cubes,
num_sips=eff_num_sips,
target_sip=int(current_sip),
)
if scope_to_current_sip:
from kernbench.policy.placement.dp import ShardSpec as _SS
sip_stride = self._num_cubes * self._pes_per_cube
offset = int(current_sip) * sip_stride
placement = [
_SS(pe_index=s.pe_index + offset,
offset_bytes=s.offset_bytes, nbytes=s.nbytes)
for s in placement
]
# Infer target_pe from placement using local (within-cube) PE IDs.
# This ensures M_CPU only fans out to PEs that own shards, not all PEs.
local_pe_ids = sorted({s.pe_index % eff_num_pe for s in placement})
local_pe_ids = sorted({s.pe for s in placement})
if len(local_pe_ids) == 1:
target_pe: int | tuple[int, ...] | str = local_pe_ids[0]
elif len(local_pe_ids) == eff_num_pe and eff_num_pe == self._pes_per_cube:
@@ -669,11 +641,8 @@ class RuntimeContext:
dp = t._dp_metadata.dp_policy if t._dp_metadata else None
if dp is None:
return t.shape
if dp.sip != "replicate":
if dp.sip == "column_wise":
K = K // self._num_sips
elif dp.sip == "row_wise":
M = M // self._num_sips
# ADR-0026: DPPolicy no longer crosses SIP boundaries; cube + PE
# are the only axes that shrink the local shape.
if dp.cube != "replicate":
if dp.cube == "column_wise":
K = K // self._num_cubes
src/kernbench/runtime_api/tensor.py
+7 -6
View File
@@ -72,7 +72,7 @@ def deploy_tensor(
shape: tuple[int, ...],
dtype: str,
placement: list[ShardSpec],
allocators: dict[int, PEMemAllocator],
allocators: dict[tuple[int, int, int], PEMemAllocator],
mem_kind: Literal["hbm", "tcm"] = "hbm",
va_allocator=None,
) -> TensorHandle:
@@ -86,15 +86,15 @@ def deploy_tensor(
shards: list[TensorShard] = []
for spec in placement:
alloc = allocators[spec.pe_index]
alloc = allocators[(spec.sip, spec.cube, spec.pe)]
if mem_kind == "hbm":
pa = alloc.alloc_hbm(spec.nbytes)
else:
pa = alloc.alloc_tcm(spec.nbytes)
shards.append(TensorShard(
sip=alloc._sip_id,
cube=alloc._cube_id,
pe=alloc._pe_id,
sip=spec.sip,
cube=spec.cube,
pe=spec.pe,
pa=pa.encode(),
nbytes=spec.nbytes,
offset_bytes=spec.offset_bytes,
@@ -394,7 +394,8 @@ class Tensor:
) -> Tensor:
"""Set DP placement metadata (like torch.Tensor.to())."""
if placement is None:
placement = [ShardSpec(pe_index=0, offset_bytes=0, nbytes=self.nbytes)]
placement = [ShardSpec(sip=0, cube=0, pe=0,
offset_bytes=0, nbytes=self.nbytes)]
self._dp_metadata = DPMetadata(
placement=placement, dp_policy=dp_policy,
sip=sip, cube=cube, target_pe=target_pe,