diff --git a/docs/diagrams/cube_view.svg b/docs/diagrams/cube_view.svg
index abc3519..befc33c 100644
--- a/docs/diagrams/cube_view.svg
+++ b/docs/diagrams/cube_view.svg
@@ -124,126 +124,188 @@
-
- r0c0
- PE0, UCIe×2
-
- 256GB/s
-
- r0c1
- PE1, UCIe×1
-
- 256GB/s
-
- r0c2
-
- r0c3
-
- r0c4
- UCIe×1
-
- r0c5
- UCIe×2
-
- r1c0
- UCIe×1
-
- r1c1
-
- r1c2
-
- r1c3
-
- r1c4
- PE2
-
- 256GB/s
-
- r1c5
- PE3, UCIe×1
-
- 256GB/s
-
- r2c0
- M_CPU
-
- r2c1
-
- r2c4
-
- r2c5
-
- r3c0
- SRAM
-
- r3c1
-
- r3c4
-
- r3c5
-
- r4c0
- PE4, UCIe×1
-
- 256GB/s
-
- r4c1
- PE5
-
- 256GB/s
-
- r4c2
-
- r4c3
-
- r4c4
-
- r4c5
- UCIe×1
-
- r5c0
- UCIe×2
-
- r5c1
- UCIe×1
-
- r5c2
-
- r5c3
-
- r5c4
- PE6, UCIe×1
-
- 256GB/s
-
- r5c5
- PE7, UCIe×2
-
- 256GB/s
-
- PE0
-
-
- PE1
-
-
- PE2
-
-
- PE3
-
-
- PE4
-
-
- PE5
-
-
- PE6
-
-
- PE7
-
+
+ r0c0
+
+
+ PE0
+
+
+ UCIe-W.c0
+
+
+ UCIe-N.c0
+
+ 256GB/s
+
+ r0c1
+
+
+ PE1
+
+
+ UCIe-N.c1
+
+ 256GB/s
+
+ r0c2
+
+
+ r0c3
+
+
+ r0c4
+
+
+ UCIe-N.c2
+
+
+ r0c5
+
+
+ UCIe-E.c0
+
+
+ UCIe-N.c3
+
+
+ r1c0
+
+
+ UCIe-W.c1
+
+
+ r1c1
+
+
+ r1c2
+
+
+ r1c3
+
+
+ r1c4
+
+
+ PE2
+
+ 256GB/s
+
+ r1c5
+
+
+ PE3
+
+
+ UCIe-E.c1
+
+ 256GB/s
+
+ r2c0
+
+
+ M_CPU
+
+
+ r2c1
+
+
+ r2c4
+
+
+ r2c5
+
+
+ r3c0
+
+
+ SRAM
+
+
+ r3c1
+
+
+ r3c4
+
+
+ r3c5
+
+
+ r4c0
+
+
+ PE4
+
+
+ UCIe-W.c2
+
+ 256GB/s
+
+ r4c1
+
+
+ PE5
+
+ 256GB/s
+
+ r4c2
+
+
+ r4c3
+
+
+ r4c4
+
+
+ r4c5
+
+
+ UCIe-E.c2
+
+
+ r5c0
+
+
+ UCIe-W.c3
+
+
+ UCIe-S.c0
+
+
+ r5c1
+
+
+ UCIe-S.c1
+
+
+ r5c2
+
+
+ r5c3
+
+
+ r5c4
+
+
+ PE6
+
+
+ UCIe-S.c2
+
+ 256GB/s
+
+ r5c5
+
+
+ PE7
+
+
+ UCIe-E.c3
+
+
+ UCIe-S.c3
+
+ 256GB/s
PE Router
diff --git a/src/kernbench/topology/visualizer.py b/src/kernbench/topology/visualizer.py
index 60cdf8b..e3f4105 100644
--- a/src/kernbench/topology/visualizer.py
+++ b/src/kernbench/topology/visualizer.py
@@ -529,120 +529,157 @@ def _render_cube_view_svg(view: ViewGraph, spec: dict) -> str:
)
break
- # ── Router nodes ──
- r_size = 10 # px radius
- pe_routers: dict[str, str] = {} # rkey → pe label
+ # ── Router nodes + attached component blocks ──
+ r_size = 8 # px radius for router circle
+ blk_w, blk_h = 32, 16 # px for component blocks
+
+ # Component style definitions
+ _COMP_STYLE = {
+ "pe": {"fill": "#2d1f3d", "stroke": "#a855f7", "text": "#a855f7"},
+ "mcpu": {"fill": "#451a03", "stroke": "#f59e0b", "text": "#f59e0b"},
+ "sram": {"fill": "#1c1917", "stroke": "#d97706", "text": "#d97706"},
+ "ucie": {"fill": "#1e1b4b", "stroke": "#8b5cf6", "text": "#8b5cf6"},
+ }
+
for rkey, rval in routers.items():
if rval is None:
continue
rx, ry = rval["pos_mm"]
px, py = mm2px(rx, ry)
attach = rval.get("attach", [])
+ is_top = ry < cube_h / 2
- # Determine router type by attachments
- has_pe = any(a.endswith(".dma") for a in attach)
- has_mcpu = "m_cpu" in attach
- has_sram = "sram" in attach
- has_ucie = any(a.startswith("ucie_") for a in attach)
-
- if has_pe:
- fill, stroke = "#3b82f6", "#1d4ed8"
- pe_name = [a for a in attach if a.endswith(".dma")][0].split(".")[0]
- pe_routers[rkey] = pe_name.upper()
- elif has_mcpu:
- fill, stroke = "#f59e0b", "#d97706"
- elif has_sram:
- fill, stroke = "#f59e0b", "#d97706"
- elif has_ucie:
- fill, stroke = "#8b5cf6", "#6d28d9"
- else:
- fill, stroke = "#334155", "#475569"
-
+ # ── Router circle ──
+ has_attach = len(attach) > 0
+ r_fill = "#475569" if has_attach else "#334155"
+ r_stroke = "#64748b" if has_attach else "#475569"
parts.append(
f' '
+ f'fill="{r_fill}" stroke="{r_stroke}" stroke-width="1"/>'
)
- # Router label
parts.append(
f' '
+ f'font-family="monospace" font-size="6" fill="white">'
f'{rkey}'
)
- # Attachment labels below router
- label_parts = []
- if has_pe:
- label_parts.append(pe_routers[rkey])
- if has_mcpu:
- label_parts.append("M_CPU")
- if has_sram:
- label_parts.append("SRAM")
- ucie_items = [a for a in attach if a.startswith("ucie_")]
- if ucie_items:
- label_parts.append(f"UCIe×{len(ucie_items)}")
- if label_parts:
+ # ── Router → HBM_CTRL line (all routers connect to HBM) ──
+ hbm_edge_y = hbm_y if py < hbm_y else hbm_y + hbm_h
+ r_edge_y = py + r_size if py < hbm_y else py - r_size
+ # Only draw if not inside HBM zone
+ if abs(r_edge_y - hbm_edge_y) > 10:
parts.append(
- f' '
- f'{", ".join(label_parts)}'
+ f' '
)
- # ── PE → HBM connection line (for PE routers) ──
- if has_pe:
- # Draw line from router to HBM zone edge
- target_y = hbm_y if py < hbm_y else hbm_y + hbm_h
+ # ── Attached component blocks ──
+ # Collect components to draw, positioned outward from router
+ blocks: list[tuple[str, str, dict]] = [] # (label, kind, style)
+ pe_items = [a for a in attach if a.endswith(".dma")]
+ if pe_items:
+ pe_name = pe_items[0].split(".")[0].upper()
+ blocks.append((pe_name, "pe", _COMP_STYLE["pe"]))
+ if "m_cpu" in attach:
+ blocks.append(("M_CPU", "mcpu", _COMP_STYLE["mcpu"]))
+ if "sram" in attach:
+ blocks.append(("SRAM", "sram", _COMP_STYLE["sram"]))
+ ucie_items = [a for a in attach if a.startswith("ucie_")]
+ for ui in ucie_items:
+ direction = ui.split(".")[0].replace("ucie_", "").upper()
+ conn = ui.split(".")[1] if "." in ui else ""
+ blocks.append((f"UCIe-{direction}.{conn}", "ucie", _COMP_STYLE["ucie"]))
+
+ # Position blocks outward from router (away from cube center)
+ for bi, (label, kind, style) in enumerate(blocks):
+ # Determine placement direction: PE/components go outward
+ # Use left/right offset for multiple blocks on same router
+ offset_x = (bi - (len(blocks) - 1) / 2) * (blk_w + 4)
+
+ if kind == "ucie":
+ # UCIe: place at cube edge direction
+ direction = label.split("-")[1].split(".")[0] if "-" in label else ""
+ if direction == "N":
+ bx, by = px + offset_x - blk_w / 2, pad - blk_h - 4
+ elif direction == "S":
+ by_base = pad + cube_h * scale
+ bx, by = px + offset_x - blk_w / 2, by_base + 4
+ elif direction == "W":
+ bx, by = pad - blk_w - 4, py + offset_x - blk_h / 2
+ elif direction == "E":
+ bx_base = pad + cube_w * scale
+ bx, by = bx_base + 4, py + offset_x - blk_h / 2
+ else:
+ bx, by = px + offset_x - blk_w / 2, py - r_size - blk_h - 4
+ else:
+ # PE/M_CPU/SRAM: place above (top half) or below (bottom half)
+ bx = px + offset_x - blk_w / 2
+ if is_top:
+ by = py - r_size - blk_h - 4 - bi * (blk_h + 2)
+ else:
+ by = py + r_size + 4 + bi * (blk_h + 2)
+
+ # Block rect
parts.append(
- f' '
+ f' '
)
- # BW annotation
- bw_y = (py + r_size + target_y) / 2
+ # Label
+ font_sz = 6 if len(label) > 6 else 7
parts.append(
- f' {_escape(label)}'
+ )
+ # Connector line: block → router
+ if kind == "ucie":
+ # Line from block edge toward router
+ if direction == "N":
+ parts.append(
+ f' '
+ )
+ elif direction == "S":
+ parts.append(
+ f' '
+ )
+ elif direction == "W":
+ parts.append(
+ f' '
+ )
+ elif direction == "E":
+ parts.append(
+ f' '
+ )
+ else:
+ # Vertical connector
+ ly1 = by + blk_h if is_top else by
+ ly2 = py - r_size if is_top else py + r_size
+ parts.append(
+ f' '
+ )
+
+ # ── PE router → HBM BW annotation ──
+ if pe_items:
+ bw_x = px + 14
+ bw_y = (r_edge_y + hbm_edge_y) / 2
+ parts.append(
+ f' '
f'{agg_bw:.0f}GB/s'
)
- # ── PE blocks (drawn next to their router) ──
- pe_w, pe_h = 30, 18 # px
- for rkey, rval in routers.items():
- if rval is None:
- continue
- attach = rval.get("attach", [])
- pe_dma = [a for a in attach if a.endswith(".dma")]
- if not pe_dma:
- continue
- pe_name = pe_dma[0].split(".")[0] # "pe0"
- pe_label = pe_name.upper()
- rx, ry = rval["pos_mm"]
- px, py = mm2px(rx, ry)
-
- # Position PE block: top-half PEs above router, bottom-half below
- is_top = ry < cube_h / 2
- pe_x = px - pe_w / 2
- pe_y = py - r_size - pe_h - 4 if is_top else py + r_size + 4
-
- parts.append(
- f' '
- )
- parts.append(
- f' '
- f'{pe_label}'
- )
- # PE ↔ router link
- link_y1 = pe_y + pe_h if is_top else pe_y
- link_y2 = py - r_size if is_top else py + r_size
- parts.append(
- f' '
- )
-
# ── Legend ──
ly = h_px - 35
legend_items = [