diff --git a/docs/diagrams/cube_view.svg b/docs/diagrams/cube_view.svg
index a51b065..54b0cbf 100644
--- a/docs/diagrams/cube_view.svg
+++ b/docs/diagrams/cube_view.svg
@@ -133,45 +133,24 @@
r0c0
-
- PE0
-
-
- UCIe-W.c0
-
-
- UCIe-N.c0
-
+
+ PE0
+
r0c1
-
- PE1
-
-
- UCIe-N.c1
-
+
+ PE1
+
r0c2
r0c3
r0c4
-
- UCIe-N.c2
-
r0c5
-
- UCIe-E.c0
-
-
- UCIe-N.c3
-
r1c0
-
- UCIe-W.c1
-
r1c1
@@ -180,21 +159,18 @@
r1c3
r1c4
-
- PE2
+
+ PE2
r1c5
-
- PE3
-
-
- UCIe-E.c1
-
+
+ PE3
+
r2c0
-
- M_CPU
+
+ M_CPU
r2c1
@@ -204,8 +180,8 @@
r2c5
r3c0
-
- SRAM
+
+ SRAM
r3c1
@@ -215,16 +191,13 @@
r3c5
r4c0
-
- PE4
-
-
- UCIe-W.c2
-
+
+ PE4
+
r4c1
-
- PE5
+
+ PE5
r4c2
@@ -234,45 +207,24 @@
r4c4
r4c5
-
- UCIe-E.c2
-
r5c0
-
- UCIe-W.c3
-
-
- UCIe-S.c0
-
r5c1
-
- UCIe-S.c1
-
r5c2
r5c3
r5c4
-
- PE6
-
-
- UCIe-S.c2
-
+
+ PE6
+
r5c5
-
- PE7
-
-
- UCIe-E.c3
-
-
- UCIe-S.c3
-
+
+ PE7
+
256GB/s
@@ -289,6 +241,62 @@
256GB/s
256GB/s
+
+ UCIe-W
+
+ c0
+
+
+ c1
+
+
+ c2
+
+
+ c3
+
+
+ UCIe-N
+
+ c0
+
+
+ c1
+
+
+ c2
+
+
+ c3
+
+
+ UCIe-E
+
+ c0
+
+
+ c1
+
+
+ c2
+
+
+ c3
+
+
+ UCIe-S
+
+ c0
+
+
+ c1
+
+
+ c2
+
+
+ c3
+
PE Router
diff --git a/src/kernbench/topology/visualizer.py b/src/kernbench/topology/visualizer.py
index e1df01d..e19de17 100644
--- a/src/kernbench/topology/visualizer.py
+++ b/src/kernbench/topology/visualizer.py
@@ -599,11 +599,7 @@ def _render_cube_view_svg(view: ViewGraph, spec: dict) -> str:
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"]))
+ # UCIe handled separately below
# Position blocks outward from router (away from cube center)
for bi, (label, kind, style) in enumerate(blocks):
@@ -611,29 +607,7 @@ def _render_cube_view_svg(view: ViewGraph, spec: dict) -> str:
# 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 flush against cube edge at router position
- direction = label.split("-")[1].split(".")[0] if "-" in label else ""
- ucie_w, ucie_h = 22, 10 # smaller blocks for UCIe ports
- if direction == "N":
- bx = px - ucie_w / 2
- by = pad - ucie_h # flush against top edge
- blk_w, blk_h = ucie_w, ucie_h
- elif direction == "S":
- bx = px - ucie_w / 2
- by = pad + cube_h * scale # flush against bottom edge
- blk_w, blk_h = ucie_w, ucie_h
- elif direction == "W":
- bx = pad - ucie_w # flush against left edge
- by = py - ucie_h / 2
- blk_w, blk_h = ucie_w, ucie_h
- elif direction == "E":
- bx = pad + cube_w * scale # flush against right edge
- by = py - ucie_h / 2
- blk_w, blk_h = ucie_w, ucie_h
- else:
- bx, by = px - blk_w / 2, py - r_size - blk_h - 4
- elif kind in ("mcpu", "sram"):
+ if kind in ("mcpu", "sram"):
# M_CPU/SRAM: place to the left of router (avoid mesh overlap)
bx = px - r_size - blk_w - 6
by = py - blk_h / 2 + bi * (blk_h + 2)
@@ -659,33 +633,7 @@ def _render_cube_view_svg(view: ViewGraph, spec: dict) -> str:
f'font-weight="bold" fill="{style["text"]}">{_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' '
- )
- elif kind in ("mcpu", "sram"):
+ if kind in ("mcpu", "sram"):
# Horizontal connector (block right edge → router left edge)
parts.append(
f' str:
f'{agg_bw:.0f}GB/s'
)
+ # ── UCIe port components (inside cube boundary, on edges) ──
+ # Collect UCIe connections per direction
+ ucie_by_dir: dict[str, list[tuple[str, str, float, float]]] = {} # dir → [(conn, rkey, rx, ry)]
+ for rkey, rval in routers.items():
+ if rval is None:
+ continue
+ rx, ry = rval["pos_mm"]
+ for a in rval.get("attach", []):
+ if not a.startswith("ucie_"):
+ continue
+ parts_a = a.split(".")
+ direction = parts_a[0].replace("ucie_", "").upper()
+ conn = parts_a[1] if len(parts_a) > 1 else "c0"
+ ucie_by_dir.setdefault(direction, []).append((conn, rkey, rx, ry))
+
+ ucie_colors = ["#818cf8", "#a78bfa", "#c084fc", "#e879f9"]
+ ucie_port_inset = 3 # mm inset from cube edge
+
+ for direction, conns in ucie_by_dir.items():
+ conns.sort(key=lambda x: x[0]) # sort by conn name
+ n_conn = len(conns)
+
+ # UCIe component box position (inside cube, along edge)
+ if direction == "N":
+ # Horizontal bar along top edge
+ # Find X span from attached routers
+ xs = [mm2px(rx, ry)[0] for _, _, rx, ry in conns]
+ ux1 = min(xs) - 15
+ ux2 = max(xs) + 15
+ uw = ux2 - ux1
+ uh = 18
+ ux = ux1
+ uy = pad + ucie_port_inset * scale
+ elif direction == "S":
+ xs = [mm2px(rx, ry)[0] for _, _, rx, ry in conns]
+ ux1 = min(xs) - 15
+ ux2 = max(xs) + 15
+ uw = ux2 - ux1
+ uh = 18
+ ux = ux1
+ uy = pad + cube_h * scale - ucie_port_inset * scale - uh
+ elif direction == "W":
+ ys = [mm2px(rx, ry)[1] for _, _, rx, ry in conns]
+ uy1 = min(ys) - 15
+ uy2 = max(ys) + 15
+ uw = 18
+ uh = uy2 - uy1
+ ux = pad + ucie_port_inset * scale
+ uy = uy1
+ elif direction == "E":
+ ys = [mm2px(rx, ry)[1] for _, _, rx, ry in conns]
+ uy1 = min(ys) - 15
+ uy2 = max(ys) + 15
+ uw = 18
+ uh = uy2 - uy1
+ ux = pad + cube_w * scale - ucie_port_inset * scale - uw
+ uy = uy1
+ else:
+ continue
+
+ # UCIe component background
+ parts.append(
+ f' '
+ )
+ # UCIe label
+ if direction in ("N", "S"):
+ parts.append(
+ f' '
+ f'UCIe-{direction}'
+ )
+ else:
+ parts.append(
+ f' '
+ f'UCIe-{direction}'
+ )
+
+ # Connection port boxes inside UCIe component
+ for ci, (conn, rkey, crx, cry) in enumerate(conns):
+ c_color = ucie_colors[ci % len(ucie_colors)]
+ if direction in ("N", "S"):
+ # Horizontal layout
+ cw = max((uw - 4) / n_conn - 1, 6)
+ ch = uh - 4
+ cx = ux + 2 + ci * (cw + 1)
+ cy = uy + 2
+ else:
+ # Vertical layout
+ cw = uw - 4
+ ch = max((uh - 4) / n_conn - 1, 6)
+ cx = ux + 2
+ cy = uy + 2 + ci * (ch + 1)
+
+ parts.append(
+ f' '
+ )
+ # Connection label
+ lx = cx + cw / 2
+ ly_t = cy + ch / 2 + 3
+ parts.append(
+ f' '
+ f'{conn}'
+ )
+
+ # Connector line: port box → attached router
+ rpx, rpy = mm2px(crx, cry)
+ if direction == "N":
+ parts.append(
+ f' '
+ )
+ elif direction == "S":
+ parts.append(
+ f' '
+ )
+ elif direction == "W":
+ parts.append(
+ f' '
+ )
+ elif direction == "E":
+ parts.append(
+ f' '
+ )
+
# ── Legend ──
ly = h_px - 35
legend_items = [