mrm-trace 0.1.0__py3-none-any.whl

mrm_trace/__init__.py

@@ -0,0 +1,7 @@
+"""mrm-trace: LLM inference memory trace platform for MRM research."""
+
+__version__ = "0.1.0"
+
+from mrm_trace.api import Experiment
+
+__all__ = ["Experiment"]

mrm_trace/analyser/__init__.py

@@ -0,0 +1,53 @@
+"""Analysis subsystem for mrm-trace — Phase 6."""
+
+from mrm_trace.analyser.iai import compute_iai
+from mrm_trace.analyser.locality import compute_locality
+from mrm_trace.analyser.read_freq import compute_read_freq
+from mrm_trace.analyser.retention import compute_retention
+from mrm_trace.analyser.suitability import classify_suitability
+from mrm_trace.analyser.working_set import compute_working_set
+from mrm_trace.analyser.write_once import compute_write_once
+
+__all__ = [
+    "compute_retention",
+    "compute_write_once",
+    "compute_read_freq",
+    "compute_working_set",
+    "compute_locality",
+    "compute_iai",
+    "classify_suitability",
+    "compute_all",
+]
+
+
+def compute_all(trace) -> dict:
+    """
+    Run all analysis modules on a trace DataFrame.
+
+    Returns a dict with keys:
+      retention_per_region, retention_summary,
+      write_once, read_freq,
+      working_set_per_region, working_set_summary,
+      locality_per_region, locality_summary,
+      iai, suitability
+    """
+    ret_per_region, ret_summary = compute_retention(trace)
+    wo = compute_write_once(trace)
+    rf = compute_read_freq(trace)
+    ws_per_region, ws_summary = compute_working_set(trace)
+    loc_per_region, loc_summary = compute_locality(trace)
+    iai = compute_iai(trace)
+    suit = classify_suitability(ret_summary, wo, rf)
+
+    return {
+        "retention_per_region": ret_per_region,
+        "retention_summary": ret_summary,
+        "write_once": wo,
+        "read_freq": rf,
+        "working_set_per_region": ws_per_region,
+        "working_set_summary": ws_summary,
+        "locality_per_region": loc_per_region,
+        "locality_summary": loc_summary,
+        "iai": iai,
+        "suitability": suit,
+    }

mrm_trace/analyser/iai.py

@@ -0,0 +1,67 @@
+"""
+Inter-access interval (IAI) analysis — time between consecutive accesses.
+
+Groups by exact address (not address_page) to correctly handle sub-page KV
+blocks at 256-byte stride where multiple blocks share a 4K page.
+"""
+
+import numpy as np
+import pandas as pd
+
+
+def compute_iai(trace: pd.DataFrame) -> pd.DataFrame:
+    """
+    Compute inter-access interval distribution per region_type.
+
+    IAI = time gap between consecutive accesses to the *same exact address*,
+    sorted by timestamp_ns within each address group.
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_type, n_intervals, iai_p50_ns, iai_p90_ns,
+                 iai_p99_ns, iai_mean_ns
+        One row per region_type (regions with < 2 accesses to any address
+        contribute 0 intervals and are excluded from the summary if they
+        have no intervals at all).
+    """
+    trace_sorted = trace.sort_values(["address", "timestamp_ns"])
+    all_intervals = []
+
+    for (address, region_type), grp in trace_sorted.groupby(
+        ["address", "region_type"], sort=False
+    ):
+        ts = grp["timestamp_ns"].to_numpy(dtype=np.int64)
+        if len(ts) < 2:
+            continue
+        intervals = np.diff(ts)
+        intervals = intervals[intervals > 0]  # ignore same-timestamp duplicates
+        if len(intervals) == 0:
+            continue
+        all_intervals.append(
+            pd.DataFrame({"region_type": region_type, "iai_ns": intervals})
+        )
+
+    if not all_intervals:
+        return pd.DataFrame(
+            columns=["region_type", "n_intervals", "iai_p50_ns",
+                     "iai_p90_ns", "iai_p99_ns", "iai_mean_ns"]
+        )
+
+    combined = pd.concat(all_intervals, ignore_index=True)
+
+    rows = []
+    for rtype, grp in combined.groupby("region_type"):
+        s = grp["iai_ns"]
+        rows.append(
+            {
+                "region_type": rtype,
+                "n_intervals": len(s),
+                "iai_p50_ns": int(s.quantile(0.50)),
+                "iai_p90_ns": int(s.quantile(0.90)),
+                "iai_p99_ns": int(s.quantile(0.99)),
+                "iai_mean_ns": float(s.mean()),
+            }
+        )
+
+    return pd.DataFrame(rows)

mrm_trace/analyser/locality.py

@@ -0,0 +1,83 @@
+"""
+Access locality analysis — stride distribution and same-page access fraction.
+
+Locality measures spatial regularity: small strides and high same-page fraction
+indicate cache-friendly, predictable access patterns.
+"""
+
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+
+
+def compute_locality(trace: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Compute locality statistics per region.
+
+    Stride = |address[i+1] - address[i]| within consecutive accesses to same region.
+    Same-page fraction = fraction of consecutive accesses that stay on the same page.
+
+    Returns
+    -------
+    per_region : DataFrame
+        Columns: region_id, region_type, n_accesses, mean_stride_bytes,
+                 median_stride_bytes, same_page_fraction
+    type_summary : DataFrame
+        Columns: region_type, mean_stride_bytes, median_stride_bytes,
+                 same_page_fraction
+    """
+    rows = []
+    trace_sorted = trace.sort_values(["region_id", "timestamp_ns"])
+
+    for region_id, grp in trace_sorted.groupby("region_id"):
+        region_type = grp["region_type"].iloc[0]
+        n = len(grp)
+
+        if n < 2:
+            rows.append(
+                {
+                    "region_id": region_id,
+                    "region_type": region_type,
+                    "n_accesses": n,
+                    "mean_stride_bytes": 0.0,
+                    "median_stride_bytes": 0.0,
+                    "same_page_fraction": 0.0,
+                }
+            )
+            continue
+
+        addrs = grp["address"].to_numpy()
+        pages = grp["address_page"].to_numpy()
+        strides = np.abs(np.diff(addrs.astype(np.int64)))
+        same_page = np.sum(pages[1:] == pages[:-1])
+
+        rows.append(
+            {
+                "region_id": region_id,
+                "region_type": region_type,
+                "n_accesses": n,
+                "mean_stride_bytes": float(strides.mean()),
+                "median_stride_bytes": float(np.median(strides)),
+                "same_page_fraction": float(same_page / len(strides)),
+            }
+        )
+
+    per_region = pd.DataFrame(rows)
+    if per_region.empty:
+        return per_region, pd.DataFrame(
+            columns=["region_type", "mean_stride_bytes", "median_stride_bytes",
+                     "same_page_fraction"]
+        )
+
+    type_summary = (
+        per_region.groupby("region_type")
+        .agg(
+            mean_stride_bytes=("mean_stride_bytes", "mean"),
+            median_stride_bytes=("median_stride_bytes", "mean"),
+            same_page_fraction=("same_page_fraction", "mean"),
+        )
+        .reset_index()
+    )
+
+    return per_region, type_summary

mrm_trace/analyser/read_freq.py

@@ -0,0 +1,40 @@
+"""
+Read frequency analysis — total accesses, read fraction, reads-per-write.
+"""
+
+import pandas as pd
+
+
+def compute_read_freq(trace: pd.DataFrame) -> pd.DataFrame:
+    """
+    Compute read frequency statistics per region.
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_id, region_type, total_reads, total_writes,
+                 total_accesses, read_fraction, reads_per_write
+    """
+    region_types = trace.groupby("region_id")["region_type"].first()
+
+    counts = trace.groupby(["region_id", "op_type"]).size().unstack(fill_value=0)
+
+    # Ensure both columns exist even if all ops are one type
+    for col in ("load", "store"):
+        if col not in counts.columns:
+            counts[col] = 0
+
+    counts = counts.rename(columns={"load": "total_reads", "store": "total_writes"})
+    counts["total_accesses"] = counts["total_reads"] + counts["total_writes"]
+    counts["read_fraction"] = (counts["total_reads"] / counts["total_accesses"]).fillna(0.0)
+    counts["reads_per_write"] = (
+        counts["total_reads"] / counts["total_writes"].replace(0, float("nan"))
+    ).fillna(0.0)
+
+    result = counts.reset_index()
+    result["region_type"] = result["region_id"].map(region_types)
+
+    return result[
+        ["region_id", "region_type", "total_reads", "total_writes",
+         "total_accesses", "read_fraction", "reads_per_write"]
+    ]

mrm_trace/analyser/retention.py

@@ -0,0 +1,83 @@
+"""
+Retention analysis — how long each region is live between first write and last read.
+
+Per-region: first_write_ns, last_read_ns, retention_ns, retention_s
+Summary:    per region_type percentile stats (p50, p90, p99, mean)
+"""
+
+from typing import Tuple
+
+import pandas as pd
+
+
+def compute_retention(trace: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Compute retention duration for each region.
+
+    Returns
+    -------
+    per_region : DataFrame
+        Columns: region_id, region_type, first_write_ns, last_read_ns,
+                 retention_ns, retention_s
+    summary : DataFrame
+        Columns: region_type, n_regions, retention_p50_s, retention_p90_s,
+                 retention_p99_s, retention_mean_s
+    """
+    stores = trace[trace["op_type"] == "store"]
+    loads = trace[trace["op_type"] == "load"]
+
+    first_write = (
+        stores.groupby("region_id")["timestamp_ns"].min().rename("first_write_ns")
+    )
+
+    last_read = (
+        loads.groupby("region_id")["timestamp_ns"].max().rename("last_read_ns")
+    )
+
+    region_types = trace.groupby("region_id")["region_type"].first()
+
+    per_region = pd.concat([region_types, first_write, last_read], axis=1).reset_index()
+    per_region.columns = ["region_id", "region_type", "first_write_ns", "last_read_ns"]
+
+    # Regions that were never written — use first access as write time
+    no_write_mask = per_region["first_write_ns"].isna()
+    if no_write_mask.any():
+        fallback = (
+            trace.groupby("region_id")["timestamp_ns"].min().rename("first_write_ns")
+        )
+        per_region.loc[no_write_mask, "first_write_ns"] = per_region.loc[
+            no_write_mask, "region_id"
+        ].map(fallback)
+
+    # Regions with no reads — retention = 0
+    per_region["last_read_ns"] = per_region["last_read_ns"].fillna(
+        per_region["first_write_ns"]
+    )
+
+    per_region["retention_ns"] = (
+        per_region["last_read_ns"] - per_region["first_write_ns"]
+    ).clip(lower=0).astype("int64")
+    per_region["retention_s"] = per_region["retention_ns"] / 1e9
+
+    per_region = per_region.astype({"first_write_ns": "int64", "last_read_ns": "int64"})
+
+    # Summary per region_type
+    def _pct(g: pd.Series, q: float) -> float:
+        return float(g.quantile(q))
+
+    summary_rows = []
+    for rtype, grp in per_region.groupby("region_type"):
+        s = grp["retention_s"]
+        summary_rows.append(
+            {
+                "region_type": rtype,
+                "n_regions": len(grp),
+                "retention_p50_s": _pct(s, 0.50),
+                "retention_p90_s": _pct(s, 0.90),
+                "retention_p99_s": _pct(s, 0.99),
+                "retention_mean_s": float(s.mean()),
+            }
+        )
+    summary = pd.DataFrame(summary_rows)
+
+    return per_region, summary

mrm_trace/analyser/suitability.py

@@ -0,0 +1,63 @@
+"""
+MRM suitability classification.
+
+Thresholds
+----------
+high_mrm   : write_once_ratio >= 0.8  AND  retention_p99_s >= 10.0
+medium_mrm : write_once_ratio >= 0.5  AND  retention_p50_s >= 1.0
+low_mrm    : everything else
+"""
+
+import pandas as pd
+
+
+def classify_suitability(
+    retention_summary: pd.DataFrame,
+    write_once_per_region: pd.DataFrame,
+    read_freq_per_region: pd.DataFrame,
+    retention_threshold_high_s: float = 10.0,
+    retention_threshold_medium_s: float = 1.0,
+    write_once_threshold_high: float = 0.8,
+    write_once_threshold_medium: float = 0.5,
+) -> pd.DataFrame:
+    """
+    Classify each region_type as high_mrm, medium_mrm, or low_mrm.
+
+    Parameters
+    ----------
+    retention_summary : output of compute_retention()[1]
+    write_once_per_region : output of compute_write_once()
+    read_freq_per_region : output of compute_read_freq()
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_type, write_once_ratio_mean, retention_p99_s,
+                 retention_p50_s, suitability
+    """
+    wo_by_type = (
+        write_once_per_region.groupby("region_type")["write_once_ratio"]
+        .mean()
+        .rename("write_once_ratio_mean")
+        .reset_index()
+    )
+
+    merged = retention_summary.merge(wo_by_type, on="region_type", how="left")
+    merged["write_once_ratio_mean"] = merged["write_once_ratio_mean"].fillna(0.0)
+
+    def _classify(row: pd.Series) -> str:
+        wo = row["write_once_ratio_mean"]
+        p99 = row["retention_p99_s"]
+        p50 = row["retention_p50_s"]
+        if wo >= write_once_threshold_high and p99 >= retention_threshold_high_s:
+            return "high_mrm"
+        if wo >= write_once_threshold_medium and p50 >= retention_threshold_medium_s:
+            return "medium_mrm"
+        return "low_mrm"
+
+    merged["suitability"] = merged.apply(_classify, axis=1)
+
+    return merged[
+        ["region_type", "write_once_ratio_mean", "retention_p99_s",
+         "retention_p50_s", "suitability"]
+    ]

mrm_trace/analyser/working_set.py

@@ -0,0 +1,49 @@
+"""
+Working set size analysis — unique pages and bytes touched per region.
+"""
+
+from typing import Tuple
+
+import pandas as pd
+
+_PAGE_BYTES = 4096
+
+
+def compute_working_set(trace: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Compute working set size per region and a summary per region_type.
+
+    Returns
+    -------
+    per_region : DataFrame
+        Columns: region_id, region_type, unique_pages, working_set_bytes
+    type_summary : DataFrame
+        Columns: region_type, n_regions, total_pages, total_bytes,
+                 mean_pages_per_region, max_pages_per_region
+    """
+    region_types = trace.groupby("region_id")["region_type"].first()
+
+    per_region = (
+        trace.groupby("region_id")["address_page"]
+        .nunique()
+        .rename("unique_pages")
+        .reset_index()
+    )
+    per_region["region_type"] = per_region["region_id"].map(region_types)
+    per_region["working_set_bytes"] = per_region["unique_pages"] * _PAGE_BYTES
+
+    type_summary_rows = []
+    for rtype, grp in per_region.groupby("region_type"):
+        type_summary_rows.append(
+            {
+                "region_type": rtype,
+                "n_regions": len(grp),
+                "total_pages": int(grp["unique_pages"].sum()),
+                "total_bytes": int(grp["working_set_bytes"].sum()),
+                "mean_pages_per_region": float(grp["unique_pages"].mean()),
+                "max_pages_per_region": int(grp["unique_pages"].max()),
+            }
+        )
+    type_summary = pd.DataFrame(type_summary_rows)
+
+    return per_region, type_summary

mrm_trace/analyser/write_once.py

@@ -0,0 +1,50 @@
+"""
+Write-once analysis — what fraction of written addresses are written exactly once.
+
+High write-once ratio indicates write-once memory behaviour ideal for MRM.
+"""
+
+import pandas as pd
+
+
+def compute_write_once(trace: pd.DataFrame) -> pd.DataFrame:
+    """
+    Compute write-once ratio per region.
+
+    An address is "write-once" if it appears in exactly one store access.
+
+    Returns
+    -------
+    DataFrame
+        Columns: region_id, region_type, n_written_addresses,
+                 n_write_once_addresses, write_once_ratio
+    """
+    stores = trace[trace["op_type"] == "store"].copy()
+
+    if stores.empty:
+        return pd.DataFrame(
+            columns=[
+                "region_id", "region_type", "n_written_addresses",
+                "n_write_once_addresses", "write_once_ratio",
+            ]
+        )
+
+    write_counts = (
+        stores.groupby(["region_id", "address"]).size().reset_index(name="write_count")
+    )
+    region_types = trace.groupby("region_id")["region_type"].first().reset_index()
+
+    per_region = write_counts.groupby("region_id").agg(
+        n_written_addresses=("address", "count"),
+        n_write_once_addresses=("write_count", lambda x: (x == 1).sum()),
+    ).reset_index()
+
+    per_region = per_region.merge(region_types, on="region_id", how="left")
+    per_region["write_once_ratio"] = (
+        per_region["n_write_once_addresses"] / per_region["n_written_addresses"]
+    ).clip(0.0, 1.0)
+
+    return per_region[
+        ["region_id", "region_type", "n_written_addresses",
+         "n_write_once_addresses", "write_once_ratio"]
+    ]

mrm_trace/api.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import List, Union
+
+from mrm_trace.config.loader import load_experiment
+from mrm_trace.config.schema import ExperimentConfig, RunConfig
+from mrm_trace.config.validators import (
+    expand_sweep,
+    validate_config_semantics,
+    validate_environment,
+    validate_model_paths,
+)
+
+
+class Experiment:
+    """Python API for defining and executing mrm-trace experiments."""
+
+    def __init__(self, config: ExperimentConfig) -> None:
+        self._config = config
+
+    @classmethod
+    def from_yaml(cls, path: Union[str, Path]) -> "Experiment":
+        """Load an experiment from a YAML config file."""
+        return cls(load_experiment(path))
+
+    def validate(self) -> bool:
+        """Validate config and environment. Returns True if no errors."""
+        all_ok = True
+        for check_fn in (validate_config_semantics, validate_model_paths, validate_environment):
+            result = check_fn(self._config)
+            if not result.is_valid:
+                all_ok = False
+        return all_ok
+
+    def plan(self) -> List[RunConfig]:
+        """Materialise and return the full list of runs for this experiment."""
+        return expand_sweep(self._config)
+
+    def run(self) -> None:
+        """Execute all planned runs. Implemented in Phase 3."""
+        raise NotImplementedError("run() is available from Phase 3 onwards")
+
+    def analyse(self) -> None:
+        """Parse, label, and compute metrics. Implemented in Phase 4–6."""
+        raise NotImplementedError("analyse() is available from Phase 4 onwards")
+
+    def report(self, output_dir: Union[str, Path] = "reports/") -> None:
+        """Generate aggregated outputs and figures. Implemented in Phase 8."""
+        raise NotImplementedError("report() is available from Phase 8 onwards")
+
+    @property
+    def config(self) -> ExperimentConfig:
+        return self._config