MP config: load protected channel indices

scmp_kernels/mp/__init__.py

@@ -6,6 +6,8 @@
 from .config import (
     MPConfig,
     AdaptiveMPConfig,
+    compute_row_metric,
+    ROW_METRIC_NAMES,
     FreeBoundaryMPConfig,
     RangeMPConfig,
     RowAssignment,
@@ -28,6 +30,8 @@
 __all__ = [
     "MPConfig",
     "AdaptiveMPConfig",
+    "compute_row_metric",
+    "ROW_METRIC_NAMES",
     "FreeBoundaryMPConfig",
     "RangeMPConfig",
     "RowAssignment",

scmp_kernels/sc/kernels.py

@@ -92,6 +92,7 @@ def clear_rng_cache():
     _rng_seq_cache.clear()
     _enable_table_cache.clear()
     _k_table_cache.clear()
+    _cum_indicator_cache.clear()
 
 
 
@@ -589,8 +590,44 @@ def _sc_matmul_per_head_bipolar(
 # Enable-Signal Host Functions
 # =============================================================================
 
-_enable_table_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
-_k_table_cache: dict[str, torch.Tensor] = {}
+from collections import OrderedDict as _OrderedDict
+
+# Bounded LRU: autoregressive DECODE grows the attention inner dim every step,
+# so per-(D, stoc_len) tables otherwise accumulate without bound (MP's 5 levels
+# x ~50 decode steps x ~270MB @ ctx-4096 = tens of GB -> OOM). The cap is far
+# above any fixed-ctx (PPL) working set, so eviction never occurs there and
+# results are unaffected; on decode, stale-D entries are evicted (rebuild cost
+# only). Tables are pure functions of the key, so eviction is always safe.
+_ENABLE_TABLE_CACHE_MAX = int(os.environ.get("SC_ENABLE_TABLE_CACHE_MAX", "32"))
+_enable_table_cache: "_OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = _OrderedDict()
+_k_table_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict()
+# Chunked-MLP cum_indicator, same key discipline as _enable_table_cache.
+_cum_indicator_cache: "_OrderedDict[str, torch.Tensor]" = _OrderedDict()
+
+
+def _cum_cache_enabled() -> bool:
+    """SC_CUM_INDICATOR_CACHE=0 restores the per-call rebuild.
+
+    Exists so the bit-identity control run can be produced from the SAME
+    binary: the cache must not change a single output bit, only how often the
+    table is built.
+    """
+    return os.environ.get("SC_CUM_INDICATOR_CACHE", "1") != "0"
+
+
+def _lru_get(cache, key):
+    if key in cache:
+        cache.move_to_end(key)
+        return cache[key]
+    return None
+
+
+def _lru_put(cache, key, value):
+    cache[key] = value
+    cache.move_to_end(key)
+    while len(cache) > _ENABLE_TABLE_CACHE_MAX:
+        cache.popitem(last=False)
+    return value
 
 
 def _resolve_rng_levels(sc_prec: int, rng_levels: Optional[int]) -> int:
@@ -830,11 +867,45 @@ def _get_cached_enable_tables(
         stoc_len = 2 ** sc_prec
     grid_levels = _resolve_rng_levels(sc_prec, rng_levels)
     key = _enable_table_cache_key(config, sc_prec, device) + f"|sl={stoc_len}|rng={grid_levels}"
-    if key not in _enable_table_cache:
-        _enable_table_cache[key] = build_enable_tables(
-            rng_a, rng_b, sc_prec, stoc_len, rng_levels=grid_levels
-        )
-    return _enable_table_cache[key]
+    cached = _lru_get(_enable_table_cache, key)
+    if cached is not None:
+        return cached
+    return _lru_put(_enable_table_cache, key, build_enable_tables(
+        rng_a, rng_b, sc_prec, stoc_len, rng_levels=grid_levels
+    ))
+
+
+def _get_cached_cum_indicator(
+    key: str,
+    rng_b: torch.Tensor,
+    chunk_d: int,
+    stoc_len: int,
+    V: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Cache the chunked-MLP cum_indicator across calls.
+
+    The table is a pure function of (rng_b, chunk_d, stoc_len, V) and rng_b is
+    itself served from _get_cached_sequences, so every SCLinear in the model
+    that shares (config, sc_prec, stoc_len, rng_levels) builds a byte-identical
+    table.  The non-chunked (attention) path has cached this since day one via
+    _get_cached_enable_tables; the chunked MLP path rebuilt it on EVERY call --
+    measured at 37,621 rebuilds per 2048-token window on Qwen3-30B-A3B (128
+    experts x 3 projections x 48 layers) and 983 on the 14B dense model, each
+    one a 4.3 MB int16 allocation plus a Triton launch.
+
+    Returned tensors are READ-ONLY by contract: the matmul kernels only load
+    from cum_indicator, and nothing writes to it after build_cum_indicator_kernel.
+    """
+    cached = _lru_get(_cum_indicator_cache, key)
+    if cached is not None:
+        return cached
+    cum_indicator = torch.zeros(
+        chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device)
+    build_cum_indicator_kernel[(chunk_d,)](
+        rng_b, cum_indicator, chunk_d, stoc_len, V,
+    )
+    return _lru_put(_cum_indicator_cache, key, cum_indicator)
 
 
 def _get_cached_k_table(
@@ -848,11 +919,12 @@ def _get_cached_k_table(
         stoc_len = 2 ** sc_prec
     grid_levels = _resolve_rng_levels(sc_prec, rng_levels)
     key = _enable_table_cache_key(config, sc_prec, device) + f"|k_only|sl={stoc_len}|rng={grid_levels}"
-    if key not in _k_table_cache:
-        _k_table_cache[key] = build_k_table_only(
-            rng_a, sc_prec, stoc_len, rng_levels=grid_levels
-        )
-    return _k_table_cache[key]
+    cached = _lru_get(_k_table_cache, key)
+    if cached is not None:
+        return cached
+    return _lru_put(_k_table_cache, key, build_k_table_only(
+        rng_a, sc_prec, stoc_len, rng_levels=grid_levels
+    ))
 
 
 def enable_matmul_triton(
@@ -1335,6 +1407,7 @@ def _sc_matmul_bipolar_mlp_chunked(
     a, b, sc_prec, k_table, rng_b, chunk_d,
     stoc_len=None,
     rng_levels: Optional[int] = None,
+    cum_cache_key: Optional[str] = None,
 ):
     """
     Bipolar SC matmul for MLP with internal chunk_d loop.
@@ -1361,12 +1434,20 @@ def _sc_matmul_bipolar_mlp_chunked(
     device = a.device
     output = torch.zeros(N, M, dtype=torch.float32, device=device)
 
-    # Build cum_indicator ONCE — all chunks share the same RNG sequences
-    cum_indicator = torch.zeros(chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device)
-    build_cum_indicator_kernel[(chunk_d,)](
-        rng_b, cum_indicator,
-        chunk_d, stoc_len, V,
-    )
+    # Build cum_indicator ONCE — all chunks share the same RNG sequences.
+    # With a cache key, reuse it across CALLS too (see
+    # _get_cached_cum_indicator); without one, keep the per-call build so any
+    # future caller is bit-identical to the old behaviour by default.
+    if cum_cache_key is not None and _cum_cache_enabled():
+        cum_indicator = _get_cached_cum_indicator(
+            cum_cache_key, rng_b, chunk_d, stoc_len, V, device)
+    else:
+        cum_indicator = torch.zeros(
+            chunk_d, stoc_len + 1, V, dtype=torch.int16, device=device)
+        build_cum_indicator_kernel[(chunk_d,)](
+            rng_b, cum_indicator,
+            chunk_d, stoc_len, V,
+        )
 
     # Tiled matmul params — adaptive tile size for small N/M
     if N <= 64 or M <= 64:
@@ -1573,6 +1654,9 @@ def _sc_matmul_per_row_mlp(
         result = _sc_matmul_bipolar_mlp_chunked(
             a, b, sc_prec, k_table, rng_b, chunk_d,
             stoc_len=stoc_len, rng_levels=grid_levels,
+            cum_cache_key=(
+                _enable_table_cache_key(config, sc_prec, a.device)
+                + f"|cum|cd={chunk_d}|sl={stoc_len}|rng={grid_levels}"),
         )
 
         if device.type != 'cuda':

scmp_kernels/sc/matmul.py

@@ -28,6 +28,7 @@
     _sc_matmul_per_head_bipolar,
 )
 from ..quant.smoothquant import apply_smoothing
+from .. import trace as _trace
 
 
 _VALID_GRANULARITIES = ("per_tensor", "per_row", "per_head")
@@ -220,6 +221,34 @@ def _sc_matmul_impl(
                 f"sc_matmul: granularity='per_head' currently only supports "
                 f"mode='bipolar', got '{mode}'.")
 
+    # ---- precision trace (see scmp_kernels/trace.py) -------------------------
+    # Host-side shape metadata only — no tensor reads, no device sync. The
+    # bool guard keeps the off-path cost at one attribute read. Placed after
+    # the validation gates so rejected calls are never recorded, and after the
+    # halving block so stoc_len is the TRUE cycle count.
+    if _trace._ENABLED:
+        if a.dim() == 3:
+            _batch, _rows, _d_in = a.shape
+        else:
+            _rows, _d_in = a.shape
+            _batch = 1
+        _trace.record_matmul(
+            rows=int(_rows), d_in=int(_d_in), d_out=int(b.shape[-2]),
+            batch=int(_batch),
+            stoc_len=int(stoc_len) if stoc_len is not None else 2 ** sc_prec,
+            sc_prec=sc_prec, mode=mode, granularity=granularity,
+            # halving only takes effect in bipolar mode (guard above) — record
+            # whether it was APPLIED, not merely requested, to keep the field
+            # consistent with the effective stoc_len/rng_levels.
+            halve=bool(halve_bipolar_stoc_len and mode == "bipolar"),
+            # RESOLVED enable-grid size (mirrors _resolve_rng_levels), so the
+            # field has one meaning across halved and non-halved runs.
+            rng_levels=(int(rng_levels) if rng_levels is not None
+                        else 2 ** sc_prec),
+            chunk_d=chunk_d,
+            smoothed=smooth_scales is not None,
+        )
+
     # ---- dispatch ------------------------------------------------------------
     if granularity == "per_tensor":
         # Compute the per-tensor range on host (one .max / .min sync).