kernbench2/docs/adr/ADR-0026-par-dppolicy-intra-device.md

# ADR-0026: DPPolicy = Intra-Device Only — remove sip/num_sips fields

## Status

Accepted (Revision 6 — cube_start added 2026-06-04; Revision 5 landed 2026-04-14)

## Context

### Goal

Clarify `DPPolicy` as a pure intra-device abstraction that only expresses
**cube × PE distribution within a single device (SIP)**. Inter-SIP
distribution (TP) is split into a separate layer (handled by ADR-0024's
`torch.ahbm.set_device(rank)` or by ADR-0027's Megatron-style parallel
layers).

## Decision

### D1. Remove `sip` + `num_sips` fields from `DPPolicy`

```python
@dataclass(frozen=True)
class DPPolicy:
    """Intra-device (cube × PE) data-parallel policy.

    SIP-level placement is controlled by ``torch.ahbm.set_device(rank)``
    (ADR-0024 D3) and, for model-level TP, by Megatron-style parallel
    layers (ADR-0027). DPPolicy does not cross SIP boundaries.
    """
    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
    cube_start: int = 0   # Revision 6 — first cube in SIP; see D8
```

Removed fields: `sip`, `num_sips`.
Added fields (Revision 6): `cube_start` — see D8.

### D2. `ShardSpec` — structural (sip, cube, pe) coordinates, `pe_index` fully removed

The current `ShardSpec.pe_index` is a **global flat index**
(`sip × cubes × pes + cube × pes + pe`). This is the form ADR-0024 D4
flagged as "abstraction leakage".

This ADR **redefines ShardSpec in structural coordinates** and **does
not even leave `pe_index` as a property**:

```python
# src/kernbench/policy/placement/dp.py (after)
@dataclass(frozen=True)
class ShardSpec:
    """Structural shard placement — intra-SIP (cube × PE) coord.

    Global-flat `pe_index` was removed in ADR-0026. Callers must use
    structural coords (sip, cube, pe) directly. If a flat integer key is
    needed (e.g. dict lookup), compute it explicitly at the call site.
    """
    sip: int              # structural — which SIP this shard lives on
    cube: int             # local within SIP
    pe: int               # local within cube
    offset_bytes: int
    nbytes: int
```

**Core principle**:
- The identity of ShardSpec is the `(sip, cube, pe)` 3-tuple.
- **No `pe_index` property either** — blocks silent semantics drift.
- Existing callers expecting global-flat get an **immediate
  `AttributeError`** on `.pe_index` access → forced migration to
  structural coordinates.
- Local contexts that genuinely need a flat integer key (e.g. internal
  dict lookup) explicitly compute
  `spec.sip * N_CUBES * N_PE + spec.cube * N_PE + spec.pe` at the call
  site.

**Justification for removing the property**: KernBench is an internal
project with a limited number of call sites. Explicit breakage
(AttributeError) is much safer than the risk of silent drift (semantics
change while the type stays int).

### D3. `resolve_dp_policy` takes `target_sip` and produces structural coordinates

Implements the contract of ADR-0024 D4. No post-hoc shifting.

```python
# src/kernbench/policy/placement/dp.py (after)

@dataclass(frozen=True)
class _LocalPeShard:
    """Internal — return value of the PE resolver. Cube-local PE id + payload."""
    local_pe: int                  # cube-local PE index (0..num_pe-1)
    offset_bytes: int
    nbytes: int


def resolve_dp_policy(
    policy: DPPolicy,
    *,
    shape: tuple[int, int],
    itemsize: int,
    num_pe: int,
    num_cubes: int = 1,
    target_sip: int,       # NEW — explicitly state which SIP to place on
) -> list[ShardSpec]:
    """2-level resolution (cube × PE) on a specified SIP.

    Returns ShardSpecs with structural coords (sip=target_sip, cube, pe).
    No SIP-level split — DPPolicy is intra-device only.
    """
    resolver = _PE_RESOLVERS[policy.pe]
    all_shards: list[ShardSpec] = []

    # Level 1: cube within SIP
    cube_splits = _split_shape(policy.cube, shape, num_cubes, itemsize)

    for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
        # Level 2: PE within cube — resolver returns _LocalPeShard (local_pe)
        local_shards = resolver(shape=cube_shape, itemsize=itemsize,
                                 num_pe=num_pe)

        for ls in local_shards:
            all_shards.append(ShardSpec(
                sip=target_sip,                              # from caller (current_device)
                cube=policy.cube_start + cube_id,            # Rev 6: shifted by D8 cube_start; default 0 = local within SIP
                pe=ls.local_pe,                              # local within cube (explicit name)
                offset_bytes=cube_offset + ls.offset_bytes,
                nbytes=ls.nbytes,
            ))

    return all_shards
```

**Internal resolvers** (`column_wise`, `row_wise`, `replicate`) return a
list of `_LocalPeShard` — the `local_pe` field name makes it **explicit
that this is a "cube-local PE identifier"**. This resolves the previous
confusion with the name `ShardSpec.pe_index`.

**Naming convention summary** (whole ADR):
- `ShardSpec.pe`: the final external API — cube-local PE (structural coord)
- `_LocalPeShard.local_pe`: the same meaning at the internal resolver stage
- `pe_index`: **removed**. Not retained anywhere, internal or external
  (additional benefit of preventing silent drift: the name does not
  reappear).

### D4. `_create_tensor` — placement directly in structural coordinates

Continuation of ADR-0024 D4. Post-hoc shifting removed; structural
coordinates are specified directly at the `resolve_dp_policy` call site.

```python
# context.py _create_tensor (after)
current_sip = self.ahbm.current_device()
if current_sip is None:
    # Single-driver fallback (consistent with ADR-0024 D2).
    # In launcher-based code, forgetting set_device() silently sticks the
    # tensor on SIP 0 — emit a warning in debug mode.
    if os.environ.get("KERNBENCH_DEBUG"):
        import warnings
        warnings.warn(
            "torch.ahbm.current_device() is None; defaulting to SIP 0. "
            "If this is a multi-rank launcher context, you likely forgot "
            "torch.ahbm.set_device(rank) inside the worker.",
            stacklevel=2,
        )
    current_sip = 0

placement = resolve_dp_policy(
    dp,
    shape=shape_2d,
    itemsize=itemsize,
    num_pe=eff_num_pe,
    num_cubes=eff_num_cubes,
    target_sip=current_sip,          # ← structural coord specified up front
)

# Each ShardSpec in placement already carries (sip=current_sip, cube=local, pe=local).
# The old post-hoc shifting block is removed entirely.
```

**Every** tensor is placed on the current device's SIP. If you need a
multi-SIP tensor, use the TP primitive of ADR-0027.

**Trade-off of the single-driver fallback**: When set_device is not
called, defaulting to SIP 0 is kept for compatibility with existing
single-driver tests. With `KERNBENCH_DEBUG=1`, a warning is emitted so
that accidentally omitting set_device in a launcher context — which would
silently place the tensor on the wrong SIP — can be detected.

### D5. Downstream — allocator lookup by structural tuple key

Existing `deploy_tensor` (`src/kernbench/runtime_api/tensor.py`):

```python
for spec in placement:
    alloc = allocators[spec.pe_index]       # ← AttributeError (property removed)
```

With `pe_index` gone, migration to structural coordinates is **forced**:

```python
for spec in placement:
    alloc = allocators[(spec.sip, spec.cube, spec.pe)]
```

The dict population in `_ensure_allocators` is also tuple-keyed:

```python
# context.py _ensure_allocators (after)
for sip_id in sip_range:
    for cube_id in range(cubes_per_sip):
        for pe_id in range(pes_per_cube):
            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,
            )
```

`_free_tensor` is the same: the old
`flat_idx = sip * ... + cube * ... + pe` computation block is removed,
and `(shard.sip, shard.cube, shard.pe)` is used directly.

**Tuple vs dataclass `PEIdentity`**: Recommend the tuple — it is simple
and hashable out of the box. A `PEIdentity` value object has the upside
of an explicit type, but the boilerplate is large and it is currently
the only key of the allocator dict, so it would be over-engineering.
Keep the tuple.

### D7. Backward compatibility — none (cleanup ADR)

This ADR is a **breaking change**.

1. `DPPolicy(sip=...)` or `DPPolicy(num_sips=...)` → `TypeError`
2. `ShardSpec.pe_index` access → `AttributeError`

Both are **immediate, explicit breakage**. No deprecation warning /
fallback path. KernBench is an internal project with a bounded set of
call sites, so migration happens in one pass.

**Blocking silent drift** is the main upside of fully removing the
property: code that expected a global flat could otherwise silently
receive a SIP-local result and index incorrectly — that possibility is
eliminated.

### D8. Add `cube_start: int = 0` for disjoint cube sub-meshes within a SIP

**Revision 6 addition (2026-06-04)**.

**Purpose**: address a disjoint cube sub-mesh within one SIP (e.g.
cubes 8..15 alongside cubes 0..7). Required by the GQA Llama-70B
8-KV-group headline target — two 2×4 KV-groups per SIP × 4 SIPs = 64
cubes — where the second KV-group on each SIP must land on cubes 8..15
instead of the default 0..7.

**Semantics**: `resolve_dp_policy` returns `ShardSpec` with
`cube = policy.cube_start + cube_id`, where `cube_id` iterates
`0..num_cubes-1` within the launch. Selected cubes lie in
`[cube_start, cube_start + num_cubes)` within the target SIP.

**Default**: `cube_start = 0` preserves every existing call site
bit-for-bit (`ShardSpec.cube ∈ [0, num_cubes)` as before). The CCL
milestone bench's full-SIP `DPPolicy(num_cubes=16)` continues to
produce cubes 0..15.

**Constraint**: intra-device invariant preserved.
`cube_start ∈ [0, cubes_per_sip)` and
`cube_start + num_cubes ≤ cubes_per_sip`. SIP boundary crossing remains
the job of `ahbm.set_device(rank)` (ADR-0024).

**Why scalar** (not 2D `cube_mesh_origin` or arbitrary `cube_ids` list):
consumer kernels (e.g. `_attention_mesh_mlo_2d`) assume row-major
contiguous cubes; scalar `cube_start` pairs naturally with `num_cubes`
(range = `[start, start + count)`); zero migration churn at default 0.
More general designs can be added on top later if a non-contiguous use
case appears.

**Backward compatibility**: additive change with default value. No
existing call site is forced to change. The "breaking change" stance
in D7 applies only to the original sip/num_sips removal — `cube_start`
does NOT break existing callers.

**Kernel-side note**: kernbench's `tl.program_id(axis=1)` returns the
physical cube id (ADR-0022), not a launch-local rank. Kernels that
derive ring positions from `program_id(axis=1)` must subtract
`cube_start` to recover launch-local rank when `cube_start > 0` —
otherwise they compute out-of-bounds sub-mesh positions and deadlock.
See `_attention_mesh_mlo_2d.py` for the reference pattern. A future
revision of ADR-0022 could move `program_id(axis=1)` to launch-local
rank semantics; that change would let kernels drop the explicit
subtraction.

## Dependencies

- **ADR-0024** (launcher): `set_device(rank)` and current-device scoping
  provide the SIP placement mechanism. This ADR sits on top and narrows
  DPPolicy to pure intra-device.
- **ADR-0027** (Megatron TP): the alternative path when a tensor spans
  multiple SIPs. After this ADR is applied, multi-SIP use cases move to
  ADR-0027.

---

## Non-goals

- **Redesign of `DPPolicy.cube` / `pe`**: existing
  replicate/column_wise/row_wise semantics are kept.
- **Tiling policy consolidation**: `tiled_column_major` /
  `tiled_row_major` stay as they are.
- **New multi-device tensor abstraction**: a DTensor-like is ADR-0028.

---

## Open questions

- **Default value of current_sip in `_create_tensor`**: for calls without
  set_device, whether to fall back to rank=0 (SIP 0) or to raise an
  error. The recommendation is fallback (compatibility with existing
  single-driver tests).
- **Scope of `test_sip_parallel.py` rewrite**: porting the existing unit
  tests to the launcher base while preserving their intent requires
  additional fixtures. Scoped as separate work.
- **Meaning of `num_sips=None` on `DPPolicy`**: once the field is gone,
  the concept of `num_sips` disappears entirely. The explicit answer for
  expressing multi-SIP is to use the TP primitive of ADR-0027.

**Resolved (items that were open in earlier revs)**:
- ~~Whether to keep the `ShardSpec.pe_index` property~~ → **fully
  removed** (D2)
- ~~Form of `_ensure_allocators` dict key~~ → **tuple `(sip, cube, pe)`**
  (D5)

---

## Consequences

### Positive

- **Clean conceptual separation**: DPPolicy = intra-device, TP =
  inter-device.
- **API simplification**: about a 33% reduction in DPPolicy constructor
  fields.
- **Structural-coordinate consistency**: ShardSpec is expressed as a
  `(sip, cube, pe)` tuple → abstraction leakage resolved (the ADR-0024
  D4 contract is satisfied).
- **Clear meaning of `pe_index`**: the single interpretation is
  SIP-local. If global-flat is needed, it must be made explicit.
- **Launcher-model consistency**: ADR-0024's "1 worker per SIP" model is
  the sole SIP-boundary control mechanism.

### Negative

- **Breaking change (explicit)**: `DPPolicy(sip=...)` → `TypeError`,
  `spec.pe_index` → `AttributeError`. All callers need to be fixed at
  once.
- **ShardSpec schema change**: a single `pe_index` field becomes three
  fields `sip`/`cube`/`pe`. Cascading edits downstream (`deploy_tensor`,
  `_free_tensor`, `_ensure_allocators`, `allocators` dict key, etc.).
- **No silent drift**: with the property fully removed, runtime failure
  is immediate → migration leakage is blocked at the source. (Not a
  negative but an explicit tradeoff.)
- The cost of rewriting `test_sip_parallel.py`.

### Neutral

- The meaning of the existing `cube` / `pe` fields is unchanged.