policy: add DPPolicy.cube_start for disjoint cube sub-meshes within a SIP
Adds an optional ``cube_start: int = 0`` field to ``DPPolicy``. In ``resolve_dp_policy`` the generated ``ShardSpec.cube`` is now ``policy.cube_start + cube_id`` instead of just ``cube_id``. With the default value (0) every existing call site resolves to identical ``ShardSpec.cube`` values (cubes 0..num_cubes-1). Why this is needed: the GQA Llama-70B 8-KV-group headline target lays two 2x4 KV-groups per SIP — the first on cubes 0..7 (rows 0..1), the second on cubes 8..15 (rows 2..3). Without ``cube_start``, ``DPPolicy`` can only address the first 8 row-major cubes, so a second launch on the same SIP overlaps the first. ``cube_start=8`` selects the second 2x4 sub-mesh directly. Design choice: scalar ``cube_start`` (vs a 2D ``cube_mesh_origin`` or arbitrary ``cube_ids`` list) was picked because (a) the 2D mesh-mlo kernel already assumes row-major contiguous cubes, (b) it pairs naturally with ``num_cubes`` (range = ``[start, start+count)``), and (c) zero migration churn — every existing call site stays unchanged. Scope: 5 production LOC (one field + one expression in resolve). No kernel changes here; kernels that consume ``program_id(axis=1)`` for ring arithmetic need their own follow-up (see ADR-0022 contract). Tests: 7 new unit tests in ``test_dppolicy_cube_start.py`` covering default behavior preservation (cube_start=0), shifted ranges (cube_start=8 → cubes 8..15), full-SIP CCL pattern, replicate cube policy, shard-count preservation, and uniqueness. ADR-0026 regression suite (12 tests) stays green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -33,12 +33,20 @@ class DPPolicy:
|
||||
Optional overrides (``None`` = use topology dimensions):
|
||||
- num_pes: override PEs per cube
|
||||
- num_cubes: override cubes per SIP
|
||||
|
||||
Cube range:
|
||||
- cube_start: index of the first cube in the SIP to place shards
|
||||
on. Defaults to 0. Enables disjoint sub-meshes 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).
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -125,7 +133,7 @@ def resolve_dp_policy(
|
||||
for ls in local_shards:
|
||||
all_shards.append(ShardSpec(
|
||||
sip=target_sip,
|
||||
cube=cube_id,
|
||||
cube=policy.cube_start + cube_id,
|
||||
pe=ls.local_pe,
|
||||
offset_bytes=cube_offset + ls.offset_bytes,
|
||||
nbytes=ls.nbytes,
|
||||
|
||||
Reference in New Issue
Block a user