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

@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:

# 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.

# 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.

# 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):

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

With pe_index gone, migration to structural coordinates is forced:

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

The dict population in _ensure_allocators is also tuple-keyed:

# 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.