analytical-viz: color-code CP groups in PE layout when cp_placement=pe

When cp_placement=pe packs multiple CP ranks intra-cube, the PE-level
layout now:

  1. Colors each PE's background by its CP rank (8-color palette wraps
     for cp > 8). Cp_placement=cube keeps the historical light-blue
     styling (only one CP rank per cube tile).

  2. Adds "TP=N  |  CP=M" to the per-PE header so users can read the
     (cp_rank, tp_rank) pair without inferring from position.

  3. Fixes a head-assignment bug: _q_heads_for_pe / _kv_heads_for_pe /
     _per_pe_bytes used to be called with the raw PE-in-group index,
     which treated every PE as a distinct TP rank. Under cp_placement=pe
     this gave wrong Q/KV head lists for PEs beyond the first TP group
     (PE 2..7 with tp=2, cp=4 would ask for head slots 32..127 in a
     32-head model). Now called with tp_rank = pe_id % tp, so all CP
     ranks in the cube share the correct head split.

Verified via a headless matplotlib smoke test with CP=4/TP=2 on Qwen 3
8B: 8 PE patches + 1 cube patch → 5 distinct facecolors (cube + 4 CP
ranks), no errors.

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

tests/analytical_visualization/pe_weight_layout.py

@@ -136,6 +136,26 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
         pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
         pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
 
+        # When cp_placement=pe, CP ranks are packed intra-cube along with
+        # TP ranks. Color-code by CP rank so it's obvious which PEs belong
+        # to which sequence-shard group. Palette wraps beyond 8 CP ranks.
+        _cp_place = cfg.topo.cp_placement
+        _cp_val = cfg.topo.cp
+        _cp_palette = [
+            "#f0f7ff",  # baseline blue (also the fallback for cp_placement=cube)
+            "#e8f5e9",  # light green
+            "#fff3e0",  # light orange
+            "#f3e5f5",  # light purple
+            "#ffebee",  # light red
+            "#e0f7fa",  # light cyan
+            "#fce4ec",  # light pink
+            "#f9fbe7",  # light lime
+        ]
+        _cp_edge_palette = [
+            "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
+            "#c62828", "#0097a7", "#c2185b", "#9e9d24",
+        ]
+
         for pr in range(PE_ROWS):
             for pc in range(PE_COLS):
                 pe_local = pr * PE_COLS + pc
@@ -145,17 +165,36 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
                 py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
 
+                # Which (cp_rank, tp_rank) does this PE hold?
+                # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
+                # When cp_placement=cube: every PE in this cube is one CP rank
+                # (drawn on its own tile), tp_rank = pe_id % tp.
+                if _cp_place == "pe" and _cp_val > 1:
+                    _cp_rank = pe_id_in_group // tp
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = _cp_palette[_cp_rank % len(_cp_palette)]
+                    _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
+                    _cp_label = f"  |  CP={_cp_rank}"
+                else:
+                    _cp_rank = None
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = "#f0f7ff"
+                    _ec = "#3a86ff"
+                    _cp_label = ""
+
                 pe_rect = patches.Rectangle(
                     (px, py), pe_w, pe_h,
-                    facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
+                    facecolor=_fc, edgecolor=_ec, linewidth=1.2,
                 )
                 ax.add_patch(pe_rect)
 
-                q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q)
+                # Head assignment uses TP rank (not raw pe_id) so that under
+                # cp_placement=pe, all CP ranks share the same head split.
+                q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
                 kv_heads, kv_note = _kv_heads_for_pe(
-                    pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
+                    _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
                 )
-                bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
+                bytes_ = _per_pe_bytes(cfg, _tp_rank)
                 weights_gb = sum(v for k, v in bytes_.items()
                                  if k not in ("KV cache", "Transient")) / 1e9
                 kv_gb = bytes_["KV cache"] / 1e9
@@ -172,7 +211,7 @@ def draw_pe_layout(cfg: FullConfig, ax=None):
                 ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
                           + bytes_["W_down"]) / 1e6
                 # Header (PE id and heads) bigger, then per-tensor breakdown
-                header = (f"PE {pe_id_in_group}\n"
+                header = (f"PE {pe_id_in_group}  TP={_tp_rank}{_cp_label}\n"
                           f"Q heads: {q_str}\n"
                           f"KV head: {kv_str}")
                 ax.text(px + 0.08, py + pe_h - 0.05, header,