Commit Graph

1 Commits

Author SHA1 Message Date
mukesh e2fe33180d 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>
2026-06-04 12:39:30 -07:00