ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

analytical-viz: topology map — color PEs by CP rank when cp_placement=pe

Browse Source 
The topology map previously colored a whole cube (border + all PE
fills) by the single cp_rank assigned to that cube in the cube→(pp,cp)
mapping. Under cp_placement=cube this is right (each cube = one CP
rank). Under cp_placement=pe, however, multiple CP ranks are PACKED
into the same cube's PEs, so the whole-cube coloring makes every PE
look identical (defaulting to cp_rank=0's palette entry — bright red).

Fix: in the per-PE loop, if cp_packed = (cp_placement=="pe" and cp>1),
compute each PE's own cp_rank = pe_id // tp and look up its color via
_cp_color(pp_stage, pe_cp_rank, cp_size). Border still uses the
cube-level color (cp_rank=0), so the outer bounding box is unchanged,
but the interior PEs now show all four (or however many) CP-rank
colors visibly.

For cp_placement=cube: unchanged (single pe_fill per cube).

Verified with a headless render at Qwen 3 8B, CP=4, TP=2,
cp_placement=pe: 145 patches → 8 distinct facecolors (4 for the CP
ranks in the packed cube + inactive/border tints), where before it
was fewer distinct colors.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-28 22:07:50 -07:00
parent da1909fdf2
commit 6f61657ba4
1 changed files with 16 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/topology_map.py
+16 -1
View File
@@ -258,12 +258,27 @@ def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
pes_used = cfg.topo.pes_per_cube_used if is_used else 0 pes_used = cfg.topo.pes_per_cube_used if is_used else 0
# When cp_placement=pe, multiple CP ranks live inside this cube's PEs
# (cp_rank = pe_id // tp). Color each PE by its cp_rank so all four
# groups are visible; otherwise fall back to the whole-cube pe_fill.
_cp_packed = is_used and cfg.topo.cp_placement == "pe" and cfg.topo.cp > 1
for pr in range(PE_ROWS): for pr in range(PE_ROWS):
for pc in range(PE_COLS): for pc in range(PE_COLS):
pe_id = pr * PE_COLS + pc pe_id = pr * PE_COLS + pc
px = x + pe_pad + pc * (pe_w + pe_gap) px = x + pe_pad + pc * (pe_w + pe_gap)
py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap) py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
fill = pe_fill if pe_id < pes_used else _INACTIVE_PE if pe_id >= pes_used:
fill = _INACTIVE_PE
elif _cp_packed:
_pe_cp_rank = pe_id // cfg.topo.tp
_pe_color, _pe_pe_fill = _cp_color(
info[0] if is_used else 0,
_pe_cp_rank,
cfg.topo.cp,
)
fill = _pe_pe_fill
else:
fill = pe_fill
pe_rect = patches.Rectangle( pe_rect = patches.Rectangle(
(px, py), pe_w, pe_h, (px, py), pe_w, pe_h,
facecolor=fill, edgecolor="#666", linewidth=0.3, facecolor=fill, edgecolor="#666", linewidth=0.3,