feat(v8-r03): memory.matrix RPC + MemoryMatrixTab sidebar grid

Closes cppmega-mlx-yz6n.3. New RPC `memory.matrix(spec, topologies?,
precisions?, headroom?)` runs verify_and_estimate once in bf16 and
post-scales weights/grads/optimizer/activations/edge_handoff per
precision (fp32=4, bf16=fp16=2, fp8=1, mxfp4=0.5 bytes/element). KV
cache stays separate since kv_cache_dtype_bytes is independent of the
weight dtype.

Default axes: 4 topologies (h100_8x / m3_ultra_solo / gb10_quarter /
tpu_v6e_8) × 5 precisions = 20 cells, each with fits flag against the
topology's total_hbm_bytes × headroom.

UI: new MemoryMatrixTab in the sidebar (key "memory") renders a
colour-coded grid — green if fits, red otherwise — with title-tooltip
showing the per-component breakdown. Refetches on specPayload change.

Tests: 10/10 new pytest (axes, precision ordering, breakdown sum,
unknown-axis rejection, dispatch e2e), 3/3 new vitest (grid mount,
refetch on payload change, error banner). Regression: 893 pytest + 458
vitest green.

Author: apstenku123

cppmega_v4/jsonrpc/dispatcher.py

@@ -75,6 +75,10 @@
     CacheMetricsParams,
     cache_metrics,
 )
+from cppmega_v4.jsonrpc.memory_matrix_method import (
+    MemoryMatrixParams,
+    memory_matrix,
+)
 from cppmega_v4.jsonrpc.tokenizer_roundtrip_text_method import (
     TokenizerRoundtripTextParams,
     roundtrip_text,
@@ -129,6 +133,10 @@
         ScaleDownParams,
         lambda p, c: scale_down_method(p, cache=c),
     ),
+    "memory.matrix": (
+        MemoryMatrixParams,
+        lambda p, c: memory_matrix(p, cache=c),
+    ),
     "probe.run": (
         ProbeRunParams,
         lambda p, c: probe_run(p, cache=c),

cppmega_v4/jsonrpc/memory_matrix_method.py

@@ -0,0 +1,195 @@
+"""V8-R03: ``memory.matrix`` RPC — compute a (topology × precision) grid
+of memory estimates and per-cell fit verdicts.
+
+Used by the V8 MemoryMatrix sidebar tab to show the user at-a-glance
+which device-precision combinations fit their current canvas spec.
+
+Backend strategy
+----------------
+
+For each cell we run :func:`cppmega_v4.spec.verify_and_estimate` once
+in bf16, then post-scale dtype-sensitive components (weights,
+gradients, activations, optimizer state) by ``precision_bytes / 2``.
+``mxfp4`` is 0.5 bytes per element — supported because the scaling is
+done in float space and rounded only on the final ``total_bytes``.
+
+This is the same approach the V8 spec §3 nominates: one verify per
+matrix, post-scale per precision, so the round-trip is fast enough to
+re-run on every spec edit.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from cppmega_v4.fusion.brick_graph import from_block_specs
+from cppmega_v4.jsonrpc.cache import LRUCache
+from cppmega_v4.jsonrpc.methods import (
+    _cache_lookup, _cache_store, _graph_to_specs,
+)
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.parallelism import topology as _topo
+from cppmega_v4.spec.api import verify_and_estimate
+
+
+__all__ = [
+    "MemoryMatrixParams",
+    "MemoryMatrixCell",
+    "MemoryMatrixResult",
+    "memory_matrix",
+    "PRECISION_BYTES",
+    "TOPOLOGY_BUILDERS",
+]
+
+
+# Precision -> bytes per element. mxfp4 is half a byte (4-bit mantissa
+# with shared e4m3 scale per block — the scale overhead is folded into
+# the headroom and is < 5% at block_size=16).
+PRECISION_BYTES: dict[str, float] = {
+    "fp32":  4.0,
+    "bf16":  2.0,
+    "fp16":  2.0,
+    "fp8":   1.0,
+    "mxfp4": 0.5,
+}
+
+
+# Builder lookup for topologies the matrix supports. Wraps the
+# zero-arg factories from cppmega_v4.parallelism.topology.
+TOPOLOGY_BUILDERS: dict[str, Any] = {
+    "h100_8x":       lambda: _topo.h100_8x(),
+    "h200_8x":       lambda: _topo.h200_8x(),
+    "a100_8x":       lambda: _topo.a100_8x(),
+    "b100_8x":       lambda: _topo.b100_8x(),
+    "gb10_quarter":  lambda: _topo.gb10_quarter(),
+    "tpu_v6e_8":     lambda: _topo.tpu_v6e_8(),
+    "tpu_v5p_4":     lambda: _topo.tpu_v5p_4(),
+    "m3_ultra_solo": lambda: _topo.m3_ultra_solo(),
+}
+
+
+_DEFAULT_TOPOLOGIES = (
+    "h100_8x", "m3_ultra_solo", "gb10_quarter", "tpu_v6e_8",
+)
+_DEFAULT_PRECISIONS = ("fp32", "bf16", "fp16", "fp8", "mxfp4")
+
+
+class MemoryMatrixParams(BaseModel):
+    """Input — a VerifyParams payload + the topologies/precisions axes."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    spec: VerifyParams
+    topologies: list[str] | None = None
+    precisions: list[str] | None = None
+    headroom: float = 0.9
+
+
+class MemoryMatrixCell(BaseModel):
+    """One cell of the matrix — wire-form."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    topology: str
+    precision: str
+    bytes: int
+    device_hbm_bytes: int
+    fits: bool
+    headroom: float
+    breakdown: dict[str, int]
+
+
+class MemoryMatrixResult(BaseModel):
+    """The full matrix — one row per (topology, precision) pair."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    cells: list[MemoryMatrixCell]
+    topologies: list[str]
+    precisions: list[str]
+
+
+def _scale_int(value: int, factor: float) -> int:
+    """Round half-away-from-zero so e.g. ``2 * 0.5 = 1`` not ``0``."""
+    return int(math.floor(value * factor + 0.5))
+
+
+def memory_matrix(
+    params: MemoryMatrixParams, *, cache: LRUCache | None = None,
+) -> MemoryMatrixResult:
+    """Build the (topology × precision) memory matrix for ``params.spec``."""
+    key, hit = _cache_lookup(cache, "memory.matrix", params)
+    if hit is not None:
+        return hit
+
+    topologies = list(params.topologies or _DEFAULT_TOPOLOGIES)
+    precisions = list(params.precisions or _DEFAULT_PRECISIONS)
+    for t in topologies:
+        if t not in TOPOLOGY_BUILDERS:
+            raise ValueError(
+                f"unknown topology {t!r}; choose from "
+                f"{sorted(TOPOLOGY_BUILDERS)}")
+    for p in precisions:
+        if p not in PRECISION_BYTES:
+            raise ValueError(
+                f"unknown precision {p!r}; choose from "
+                f"{sorted(PRECISION_BYTES)}")
+    if not 0.0 < params.headroom <= 1.0:
+        raise ValueError(
+            f"headroom must be in (0, 1], got {params.headroom}")
+
+    # One verify_and_estimate call gives us the full MemoryReport with
+    # aggregate totals; we then post-scale per precision in float space.
+    specs = _graph_to_specs(params.spec.graph)
+    hidden = params.spec.dim_env.get("H", 64)
+    graph = from_block_specs(specs, hidden_size=hidden, instantiate=False)
+    base = verify_and_estimate(
+        graph,
+        dim_env=params.spec.dim_env,
+        training=params.spec.training,
+    ).memory
+
+    base_weights_b   = base.weights_bytes
+    base_grads_b     = base.grads_bytes
+    base_optim_b     = base.optimizer_bytes
+    base_act_b       = base.activations_bytes
+    base_kv_b        = base.kv_cache_bytes
+    base_edge_b      = base.edge_handoff_bytes
+    # KV cache uses kv_cache_dtype_bytes (independent of weight dtype),
+    # so we don't scale it. Edge handoff is bytes-of-tensors-in-transit
+    # and IS dtype-sensitive.
+
+    cells: list[MemoryMatrixCell] = []
+    for t_name in topologies:
+        topo = TOPOLOGY_BUILDERS[t_name]()
+        hbm = topo.total_hbm_bytes
+        for p_name in precisions:
+            scale = PRECISION_BYTES[p_name] / 2.0  # bf16 baseline
+            w  = _scale_int(base_weights_b, scale)
+            g  = _scale_int(base_grads_b, scale)
+            o  = _scale_int(base_optim_b, scale)
+            a  = _scale_int(base_act_b, scale)
+            e  = _scale_int(base_edge_b, scale)
+            total = w + g + o + a + base_kv_b + e
+            cells.append(MemoryMatrixCell(
+                topology=t_name, precision=p_name,
+                bytes=total, device_hbm_bytes=hbm,
+                fits=total <= int(hbm * params.headroom),
+                headroom=params.headroom,
+                breakdown={
+                    "weights":     w,
+                    "grads":       g,
+                    "optimizer":   o,
+                    "activations": a,
+                    "kv_cache":    base_kv_b,
+                    "edge_handoff": e,
+                },
+            ))
+
+    out = MemoryMatrixResult(
+        cells=cells, topologies=topologies, precisions=precisions)
+    _cache_store(cache, key, out)
+    return out

cppmega_v4/jsonrpc/schema.py

@@ -808,6 +808,7 @@ class PipelineRunResult(BaseModel):
     "side_channels.preview",
     "platform.get_info",
     "architectures.scale_down",
+    "memory.matrix",
 })
 
 

tests/v4/test_memory_matrix.py

@@ -0,0 +1,146 @@
+"""V8-R03 unit tests: memory.matrix RPC.
+
+Asserts:
+  * default 4×5 matrix when topologies/precisions omitted
+  * per-precision bytes scale linearly with precision_bytes/2 (bf16 baseline)
+  * mxfp4 < fp8 < bf16 == fp16 < fp32 for the same topology
+  * fits flag flips when bytes exceed device_hbm_bytes × headroom
+  * unknown topology / precision rejected
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_v4.jsonrpc.dispatcher import dispatch
+from cppmega_v4.jsonrpc.memory_matrix_method import (
+    MemoryMatrixParams, memory_matrix, PRECISION_BYTES,
+    TOPOLOGY_BUILDERS,
+)
+from cppmega_v4.jsonrpc.schema import JsonRpcRequest, VerifyParams
+
+
+def _mini_spec() -> dict:
+    """Minimal canvas spec: attention + mlp, hidden=256."""
+    return {
+        "graph": {
+            "nodes": [
+                {"id": "a", "kind": "attention", "params": {}},
+                {"id": "b", "kind": "mlp", "params": {}},
+            ],
+            "edges": [{"src": "a", "dst": "b"}],
+        },
+        "dim_env": {"H": 256, "B": 1, "S": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["b"]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 1e-3, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]},
+        ]},
+        "sharding": None,
+        "training": True,
+    }
+
+
+def test_default_axes_yield_4x5_grid():
+    res = memory_matrix(MemoryMatrixParams(
+        spec=VerifyParams(**_mini_spec()),
+    ))
+    assert len(res.cells) == 4 * 5
+    assert set(res.topologies) == {
+        "h100_8x", "m3_ultra_solo", "gb10_quarter", "tpu_v6e_8"}
+    assert set(res.precisions) == {"fp32", "bf16", "fp16", "fp8", "mxfp4"}
+
+
+def test_precision_ordering_within_a_topology():
+    res = memory_matrix(MemoryMatrixParams(
+        spec=VerifyParams(**_mini_spec()),
+        topologies=["h100_8x"],
+    ))
+    by_p = {c.precision: c.bytes for c in res.cells}
+    assert by_p["mxfp4"] < by_p["fp8"] < by_p["bf16"]
+    assert by_p["bf16"] == by_p["fp16"]
+    assert by_p["bf16"] < by_p["fp32"]
+
+
+def test_bf16_baseline_unchanged():
+    """The bf16 column must equal the raw verify_and_estimate total —
+    no rounding loss when the scale is exactly 1.0."""
+    res = memory_matrix(MemoryMatrixParams(
+        spec=VerifyParams(**_mini_spec()),
+        topologies=["h100_8x"], precisions=["bf16"],
+    ))
+    cell = res.cells[0]
+    breakdown = cell.breakdown
+    # Sum of the parts is the cell total.
+    assert (breakdown["weights"] + breakdown["grads"] + breakdown["optimizer"]
+            + breakdown["activations"] + breakdown["kv_cache"]
+            + breakdown["edge_handoff"]) == cell.bytes
+
+
+def test_fits_flag_reflects_device_hbm():
+    """All cells with bytes < gb10's HBM × headroom must fit."""
+    res = memory_matrix(MemoryMatrixParams(
+        spec=VerifyParams(**_mini_spec()),
+        topologies=["gb10_quarter"], precisions=["bf16"],
+        headroom=0.9,
+    ))
+    cell = res.cells[0]
+    assert cell.fits == (cell.bytes <= int(cell.device_hbm_bytes * 0.9))
+
+
+def test_unknown_topology_rejected():
+    with pytest.raises(ValueError, match="unknown topology"):
+        memory_matrix(MemoryMatrixParams(
+            spec=VerifyParams(**_mini_spec()),
+            topologies=["nonexistent_zzz"],
+        ))
+
+
+def test_unknown_precision_rejected():
+    with pytest.raises(ValueError, match="unknown precision"):
+        memory_matrix(MemoryMatrixParams(
+            spec=VerifyParams(**_mini_spec()),
+            precisions=["int1024"],
+        ))
+
+
+def test_invalid_headroom_rejected():
+    with pytest.raises(ValueError, match="headroom"):
+        memory_matrix(MemoryMatrixParams(
+            spec=VerifyParams(**_mini_spec()),
+            headroom=1.5,
+        ))
+
+
+def test_dispatch_end_to_end():
+    req = JsonRpcRequest(
+        jsonrpc="2.0", id="t-r03",
+        method="memory.matrix",
+        params={
+            "spec": _mini_spec(),
+            "topologies": ["m3_ultra_solo", "h100_8x"],
+            "precisions": ["bf16", "mxfp4"],
+        },
+    )
+    resp = dispatch(req)
+    assert resp.error is None, resp.error
+    r = resp.result
+    assert len(r["cells"]) == 4
+    by = {(c["topology"], c["precision"]): c["bytes"] for c in r["cells"]}
+    assert by[("h100_8x", "mxfp4")] < by[("h100_8x", "bf16")]
+
+
+def test_precision_bytes_table_matches_spec():
+    """V8 spec §3 nominates the canonical precision -> bytes map."""
+    assert PRECISION_BYTES["fp32"] == 4.0
+    assert PRECISION_BYTES["bf16"] == 2.0
+    assert PRECISION_BYTES["fp16"] == 2.0
+    assert PRECISION_BYTES["fp8"]  == 1.0
+    assert PRECISION_BYTES["mxfp4"] == 0.5
+
+
+def test_all_default_topologies_buildable():
+    """Each named topology builder produces a usable DeviceTopology."""
+    for name, factory in TOPOLOGY_BUILDERS.items():
+        topo = factory()
+        assert topo.total_hbm_bytes > 0, name