analytical-viz: PE view keeps CP rank locator visible at TP>=8

The Per-PE layout diagram used to color-code + label CP ranks only
when cp_placement=='pe' (CP packed intra-cube). At TP>=8 the TP dim
fills the cube, cp_placement gets pushed to 'cube', the intra-cube CP
branch is skipped, and the CP rank tag silently disappeared from
cubes, per-PE headers, and the figure title. Users at TP>=8 had no
way to tell which of the CP groups the diagram was drawing.

Fix: when cp_placement=='cube' and CP>1, add
  - cube header suffix:   '[CP rank 0 of {CP}]'
  - per-PE header suffix: '  |  CP=0'
  - figure title:         'CP={CP} (showing 1 of CP groups; others identical)'

Behavior at cp_placement=='pe' is unchanged — the existing per-CP-rank
palette + 'CP={rank}' label still fire. CP=1 adds no locators at all.

test_pe_weight_layout covers four cases: TP=8/CP=4 (cube placement,
locator appears), TP=2/CP=4 forced to pe placement (regression:
per-rank labels still there, cube tag not added), CP=1 (no locators),
TP=16/CP=2 (both spilled cubes carry the CP-rank-0 tag).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/analytical_visualization/pe_weight_layout.py

@@ -114,6 +114,14 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     cube_h = 6.0
     cube_gap = 0.6
 
+    # When TP fills the cube (TP>=8), cp_placement is "cube" and CP
+    # ranks span cubes — the intra-cube CP-color branch below is
+    # skipped. Surface which CP rank this figure is showing so users
+    # aren't guessing which of `CP` groups is drawn.
+    _cube_place_tag = ""
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _cube_place_tag = f"  [CP rank 0 of {cfg.topo.cp}]"
+
     for cube_idx in range(n_cubes):
         x0 = cube_idx * (cube_w + cube_gap)
         y0 = 0
@@ -125,7 +133,8 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         )
         ax.add_patch(rect)
         ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
-                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})",
+                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})"
+                + _cube_place_tag,
                 ha="center", va="bottom", fontsize=11, fontweight="bold")
 
         pe_pad_x = 0.15
@@ -180,6 +189,13 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                     _tp_rank = pe_id_in_group % tp
                     _fc = "#f0f7ff"
                     _ec = "#3a86ff"
+                    # cp_placement=="cube" with CP>1: the whole cube is
+                    # one CP rank (rank 0 by convention — see cube title
+                    # tag above). Add the locator so users know what
+                    # rank this per-PE data belongs to.
+                    if _cp_place == "cube" and _cp_val > 1:
+                        _cp_label = "  |  CP=0"
+                    else:
                         _cp_label = ""
 
                 pe_rect = patches.Rectangle(
@@ -250,9 +266,12 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
     ax.set_aspect("equal")
     ax.axis("off")
 
+    _title_cp = f", CP={cfg.topo.cp}"
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _title_cp += " (showing 1 of CP groups; others identical)"
     title = (
         f"Per-PE layout of one CP group  "
-        f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
+        f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
         f" | KV mode: {cfg.topo.kv_shard_mode}"
     )
     ax.set_title(title, fontsize=11, fontweight="bold")

tests/analytical_visualization/test_pe_weight_layout.py

@@ -0,0 +1,101 @@
+"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers.
+
+The PE-level view has always color-labeled CP ranks when they're packed
+intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP
+gets pushed to cube-level, so the intra-cube color branch is skipped —
+which used to silently drop the CP rank locator entirely. These tests
+guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are
+always present when CP > 1, regardless of placement.
+"""
+from __future__ import annotations
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tests.analytical_visualization.model_config import (
+    FullConfig, MachineParams, TopologyConfig,
+)
+from tests.analytical_visualization.model_presets import PRESETS
+from tests.analytical_visualization.pe_weight_layout import draw_pe_layout
+
+
+def _cfg(cp: int, tp: int, cp_placement: str | None = None):
+    model = PRESETS["Llama 3 8B"].model
+    kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode")
+    if cp_placement is not None:
+        kwargs["cp_placement"] = cp_placement
+    return FullConfig(
+        model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(),
+    )
+
+
+def _all_text(fig) -> str:
+    """Concatenate all text artists (labels + title) into one blob."""
+    ax = fig.gca()
+    parts = [t.get_text() for t in ax.texts]
+    parts.append(ax.get_title())
+    return "\n".join(parts)
+
+
+def test_high_tp_cp_gt_1_shows_rank_locator():
+    """TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator
+    must still appear in cube tag + per-PE header + figure title."""
+    cfg = _cfg(cp=4, tp=8)
+    assert cfg.topo.cp_placement == "cube"  # sanity
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank 0 of 4" in text, text
+        assert "CP=0" in text, text
+        assert "CP=4" in text, text  # figure title dim summary
+    finally:
+        plt.close(fig)
+
+
+def test_intra_cube_cp_placement_still_shows_per_rank_labels():
+    """cp_placement='pe' (small TP, CP>1) — the old branch already
+    colors + labels each CP rank. Regression check: the new tag does
+    NOT replace those, and each CP=0..CP-1 label is still emitted."""
+    cfg = _cfg(cp=4, tp=2)  # 4*2=8 fits intra-cube
+    assert cfg.topo.cp_placement == "cube"  # default, but auto/manual can differ
+    # Force pe placement to exercise the intra-cube branch.
+    cfg.topo.cp_placement = "pe"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        for r in range(4):
+            assert f"CP={r}" in text, (r, text)
+        # The 'CP rank 0 of N' tag is only for the cube-placement fallback.
+        assert "CP rank 0 of" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_cp_1_no_locator_added():
+    """CP=1: nothing to locate; no CP tag or label should appear."""
+    cfg = _cfg(cp=1, tp=8)
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert "CP rank" not in text, text
+        # The header 'CP=0' should not be added when CP=1.
+        # (Figure title mentions 'CP=1' — that's an OK dim summary,
+        # not a locator per-PE.)
+        assert "  |  CP=0" not in text, text
+    finally:
+        plt.close(fig)
+
+
+def test_tp_16_still_shows_cp_rank_across_spilled_cubes():
+    """TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes
+    for CP rank 0 must carry the '[CP rank 0 of 2]' tag."""
+    cfg = _cfg(cp=2, tp=16)
+    assert cfg.topo.cp_placement == "cube"
+    fig = draw_pe_layout(cfg)
+    try:
+        text = _all_text(fig)
+        assert text.count("CP rank 0 of 2") >= 2, text
+        assert "CP=0" in text, text
+    finally:
+        plt.close(fig)