Update Feather a10g-large-gt80k training runtime image

.dockerignore

@@ -1,16 +1,16 @@
-# Keep HF runtime image context deterministic and small.
-**/__pycache__/
-**/*.py[cod]
-**/.pytest_cache/
-**/.mypy_cache/
-**/.ruff_cache/
-**/.venv/
-**/target/
-**/logs/
-**/*.log
-**/*.out
-**/*.pt
-**/*.safetensors
-**/*.parquet
-**/*.npz
-**/.git/
+# Keep HF runtime image context deterministic and small.
+**/__pycache__/
+**/*.py[cod]
+**/.pytest_cache/
+**/.mypy_cache/
+**/.ruff_cache/
+**/.venv/
+**/target/
+**/logs/
+**/*.log
+**/*.out
+**/*.pt
+**/*.safetensors
+**/*.parquet
+**/*.npz
+**/.git/

Dockerfile

@@ -1,124 +1,124 @@
-FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel
-
-# Default target is HF Jobs a10g-large (NVIDIA A10G, Ampere GA102, sm_86).
-# Override at build time for other cards, e.g. --build-arg FEATHER_GPU_ARCH=sm_90a.
-ARG FEATHER_GPU_ARCH=sm_86
-ARG FEATHER_TORCH_CUDA_ARCH_LIST=8.6
-
-ENV DEBIAN_FRONTEND=noninteractive \
-    PIP_NO_CACHE_DIR=1 \
-    PYTHONUNBUFFERED=1 \
-    CARGO_HOME=/root/.cargo \
-    RUSTUP_HOME=/root/.rustup \
-    HTM_CUDA_ARCH=${FEATHER_GPU_ARCH} \
-    TORCH_CUDA_ARCH_LIST=${FEATHER_TORCH_CUDA_ARCH_LIST} \
-    PATH=/root/.cargo/bin:${PATH}
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    git curl ca-certificates build-essential pkg-config libssl-dev && \
-    rm -rf /var/lib/apt/lists/*
-
-RUN curl https://sh.rustup.rs -sSf | bash -s -- -y --profile minimal --default-toolchain stable
-
-RUN pip install --upgrade pip setuptools wheel && \
-    pip install \
-      maturin \
-      huggingface_hub \
-      datasets \
-      requests \
-      pyarrow \
-      rustbpe \
-      pandas \
-      tiktoken \
-      pydantic \
-      ninja \
-      packaging \
-      einops
-
-# Mamba-3 fused CUDA kernel stack (mandatory — NO fallback allowed).
-#
-# We install PRE-BUILT manylinux wheels from the official state-spaces/mamba
-# and Dao-AILab/causal-conv1d GitHub releases. Compiling mamba_ssm from source
-# on HF Spaces' cpu-basic builder (~16GB RAM) OOMKills even with MAX_JOBS=1 —
-# nvcc on the templated selective-scan/chunk-scan kernels needs 8–12GB per TU.
-#
-# Wheel selection for base image pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel:
-#   - Python 3.11 (cp311)                       — matches PyTorch 2.5.1 image
-#   - CUDA 12.x wheels (cu12)                   — compatible with CUDA 12.1 base
-#   - PyTorch 2.5 ABI (torch2.5)                — exact torch match
-#   - cxx11abiFALSE                             — standard PyTorch pip build
-#
-# Versions: mamba_ssm 2.3.0 + causal_conv1d 1.6.0 (matching torch2.5 ABI).
-# Both are CUDA-compiled, no build toolchain needed
-# on the Space builder.
-#
-# Step A: install the published v2.3.0 prebuilt wheel (compiled CUDA ops
-# for selective_scan, layernorm_gated, ssd_*, causal_conv1d, etc).
-RUN pip install \
-      'https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.0/causal_conv1d-1.6.0+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' \
-      'https://github.com/state-spaces/mamba/releases/download/v2.3.0/mamba_ssm-2.3.0+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' && \
-    python -c "import importlib.metadata as m; print('installed mamba_ssm=' + m.version('mamba_ssm') + ' causal_conv1d=' + m.version('causal_conv1d'))"
-
-#
-# Step B: graft the Mamba3 class + its pure-Triton ops subtree from mamba-ssm
-# main. v2.3.1 is the latest release but Mamba3 landed post-release; the new
-# files under ops/triton/mamba3/ are ALL pure Python @triton.jit kernels with
-# zero compiled-CUDA dependencies (verified: every import in that subtree is
-# triton/torch/python — no .so files, no nvcc). So we install the v2.3.1 wheel
-# (for its compiled ops) and overlay the main-branch Mamba3 sources on top.
-#
-# This avoids the source-build OOM on the cpu-basic HF Space builder and the
-# missing-file error the smoke hit on the last attempt.
-# Download grafted mamba3 module + triton ops subtree
-RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
-    BASE=https://raw.githubusercontent.com/state-spaces/mamba/main && \
-    curl -fsSL "$BASE/mamba_ssm/modules/mamba3.py" -o "$SITE/modules/mamba3.py" && \
-    mkdir -p "$SITE/ops/triton/mamba3" && \
-    for f in __init__.py angle_dt.py mamba3_mimo_rotary_step.py mamba3_mimo_utils.py mamba3_siso_bwd.py mamba3_siso_combined.py mamba3_siso_fwd.py mamba3_siso_step.py utils.py; do \
-        curl -fsSL "$BASE/mamba_ssm/ops/triton/mamba3/$f" -o "$SITE/ops/triton/mamba3/$f"; \
-    done
-
-# Replace mamba_ssm/__init__.py with a minimal one that only imports Mamba3
-# (pure-Triton, works). The shipped __init__.py eagerly imports
-# selective_scan_cuda.so which has a libtorch C++ ABI mismatch on this base
-# image ("undefined symbol: _ZN3c107WarningC1E..."). Since training only needs
-# Mamba3 (grafted from main), we skip all compiled-CUDA imports.
-COPY mamba_ssm_init.py /opt/conda/lib/python3.11/site-packages/mamba_ssm/__init__.py
-
-# Structural check (no triton init — triton has no GPU on the builder)
-RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
-    test -f "$SITE/modules/mamba3.py" && \
-    test -f "$SITE/ops/triton/mamba3/mamba3_siso_combined.py" && \
-    test -s "$SITE/__init__.py" && \
-    echo "mamba3 graft + __init__ override verified"
-
-# Optional tilelang for MIMO path — pure-python, cheap; SISO Mamba3 works without.
-RUN pip install tilelang || echo "[dockerfile] tilelang optional install failed — continuing"
-
-# Triton version decision: FORCE 3.4.0 — first line with both mamba3
-# APIs (set_allocator + tl.make_tensor_descriptor) while avoiding the 3.5.x
-# driver-discovery regression seen on HF A10G (`0 active drivers` despite
-# torch.cuda being available). torch 2.5's _inductor expects older Triton
-# internals, but mamba_ssm/__init__.py shims AttrsDescriptor as a stub
-# before any torch._inductor import path runs, so the incompatibility is
-# neutralized. Build-time assert verifies mamba3's two required APIs.
-RUN pip install --force-reinstall --no-deps 'triton==3.4.0' && \
-    python -c "import triton; from triton import language as tl; \
-               assert hasattr(triton, 'set_allocator'), 'missing triton.set_allocator'; \
-               assert hasattr(tl, 'make_tensor_descriptor'), 'missing tl.make_tensor_descriptor'; \
-               print(f'triton={triton.__version__} set_allocator+make_tensor_descriptor OK, AttrsDescriptor shimmed in mamba_ssm/__init__.py')"
-
-WORKDIR /workspace
-COPY overlay /workspace/feather
-COPY entrypoint.py /app/entrypoint.py
-WORKDIR /workspace/feather
-
-RUN python -m py_compile hydra/training.py prepare.py train.py && \
-    bash -n scripts/run_domain_expanded_pretrain.sh
-
-RUN export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} && \
-    echo "building htm_rust GPU kernels for HTM_CUDA_ARCH=${HTM_CUDA_ARCH} TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}" && \
-    maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml && \
-    pip install htm_rust/target/wheels/htm_rust-*.whl
-
-CMD ["python", "/app/entrypoint.py"]
+FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel
+
+# Default target is HF Jobs a10g-large (NVIDIA A10G, Ampere GA102, sm_86).
+# Override at build time for other cards, e.g. --build-arg FEATHER_GPU_ARCH=sm_90a.
+ARG FEATHER_GPU_ARCH=sm_86
+ARG FEATHER_TORCH_CUDA_ARCH_LIST=8.6
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=1 \
+    PYTHONUNBUFFERED=1 \
+    CARGO_HOME=/root/.cargo \
+    RUSTUP_HOME=/root/.rustup \
+    HTM_CUDA_ARCH=${FEATHER_GPU_ARCH} \
+    TORCH_CUDA_ARCH_LIST=${FEATHER_TORCH_CUDA_ARCH_LIST} \
+    PATH=/root/.cargo/bin:${PATH}
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git curl ca-certificates build-essential pkg-config libssl-dev && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y --profile minimal --default-toolchain stable
+
+RUN pip install --upgrade pip setuptools wheel && \
+    pip install \
+      maturin \
+      huggingface_hub \
+      datasets \
+      requests \
+      pyarrow \
+      rustbpe \
+      pandas \
+      tiktoken \
+      pydantic \
+      ninja \
+      packaging \
+      einops
+
+# Mamba-3 fused CUDA kernel stack (mandatory — NO fallback allowed).
+#
+# We install PRE-BUILT manylinux wheels from the official state-spaces/mamba
+# and Dao-AILab/causal-conv1d GitHub releases. Compiling mamba_ssm from source
+# on HF Spaces' cpu-basic builder (~16GB RAM) OOMKills even with MAX_JOBS=1 —
+# nvcc on the templated selective-scan/chunk-scan kernels needs 8–12GB per TU.
+#
+# Wheel selection for base image pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel:
+#   - Python 3.11 (cp311)                       — matches PyTorch 2.5.1 image
+#   - CUDA 12.x wheels (cu12)                   — compatible with CUDA 12.1 base
+#   - PyTorch 2.5 ABI (torch2.5)                — exact torch match
+#   - cxx11abiFALSE                             — standard PyTorch pip build
+#
+# Versions: mamba_ssm 2.3.0 + causal_conv1d 1.6.0 (matching torch2.5 ABI).
+# Both are CUDA-compiled, no build toolchain needed
+# on the Space builder.
+#
+# Step A: install the published v2.3.0 prebuilt wheel (compiled CUDA ops
+# for selective_scan, layernorm_gated, ssd_*, causal_conv1d, etc).
+RUN pip install \
+      'https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.0/causal_conv1d-1.6.0+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' \
+      'https://github.com/state-spaces/mamba/releases/download/v2.3.0/mamba_ssm-2.3.0+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' && \
+    python -c "import importlib.metadata as m; print('installed mamba_ssm=' + m.version('mamba_ssm') + ' causal_conv1d=' + m.version('causal_conv1d'))"
+
+#
+# Step B: graft the Mamba3 class + its pure-Triton ops subtree from mamba-ssm
+# main. v2.3.1 is the latest release but Mamba3 landed post-release; the new
+# files under ops/triton/mamba3/ are ALL pure Python @triton.jit kernels with
+# zero compiled-CUDA dependencies (verified: every import in that subtree is
+# triton/torch/python — no .so files, no nvcc). So we install the v2.3.1 wheel
+# (for its compiled ops) and overlay the main-branch Mamba3 sources on top.
+#
+# This avoids the source-build OOM on the cpu-basic HF Space builder and the
+# missing-file error the smoke hit on the last attempt.
+# Download grafted mamba3 module + triton ops subtree
+RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
+    BASE=https://raw.githubusercontent.com/state-spaces/mamba/main && \
+    curl -fsSL "$BASE/mamba_ssm/modules/mamba3.py" -o "$SITE/modules/mamba3.py" && \
+    mkdir -p "$SITE/ops/triton/mamba3" && \
+    for f in __init__.py angle_dt.py mamba3_mimo_rotary_step.py mamba3_mimo_utils.py mamba3_siso_bwd.py mamba3_siso_combined.py mamba3_siso_fwd.py mamba3_siso_step.py utils.py; do \
+        curl -fsSL "$BASE/mamba_ssm/ops/triton/mamba3/$f" -o "$SITE/ops/triton/mamba3/$f"; \
+    done
+
+# Replace mamba_ssm/__init__.py with a minimal one that only imports Mamba3
+# (pure-Triton, works). The shipped __init__.py eagerly imports
+# selective_scan_cuda.so which has a libtorch C++ ABI mismatch on this base
+# image ("undefined symbol: _ZN3c107WarningC1E..."). Since training only needs
+# Mamba3 (grafted from main), we skip all compiled-CUDA imports.
+COPY mamba_ssm_init.py /opt/conda/lib/python3.11/site-packages/mamba_ssm/__init__.py
+
+# Structural check (no triton init — triton has no GPU on the builder)
+RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
+    test -f "$SITE/modules/mamba3.py" && \
+    test -f "$SITE/ops/triton/mamba3/mamba3_siso_combined.py" && \
+    test -s "$SITE/__init__.py" && \
+    echo "mamba3 graft + __init__ override verified"
+
+# Optional tilelang for MIMO path — pure-python, cheap; SISO Mamba3 works without.
+RUN pip install tilelang || echo "[dockerfile] tilelang optional install failed — continuing"
+
+# Triton version decision: FORCE 3.4.0 — first line with both mamba3
+# APIs (set_allocator + tl.make_tensor_descriptor) while avoiding the 3.5.x
+# driver-discovery regression seen on HF A10G (`0 active drivers` despite
+# torch.cuda being available). torch 2.5's _inductor expects older Triton
+# internals, but mamba_ssm/__init__.py shims AttrsDescriptor as a stub
+# before any torch._inductor import path runs, so the incompatibility is
+# neutralized. Build-time assert verifies mamba3's two required APIs.
+RUN pip install --force-reinstall --no-deps 'triton==3.4.0' && \
+    python -c "import triton; from triton import language as tl; \
+               assert hasattr(triton, 'set_allocator'), 'missing triton.set_allocator'; \
+               assert hasattr(tl, 'make_tensor_descriptor'), 'missing tl.make_tensor_descriptor'; \
+               print(f'triton={triton.__version__} set_allocator+make_tensor_descriptor OK, AttrsDescriptor shimmed in mamba_ssm/__init__.py')"
+
+WORKDIR /workspace
+COPY overlay /workspace/feather
+COPY entrypoint.py /app/entrypoint.py
+WORKDIR /workspace/feather
+
+RUN python -m py_compile hydra/training.py prepare.py train.py && \
+    bash -n scripts/run_domain_expanded_pretrain.sh
+
+RUN export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} && \
+    echo "building htm_rust GPU kernels for HTM_CUDA_ARCH=${HTM_CUDA_ARCH} TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}" && \
+    maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml && \
+    pip install htm_rust/target/wheels/htm_rust-*.whl
+
+CMD ["python", "/app/entrypoint.py"]

entrypoint.py

@@ -1,260 +1,278 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-import os
-import subprocess
-import sys
-import time
-from http.server import BaseHTTPRequestHandler, HTTPServer
-from pathlib import Path
-from threading import Thread
-
-
-def _prepend_library_path(*paths: str) -> None:
-    """Expose injected NVIDIA driver libraries before torch/triton imports."""
-    existing = [p for p in os.environ.get('LD_LIBRARY_PATH', '').split(':') if p]
-    merged = []
-    for p in paths:
-        if p and p not in merged:
-            merged.append(p)
-    for p in existing:
-        if p not in merged:
-            merged.append(p)
-    os.environ['LD_LIBRARY_PATH'] = ':'.join(merged)
-
-
-_prepend_library_path(
-    # HF Jobs injects the host driver under /usr/local/nvidia. Prefer that
-    # over CUDA toolkit/compat libcuda stubs; using /usr/local/cuda/compat here
-    # made A10G PyTorch report Error 803 despite nvidia-smi working.
-    '/usr/local/nvidia/lib64',
-    '/usr/local/nvidia/lib',
-    '/usr/lib/x86_64-linux-gnu',
-)
-
-
-# =============================================================================
-# EARLY CUDA FABRIC MANAGER KICK (before ANY CUDA-touching imports)
-# =============================================================================
-# On HF GPU hosts, cudaGetDeviceCount can transiently return not-ready errors
-# on first use. H200 fabric-manager is the worst case; A10G is usually ready
-# immediately, but the same early kick keeps the runtime deterministic.
-# synchronizes with the container's first driver call. Once any NVML/CUDA
-# call succeeds once (even just nvidia-smi), the fabric is up for the rest
-# of the container lifetime.
-#
-# Our previous approach (wait in a subprocess before training) didn't work
-# because the "initialization failed" state persisted across calls in the
-# same container. The real fix: kick the driver exactly once with
-# nvidia-smi, which is what successfully-working baseline containers do
-# implicitly via their first torch.cuda call.
-#
-# Must happen BEFORE `import torch` (because any import that eagerly calls
-# cudaGetDeviceCount will cache the Error 802 state).
-def _early_cuda_kick() -> None:
-    deadline = time.time() + 120.0
-    attempt = 0
-    while time.time() < deadline:
-        attempt += 1
-        r = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=30)
-        if r.returncode == 0:
-            gpu_line = next((ln.strip() for ln in (r.stdout or '').splitlines() if any(g in ln for g in ('A10', 'A100', 'H100', 'H200', 'RTX'))), 'gpu=unknown')
-            print(f'[boot] nvidia-smi OK on attempt {attempt}: {gpu_line}', flush=True)
-            break
-        print(f'[boot] nvidia-smi attempt {attempt} rc={r.returncode} stderr={(r.stderr or "")[:120]}',
-              flush=True)
-        time.sleep(2)
-    # After nvidia-smi, probe torch in a subprocess so any latent error state
-    # doesn't leak into the main process's CUDA context.
-    probe = 'import torch; import sys; sys.exit(0 if torch.cuda.is_available() else 1)'
-    torch_deadline = time.time() + 120.0
-    t_attempt = 0
-    while time.time() < torch_deadline:
-        t_attempt += 1
-        r = subprocess.run([sys.executable, '-c', probe], capture_output=True, text=True, timeout=60)
-        if r.returncode == 0:
-            print(f'[boot] torch.cuda.is_available() = True after {t_attempt} probe(s)', flush=True)
-            return
-        if t_attempt == 1:
-            print(f'[boot] torch cuda probe {t_attempt}: {(r.stderr or "")[:200]}', flush=True)
-        time.sleep(2)
-    print('[boot] WARNING: torch.cuda never became ready — training will likely fail', flush=True)
-
-
-_early_cuda_kick()
-
-# Hydrate triton compilation cache from HF Hub before any triton/mamba_ssm import.
-# triton_cache_setup.py is copied next to this file by the job bash command.
-try:
-    import triton_cache_setup as _tcs
-    _tcs.setup()
-except ImportError:
-    print('[boot] triton_cache_setup not found; skipping cache hydrate', flush=True)
-
-from huggingface_hub import HfApi  # noqa: E402  (import after cuda kick)
-
-REPO_ROOT = Path('/workspace/feather')
-CACHE_ROOT = Path.home() / '.cache' / 'autoresearch'
-LOG_FILE = REPO_ROOT / 'run_domain_expanded.log'
-JOB_ID = os.environ.get('JOB_ID', 'local-job')
-OUTPUT_REPO = os.environ.get('HF_REPO_ID', 'icarus112/feather-pretrain-checkpoints')
-TOKEN = os.environ.get('HF_TOKEN')
-RUNTIME_MODE = os.environ.get('FEATHER_RUNTIME_MODE', 'space')
-APP_PORT = int(os.environ.get('PORT', '7860'))
-
-
-class _HealthHandler(BaseHTTPRequestHandler):
-    def do_GET(self):
-        if self.path in ('/', '/health', '/healthz', '/ready'):
-            payload = {
-                'status': 'ok',
-                'mode': RUNTIME_MODE,
-                'job_id': JOB_ID,
-            }
-            body = json.dumps(payload).encode('utf-8')
-            self.send_response(200)
-            self.send_header('Content-Type', 'application/json')
-            self.send_header('Content-Length', str(len(body)))
-            self.end_headers()
-            self.wfile.write(body)
-            return
-        self.send_response(404)
-        self.end_headers()
-
-    def log_message(self, format, *args):
-        return
-
-
-def _start_health_server() -> HTTPServer:
-    server = HTTPServer(('0.0.0.0', APP_PORT), _HealthHandler)
-    thread = Thread(target=server.serve_forever, daemon=True)
-    thread.start()
-    print(f'[space] health server listening on 0.0.0.0:{APP_PORT}', flush=True)
-    return server
-
-
-def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
-    if not path.exists():
-        print(f'[upload] skip missing {path}', flush=True)
-        return
-    api.upload_file(
-        path_or_fileobj=str(path),
-        path_in_repo=dest,
-        repo_id=OUTPUT_REPO,
-        repo_type='model',
-    )
-    print(f'[upload] uploaded {path} -> {OUTPUT_REPO}/{dest}', flush=True)
-
-
-def _wait_for_cuda_ready(timeout_s: int = 120) -> None:
-    """Block until CUDA is fully initialized or timeout.
-
-    On H200 hosts with NVSwitch/fabric manager, nvidia driver setup can race
-    with container start. cudaGetDeviceCount can return CUDA_ERROR_SYSTEM_NOT_READY
-    (error 802) for the first few seconds, and any import that triggers
-    @triton.autotune (e.g. mamba_ssm, torch amp utilities) blows up with
-    "0 active drivers" if it happens during that window.
-
-    We pre-init CUDA in a throwaway Python subprocess (so any error state does
-    not leak into the main training process) and retry until torch.cuda
-    reports ready.
-    """
-    import time as _t
-    probe = (
-        "import torch; "
-        "import sys; "
-        "avail = torch.cuda.is_available(); "
-        "count = torch.cuda.device_count() if avail else 0; "
-        "torch.empty(1, device='cuda') if (avail and count > 0) else None; "
-        "from triton.runtime import driver; "
-        "driver.active.get_current_device(); "
-        "sys.exit(0 if (avail and count > 0) else 1)"
-    )
-    deadline = _t.time() + timeout_s
-    attempt = 0
-    while _t.time() < deadline:
-        attempt += 1
-        r = subprocess.run(['python', '-c', probe], capture_output=True, text=True)
-        if r.returncode == 0:
-            print(f'[job] CUDA/Triton ready after {attempt} probe(s)', flush=True)
-            return
-        if attempt == 1:
-            print(f'[job] CUDA not ready yet (will retry up to {timeout_s}s): {r.stderr.strip()[:200]}', flush=True)
-        _t.sleep(2)
-    print(f'[job] CUDA still not ready after {timeout_s}s — continuing anyway (training will likely fail)', flush=True)
-
-
-def run_job_mode() -> int:
-    os.chdir(REPO_ROOT)
-    os.environ.setdefault('HYDRA_TIME_BUDGET', '43200')
-    os.environ.setdefault('HYDRA_TARGET_SHARDS', '2048')
-    os.environ.setdefault('HYDRA_DOWNLOAD_WORKERS', '16')
-    os.environ.setdefault('HYDRA_CKPT_INTERVAL', '1000')
-    os.environ.setdefault('HYDRA_RESUME_CKPT', str(CACHE_ROOT / 'latest.pt'))
-    os.environ.setdefault('FEATHER_GPU_PROFILE', 'a10g-large')
-    os.environ.setdefault('HTM_CUDA_ARCH', 'sm_86')
-    os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '8.6')
-    os.environ.setdefault('TRITON_CACHE_DIR', f"/workspace/triton_cache/{os.environ['FEATHER_GPU_PROFILE']}")
-    os.environ.setdefault('TRITON_CACHE_REPO', f"icarus112/feather-triton-cache-{os.environ['FEATHER_GPU_PROFILE']}")
-    print(f"[job] gpu_profile={os.environ['FEATHER_GPU_PROFILE']} htm_cuda_arch={os.environ['HTM_CUDA_ARCH']} torch_cuda_arch={os.environ['TORCH_CUDA_ARCH_LIST']}", flush=True)
-
-    # CUDA readiness was kicked at module import via _early_cuda_kick. Keep
-    # the wait as a second safety net — no-op if CUDA already ready.
-    _wait_for_cuda_ready()
-
-    cmd = [
-        'bash',
-        './scripts/run_domain_expanded_pretrain.sh',
-        '--target-shards', os.environ['HYDRA_TARGET_SHARDS'],
-        '--download-workers', os.environ['HYDRA_DOWNLOAD_WORKERS'],
-    ]
-    print('[job] starting Feather domain-expanded pretrain', flush=True)
-    print(f'[job] command={cmd}', flush=True)
-    proc = subprocess.run(cmd, check=False)
-
-    # Push triton compilation cache back to HF Hub for next run.
-    try:
-        import triton_cache_setup as _tcs
-        _tcs.teardown()
-    except Exception as _tcs_err:
-        print(f'[triton_cache] teardown error (non-fatal): {_tcs_err}', flush=True)
-
-    if TOKEN:
-        api = HfApi(token=TOKEN)
-        try:
-            api.create_repo(repo_id=OUTPUT_REPO, repo_type='model', private=True, exist_ok=True)
-        except Exception as e:
-            print(f'[upload] create_repo warning: {type(e).__name__}: {e}', flush=True)
-        prefix = f'jobs/{JOB_ID}'
-        try:
-            upload_artifact(api, LOG_FILE, f'{prefix}/run_domain_expanded.log')
-            upload_artifact(api, CACHE_ROOT / 'latest.pt', f'{prefix}/latest.pt')
-            upload_artifact(api, CACHE_ROOT / 'pretrain_final.pt', f'{prefix}/pretrain_final.pt')
-        except Exception as e:
-            print(f'[upload] upload warning: {type(e).__name__}: {e}', flush=True)
-    else:
-        print('[upload] HF_TOKEN not set; skipping artifact upload', flush=True)
-
-    return proc.returncode
-
-
-def run_space_mode() -> int:
-    server = _start_health_server()
-    print('[space] Feather runtime image ready', flush=True)
-    try:
-        while True:
-            time.sleep(3600)
-    finally:
-        server.shutdown()
-        server.server_close()
-
-
-def main() -> int:
-    if RUNTIME_MODE == 'job':
-        return run_job_mode()
-    return run_space_mode()
-
-
-if __name__ == '__main__':
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import time
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from threading import Thread
+
+
+def _prepend_library_path(*paths: str) -> None:
+    """Expose injected NVIDIA driver libraries before torch/triton imports."""
+    existing = [p for p in os.environ.get('LD_LIBRARY_PATH', '').split(':') if p]
+    merged = []
+    for p in paths:
+        if p and p not in merged:
+            merged.append(p)
+    for p in existing:
+        if p not in merged:
+            merged.append(p)
+    os.environ['LD_LIBRARY_PATH'] = ':'.join(merged)
+
+
+_prepend_library_path(
+    # HF Jobs injects the host driver under /usr/local/nvidia. Prefer that
+    # over CUDA toolkit/compat libcuda stubs; using /usr/local/cuda/compat here
+    # made A10G PyTorch report Error 803 despite nvidia-smi working.
+    '/usr/local/nvidia/lib64',
+    '/usr/local/nvidia/lib',
+    '/usr/lib/x86_64-linux-gnu',
+)
+
+
+# =============================================================================
+# EARLY CUDA FABRIC MANAGER KICK (before ANY CUDA-touching imports)
+# =============================================================================
+# On HF GPU hosts, cudaGetDeviceCount can transiently return not-ready errors
+# on first use. H200 fabric-manager is the worst case; A10G is usually ready
+# immediately, but the same early kick keeps the runtime deterministic.
+# synchronizes with the container's first driver call. Once any NVML/CUDA
+# call succeeds once (even just nvidia-smi), the fabric is up for the rest
+# of the container lifetime.
+#
+# Our previous approach (wait in a subprocess before training) didn't work
+# because the "initialization failed" state persisted across calls in the
+# same container. The real fix: kick the driver exactly once with
+# nvidia-smi, which is what successfully-working baseline containers do
+# implicitly via their first torch.cuda call.
+#
+# Must happen BEFORE `import torch` (because any import that eagerly calls
+# cudaGetDeviceCount will cache the Error 802 state).
+def _early_cuda_kick() -> None:
+    deadline = time.time() + 120.0
+    attempt = 0
+    while time.time() < deadline:
+        attempt += 1
+        r = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=30)
+        if r.returncode == 0:
+            gpu_line = next((ln.strip() for ln in (r.stdout or '').splitlines() if any(g in ln for g in ('A10', 'A100', 'H100', 'H200', 'RTX'))), 'gpu=unknown')
+            print(f'[boot] nvidia-smi OK on attempt {attempt}: {gpu_line}', flush=True)
+            break
+        print(f'[boot] nvidia-smi attempt {attempt} rc={r.returncode} stderr={(r.stderr or "")[:120]}',
+              flush=True)
+        time.sleep(2)
+    # After nvidia-smi, probe torch in a subprocess so any latent error state
+    # doesn't leak into the main process's CUDA context.
+    probe = 'import torch; import sys; sys.exit(0 if torch.cuda.is_available() else 1)'
+    torch_deadline = time.time() + 120.0
+    t_attempt = 0
+    while time.time() < torch_deadline:
+        t_attempt += 1
+        r = subprocess.run([sys.executable, '-c', probe], capture_output=True, text=True, timeout=60)
+        if r.returncode == 0:
+            print(f'[boot] torch.cuda.is_available() = True after {t_attempt} probe(s)', flush=True)
+            return
+        if t_attempt == 1:
+            print(f'[boot] torch cuda probe {t_attempt}: {(r.stderr or "")[:200]}', flush=True)
+        time.sleep(2)
+    print('[boot] WARNING: torch.cuda never became ready — training will likely fail', flush=True)
+
+
+_early_cuda_kick()
+
+# Hydrate triton compilation cache from HF Hub before any triton/mamba_ssm import.
+# triton_cache_setup.py is copied next to this file by the job bash command.
+try:
+    import triton_cache_setup as _tcs
+    _tcs.setup()
+except ImportError:
+    print('[boot] triton_cache_setup not found; skipping cache hydrate', flush=True)
+
+from huggingface_hub import HfApi  # noqa: E402  (import after cuda kick)
+
+REPO_ROOT = Path('/workspace/feather')
+CACHE_ROOT = Path.home() / '.cache' / 'autoresearch'
+LOG_FILE = REPO_ROOT / 'run_domain_expanded.log'
+JOB_ID = os.environ.get('JOB_ID', 'local-job')
+OUTPUT_REPO = os.environ.get('HF_REPO_ID', 'icarus112/feather-pretrain-checkpoints')
+TOKEN = os.environ.get('HF_TOKEN')
+RUNTIME_MODE = os.environ.get('FEATHER_RUNTIME_MODE', 'space')
+APP_PORT = int(os.environ.get('PORT', '7860'))
+
+
+class _HealthHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path in ('/', '/health', '/healthz', '/ready'):
+            payload = {
+                'status': 'ok',
+                'mode': RUNTIME_MODE,
+                'job_id': JOB_ID,
+            }
+            body = json.dumps(payload).encode('utf-8')
+            self.send_response(200)
+            self.send_header('Content-Type', 'application/json')
+            self.send_header('Content-Length', str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+            return
+        self.send_response(404)
+        self.end_headers()
+
+    def log_message(self, format, *args):
+        return
+
+
+def _start_health_server() -> HTTPServer:
+    server = HTTPServer(('0.0.0.0', APP_PORT), _HealthHandler)
+    thread = Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    print(f'[space] health server listening on 0.0.0.0:{APP_PORT}', flush=True)
+    return server
+
+
+def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
+    if not path.exists():
+        print(f'[upload] skip missing {path}', flush=True)
+        return
+    api.upload_file(
+        path_or_fileobj=str(path),
+        path_in_repo=dest,
+        repo_id=OUTPUT_REPO,
+        repo_type='model',
+    )
+    print(f'[upload] uploaded {path} -> {OUTPUT_REPO}/{dest}', flush=True)
+
+
+def _wait_for_cuda_ready(timeout_s: int = 120) -> None:
+    """Block until CUDA is fully initialized or timeout.
+
+    On H200 hosts with NVSwitch/fabric manager, nvidia driver setup can race
+    with container start. cudaGetDeviceCount can return CUDA_ERROR_SYSTEM_NOT_READY
+    (error 802) for the first few seconds, and any import that triggers
+    @triton.autotune (e.g. mamba_ssm, torch amp utilities) blows up with
+    "0 active drivers" if it happens during that window.
+
+    We pre-init CUDA in a throwaway Python subprocess (so any error state does
+    not leak into the main training process) and retry until torch.cuda
+    reports ready.
+    """
+    import time as _t
+    probe = (
+        "import torch; "
+        "import sys; "
+        "avail = torch.cuda.is_available(); "
+        "count = torch.cuda.device_count() if avail else 0; "
+        "torch.empty(1, device='cuda') if (avail and count > 0) else None; "
+        "from triton.runtime import driver; "
+        "driver.active.get_current_device(); "
+        "sys.exit(0 if (avail and count > 0) else 1)"
+    )
+    deadline = _t.time() + timeout_s
+    attempt = 0
+    while _t.time() < deadline:
+        attempt += 1
+        r = subprocess.run(['python', '-c', probe], capture_output=True, text=True)
+        if r.returncode == 0:
+            print(f'[job] CUDA/Triton ready after {attempt} probe(s)', flush=True)
+            return
+        if attempt == 1:
+            print(f'[job] CUDA not ready yet (will retry up to {timeout_s}s): {r.stderr.strip()[:200]}', flush=True)
+        _t.sleep(2)
+    print(f'[job] CUDA still not ready after {timeout_s}s — continuing anyway (training will likely fail)', flush=True)
+
+
+def run_job_mode() -> int:
+    os.chdir(REPO_ROOT)
+    os.environ.setdefault('HYDRA_TIME_BUDGET', '43200')
+    os.environ.setdefault('HYDRA_TARGET_SHARDS', '2048')
+    os.environ.setdefault('HYDRA_DOWNLOAD_WORKERS', '16')
+    os.environ.setdefault('HYDRA_CKPT_INTERVAL', '1000')
+    os.environ.setdefault('HYDRA_RESUME_CKPT', str(CACHE_ROOT / 'latest.pt'))
+    os.environ.setdefault('FEATHER_GPU_PROFILE', 'a10g-large')
+    os.environ.setdefault('HTM_CUDA_ARCH', 'sm_86')
+    os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '8.6')
+    os.environ.setdefault('TRITON_CACHE_DIR', f"/workspace/triton_cache/{os.environ['FEATHER_GPU_PROFILE']}")
+    os.environ.setdefault('TRITON_CACHE_REPO', f"icarus112/feather-triton-cache-{os.environ['FEATHER_GPU_PROFILE']}")
+    print(f"[job] gpu_profile={os.environ['FEATHER_GPU_PROFILE']} htm_cuda_arch={os.environ['HTM_CUDA_ARCH']} torch_cuda_arch={os.environ['TORCH_CUDA_ARCH_LIST']}", flush=True)
+
+    # CUDA readiness was kicked at module import via _early_cuda_kick. Keep
+    # the wait as a second safety net — no-op if CUDA already ready.
+    _wait_for_cuda_ready()
+
+    # ---------- FIX 1: retina.npz auto-generation ----------
+    # SemanticFoldingSDR requires retina.npz at model init. Generate it from
+    # the shard data (if not already cached) before importing the model.
+    retina_path = CACHE_ROOT / 'retina.npz'
+    if not retina_path.exists():
+        print(f'[job] retina.npz missing at {retina_path} — generating from shards...', flush=True)
+        r = subprocess.run(
+            [sys.executable, '-u', 'subsystems/sdr_retina.py'],
+            check=False, timeout=1800,
+        )
+        if r.returncode != 0 or not retina_path.exists():
+            print(f'[job] FATAL: retina generation failed (rc={r.returncode})', flush=True)
+            return 1
+        print(f'[job] retina.npz generated ({retina_path.stat().st_size} bytes)', flush=True)
+    else:
+        print(f'[job] retina.npz already cached', flush=True)
+    # --------------------------------------------------------
+
+    # ---------- FIX 2: direct train.py (no bash wrapper) ----------
+    # run_domain_expanded_pretrain.sh drops HYDRA_* env vars at the bash→py
+    # boundary. Bypass it: stream data at train-time via Nemotron path, call
+    # train.py directly with all env vars intact.
+    print('[job] starting Feather training (direct train.py, streaming data)...', flush=True)
+    train_cmd = [sys.executable, '-u', 'train.py']
+    print(f'[job] command={train_cmd}', flush=True)
+    proc = subprocess.run(train_cmd, check=False)
+    # ------------------------------------------------------------
+
+    # Push triton compilation cache back to HF Hub for next run.
+    try:
+        import triton_cache_setup as _tcs
+        _tcs.teardown()
+    except Exception as _tcs_err:
+        print(f'[triton_cache] teardown error (non-fatal): {_tcs_err}', flush=True)
+
+    if TOKEN:
+        api = HfApi(token=TOKEN)
+        try:
+            api.create_repo(repo_id=OUTPUT_REPO, repo_type='model', private=True, exist_ok=True)
+        except Exception as e:
+            print(f'[upload] create_repo warning: {type(e).__name__}: {e}', flush=True)
+        prefix = f'jobs/{JOB_ID}'
+        try:
+            upload_artifact(api, LOG_FILE, f'{prefix}/run_domain_expanded.log')
+            upload_artifact(api, CACHE_ROOT / 'latest.pt', f'{prefix}/latest.pt')
+            upload_artifact(api, CACHE_ROOT / 'pretrain_final.pt', f'{prefix}/pretrain_final.pt')
+        except Exception as e:
+            print(f'[upload] upload warning: {type(e).__name__}: {e}', flush=True)
+    else:
+        print('[upload] HF_TOKEN not set; skipping artifact upload', flush=True)
+
+    return proc.returncode
+
+
+def run_space_mode() -> int:
+    server = _start_health_server()
+    print('[space] Feather runtime image ready', flush=True)
+    try:
+        while True:
+            time.sleep(3600)
+    finally:
+        server.shutdown()
+        server.server_close()
+
+
+def main() -> int:
+    if RUNTIME_MODE == 'job':
+        return run_job_mode()
+    return run_space_mode()
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

mamba_ssm_init.py

@@ -1,69 +1,69 @@
-# mamba_ssm package init — minimal override to avoid broken selective_scan_cuda.so
-# ABI mismatch with the base image's libtorch.
-#
-# The upstream __init__.py eagerly imports selective_scan_cuda which fails on
-# pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel (undefined c10::Warning ctor
-# symbol). We only need Mamba3 (grafted from main, pure-Triton), so we skip
-# all compiled-CUDA imports here and let Mamba3 load directly.
-
-__version__ = "2.3.1+feather-graft"
-
-# selective_scan_fn / mamba_inner_fn are shimmed to None — they are NOT used
-# by the Feather training path (which is Mamba3-only). If any import path
-# hits this, it will get a clear AttributeError instead of an obscure ImportError.
-selective_scan_fn = None
-mamba_inner_fn = None
-
-# --- triton API compatibility shims -----------------------------------------
-# Version matrix is hostile: torch 2.6 pins triton==3.2.0 because torch._inductor
-# imports AttrsDescriptor from triton.compiler.compiler — removed in triton 3.4+.
-# Grafted Mamba3 (from mamba-ssm main) needs triton.set_allocator and
-# tl.make_tensor_descriptor, both added in triton 3.3+. No single triton version
-# satisfies both simultaneously. We run on triton 3.5.1 (latest, has both mamba3
-# APIs) and shim AttrsDescriptor as a stub dataclass for torch._inductor. The
-# stub is never actually invoked at runtime because the codebase does not use
-# torch.compile — but importing torch._inductor.* still requires the symbol to
-# exist at module load time.
-import triton as _triton  # noqa: E402
-if not hasattr(_triton, "set_allocator"):
-    def _noop_set_allocator(_fn):  # pragma: no cover
-        return None
-    _triton.set_allocator = _noop_set_allocator
-
-import triton.compiler.compiler as _tcc  # noqa: E402
-if not hasattr(_tcc, "AttrsDescriptor"):
-    class _AttrsDescriptorShim:
-        """Stub for torch._inductor compatibility on triton >= 3.4.
-        torch._inductor.runtime.hints imports this at module load but the
-        constructor is only called inside torch.compile paths. Accept any
-        args/kwargs so the import itself succeeds."""
-        def __init__(self, *args, **kwargs):
-            self.args = args
-            self.kwargs = kwargs
-
-        @classmethod
-        def from_hints(cls, *args, **kwargs):
-            return cls(*args, **kwargs)
-
-    _tcc.AttrsDescriptor = _AttrsDescriptorShim
-
-# triton_key: removed in triton 3.5, used by torch._inductor.codecache for
-# FxGraphCache key derivation. Return a stable string so caching still works.
-if not hasattr(_tcc, "triton_key"):
-    def _triton_key_shim():
-        import triton as _t
-        return f"triton-{_t.__version__}-shim"
-    _tcc.triton_key = _triton_key_shim
-
-# Suppress torch.compile/_dynamo errors globally — we don't rely on torch.compile
-# for performance in this codebase (Muon + mamba3 CUDA kernels already fused),
-# so fall back to eager on any dynamo failure rather than crashing. This is
-# defense-in-depth against further triton API drift.
-try:
-    import torch._dynamo  # noqa: F401 — triggers dynamo module init
-    torch._dynamo.config.suppress_errors = True
-except Exception:  # pragma: no cover
-    pass
-
-# Expose Mamba3 at top level to match `from mamba_ssm import Mamba3`.
-from mamba_ssm.modules.mamba3 import Mamba3  # noqa: E402
+# mamba_ssm package init — minimal override to avoid broken selective_scan_cuda.so
+# ABI mismatch with the base image's libtorch.
+#
+# The upstream __init__.py eagerly imports selective_scan_cuda which fails on
+# pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel (undefined c10::Warning ctor
+# symbol). We only need Mamba3 (grafted from main, pure-Triton), so we skip
+# all compiled-CUDA imports here and let Mamba3 load directly.
+
+__version__ = "2.3.1+feather-graft"
+
+# selective_scan_fn / mamba_inner_fn are shimmed to None — they are NOT used
+# by the Feather training path (which is Mamba3-only). If any import path
+# hits this, it will get a clear AttributeError instead of an obscure ImportError.
+selective_scan_fn = None
+mamba_inner_fn = None
+
+# --- triton API compatibility shims -----------------------------------------
+# Version matrix is hostile: torch 2.6 pins triton==3.2.0 because torch._inductor
+# imports AttrsDescriptor from triton.compiler.compiler — removed in triton 3.4+.
+# Grafted Mamba3 (from mamba-ssm main) needs triton.set_allocator and
+# tl.make_tensor_descriptor, both added in triton 3.3+. No single triton version
+# satisfies both simultaneously. We run on triton 3.5.1 (latest, has both mamba3
+# APIs) and shim AttrsDescriptor as a stub dataclass for torch._inductor. The
+# stub is never actually invoked at runtime because the codebase does not use
+# torch.compile — but importing torch._inductor.* still requires the symbol to
+# exist at module load time.
+import triton as _triton  # noqa: E402
+if not hasattr(_triton, "set_allocator"):
+    def _noop_set_allocator(_fn):  # pragma: no cover
+        return None
+    _triton.set_allocator = _noop_set_allocator
+
+import triton.compiler.compiler as _tcc  # noqa: E402
+if not hasattr(_tcc, "AttrsDescriptor"):
+    class _AttrsDescriptorShim:
+        """Stub for torch._inductor compatibility on triton >= 3.4.
+        torch._inductor.runtime.hints imports this at module load but the
+        constructor is only called inside torch.compile paths. Accept any
+        args/kwargs so the import itself succeeds."""
+        def __init__(self, *args, **kwargs):
+            self.args = args
+            self.kwargs = kwargs
+
+        @classmethod
+        def from_hints(cls, *args, **kwargs):
+            return cls(*args, **kwargs)
+
+    _tcc.AttrsDescriptor = _AttrsDescriptorShim
+
+# triton_key: removed in triton 3.5, used by torch._inductor.codecache for
+# FxGraphCache key derivation. Return a stable string so caching still works.
+if not hasattr(_tcc, "triton_key"):
+    def _triton_key_shim():
+        import triton as _t
+        return f"triton-{_t.__version__}-shim"
+    _tcc.triton_key = _triton_key_shim
+
+# Suppress torch.compile/_dynamo errors globally — we don't rely on torch.compile
+# for performance in this codebase (Muon + mamba3 CUDA kernels already fused),
+# so fall back to eager on any dynamo failure rather than crashing. This is
+# defense-in-depth against further triton API drift.
+try:
+    import torch._dynamo  # noqa: F401 — triggers dynamo module init
+    torch._dynamo.config.suppress_errors = True
+except Exception:  # pragma: no cover
+    pass
+
+# Expose Mamba3 at top level to match `from mamba_ssm import Mamba3`.
+from mamba_ssm.modules.mamba3 import Mamba3  # noqa: E402

overlay/.dockerignore

@@ -1,20 +1,20 @@
-.git
-.github
-.venv
-.remember
-.letta
-.claude
-__pycache__
-*.pyc
-*.pyo
-*.pyd
-*.log
-run_*.log
-run*.log
-*.txt
-WORKER_COMPLETE
-autoresearch_loop.log
-data/
-state_store/
-htm_rust/target/
-hydra-core/target/
+.git
+.github
+.venv
+.remember
+.letta
+.claude
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+*.log
+run_*.log
+run*.log
+*.txt
+WORKER_COMPLETE
+autoresearch_loop.log
+data/
+state_store/
+htm_rust/target/
+hydra-core/target/

overlay/configs/__init__.py

@@ -1,5 +1,5 @@
-from configs.hardware_config import HardwareConfig
-from configs.harness_config import HarnessConfig
-from configs.model_config import PostSemClawConfig
-
-__all__ = ["PostSemClawConfig", "HarnessConfig", "HardwareConfig"]
+from configs.hardware_config import HardwareConfig
+from configs.harness_config import HarnessConfig
+from configs.model_config import PostSemClawConfig
+
+__all__ = ["PostSemClawConfig", "HarnessConfig", "HardwareConfig"]

overlay/configs/hardware_config.py

@@ -1,104 +1,104 @@
-"""Hardware detection and memory budget configuration."""
-from __future__ import annotations
-
-import torch
-from pydantic import BaseModel, Field
-
-
-class HardwareConfig(BaseModel):
-    """Auto-detected hardware configuration with memory budgets."""
-
-    gpu_name: str = Field(default="unknown", description="GPU device name")
-    gpu_memory_mb: int = Field(default=0, description="Total GPU memory in MB")
-    gpu_vram_mb: int = Field(default=0, description="Alias for gpu_memory_mb (legacy compat)")
-    compute_capability: tuple[int, int] = Field(
-        default=(0, 0), description="CUDA compute capability"
-    )
-    peak_flops: float = Field(
-        default=12.74e12, description="Peak FP32 FLOPS for MFU calculation"
-    )
-    bf16_peak_flops: float = Field(
-        default=38.1e12, description="Peak BF16 FLOPS (RTX 3060 default)"
-    )
-
-    # Memory budget
-    model_budget_mb: int = Field(
-        default=1500, description="Max MB for model params + optimizer"
-    )
-    activation_budget_mb: int = Field(
-        default=3000, description="Max MB for activations"
-    )
-    overhead_mb: int = Field(
-        default=500, description="Reserved for CUDA context + PyTorch overhead"
-    )
-    max_vram_usage_pct: float = Field(
-        default=90.0, description="Max VRAM usage as % of total"
-    )
-    gradient_checkpointing: bool = Field(
-        default=False, description="Enable gradient checkpointing to save VRAM"
-    )
-
-    @classmethod
-    def detect(cls) -> HardwareConfig:
-        """Auto-detect hardware from current CUDA device."""
-        if not torch.cuda.is_available():
-            return cls()
-
-        device = torch.cuda.current_device()
-        props = torch.cuda.get_device_properties(device)
-        cap = (props.major, props.minor)
-        mem_mb = props.total_memory // (1024 * 1024)
-        gpu_name = props.name
-
-        # Peak FP32 FLOPS lookup by compute capability (approximate)
-        fp32_flops_table: dict[tuple[int, int], float] = {
-            (8, 6): 12.74e12,  # RTX 3060
-            (8, 9): 40.09e12,  # RTX 4090
-            (9, 0): 989.5e12,  # H100 (BF16)
-        }
-        peak = fp32_flops_table.get(cap, 12.74e12)
-
-        # BF16 peak FLOPS lookup by GPU name substring
-        bf16_flops_table: dict[str, float] = {
-            "3060": 38.1e12,
-            "3090": 71.0e12,
-            "4090": 165.2e12,
-            "A100": 312e12,
-            "H100": 989.5e12,
-            "A10G": 70.0e12,
-        }
-        bf16_peak = 38.1e12  # default to RTX 3060
-        for key, val in bf16_flops_table.items():
-            if key in gpu_name:
-                bf16_peak = val
-                break
-
-        # Memory budget: leave overhead_mb for CUDA context
-        overhead = 500
-        available = mem_mb - overhead
-        model_budget = int(available * 0.3)      # 30% for params + optimizer
-        activation_budget = int(available * 0.7)  # 70% for activations
-
-        return cls(
-            gpu_name=gpu_name,
-            gpu_memory_mb=mem_mb,
-            gpu_vram_mb=mem_mb,
-            compute_capability=cap,
-            peak_flops=peak,
-            bf16_peak_flops=bf16_peak,
-            model_budget_mb=model_budget,
-            activation_budget_mb=activation_budget,
-        )
-
-    def suggest_batch_size(self, d_model: int, seq_len: int, n_layer: int) -> int:
-        """Suggest batch size based on activation budget.
-
-        Uses rough estimate: per-sample activation ~= n_layer * seq_len * d_model
-        * 4 bytes * 2 (fwd + bwd).
-        """
-        per_sample_mb = n_layer * seq_len * d_model * 4 * 2 / (1024 * 1024)
-        if per_sample_mb <= 0:
-            return 1
-        batch = max(1, int(self.activation_budget_mb / per_sample_mb))
-        # Round down to power of 2
-        return 2 ** (batch.bit_length() - 1) if batch > 1 else 1
+"""Hardware detection and memory budget configuration."""
+from __future__ import annotations
+
+import torch
+from pydantic import BaseModel, Field
+
+
+class HardwareConfig(BaseModel):
+    """Auto-detected hardware configuration with memory budgets."""
+
+    gpu_name: str = Field(default="unknown", description="GPU device name")
+    gpu_memory_mb: int = Field(default=0, description="Total GPU memory in MB")
+    gpu_vram_mb: int = Field(default=0, description="Alias for gpu_memory_mb (legacy compat)")
+    compute_capability: tuple[int, int] = Field(
+        default=(0, 0), description="CUDA compute capability"
+    )
+    peak_flops: float = Field(
+        default=12.74e12, description="Peak FP32 FLOPS for MFU calculation"
+    )
+    bf16_peak_flops: float = Field(
+        default=38.1e12, description="Peak BF16 FLOPS (RTX 3060 default)"
+    )
+
+    # Memory budget
+    model_budget_mb: int = Field(
+        default=1500, description="Max MB for model params + optimizer"
+    )
+    activation_budget_mb: int = Field(
+        default=3000, description="Max MB for activations"
+    )
+    overhead_mb: int = Field(
+        default=500, description="Reserved for CUDA context + PyTorch overhead"
+    )
+    max_vram_usage_pct: float = Field(
+        default=90.0, description="Max VRAM usage as % of total"
+    )
+    gradient_checkpointing: bool = Field(
+        default=False, description="Enable gradient checkpointing to save VRAM"
+    )
+
+    @classmethod
+    def detect(cls) -> HardwareConfig:
+        """Auto-detect hardware from current CUDA device."""
+        if not torch.cuda.is_available():
+            return cls()
+
+        device = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(device)
+        cap = (props.major, props.minor)
+        mem_mb = props.total_memory // (1024 * 1024)
+        gpu_name = props.name
+
+        # Peak FP32 FLOPS lookup by compute capability (approximate)
+        fp32_flops_table: dict[tuple[int, int], float] = {
+            (8, 6): 12.74e12,  # RTX 3060
+            (8, 9): 40.09e12,  # RTX 4090
+            (9, 0): 989.5e12,  # H100 (BF16)
+        }
+        peak = fp32_flops_table.get(cap, 12.74e12)
+
+        # BF16 peak FLOPS lookup by GPU name substring
+        bf16_flops_table: dict[str, float] = {
+            "3060": 38.1e12,
+            "3090": 71.0e12,
+            "4090": 165.2e12,
+            "A100": 312e12,
+            "H100": 989.5e12,
+            "A10G": 70.0e12,
+        }
+        bf16_peak = 38.1e12  # default to RTX 3060
+        for key, val in bf16_flops_table.items():
+            if key in gpu_name:
+                bf16_peak = val
+                break
+
+        # Memory budget: leave overhead_mb for CUDA context
+        overhead = 500
+        available = mem_mb - overhead
+        model_budget = int(available * 0.3)      # 30% for params + optimizer
+        activation_budget = int(available * 0.7)  # 70% for activations
+
+        return cls(
+            gpu_name=gpu_name,
+            gpu_memory_mb=mem_mb,
+            gpu_vram_mb=mem_mb,
+            compute_capability=cap,
+            peak_flops=peak,
+            bf16_peak_flops=bf16_peak,
+            model_budget_mb=model_budget,
+            activation_budget_mb=activation_budget,
+        )
+
+    def suggest_batch_size(self, d_model: int, seq_len: int, n_layer: int) -> int:
+        """Suggest batch size based on activation budget.
+
+        Uses rough estimate: per-sample activation ~= n_layer * seq_len * d_model
+        * 4 bytes * 2 (fwd + bwd).
+        """
+        per_sample_mb = n_layer * seq_len * d_model * 4 * 2 / (1024 * 1024)
+        if per_sample_mb <= 0:
+            return 1
+        batch = max(1, int(self.activation_budget_mb / per_sample_mb))
+        # Round down to power of 2
+        return 2 ** (batch.bit_length() - 1) if batch > 1 else 1

overlay/configs/harness_config.py

@@ -1,78 +1,78 @@
-"""Harness configuration for HYDRA's self-evolving outer loop."""
-from typing import Literal
-
-from pydantic import BaseModel, Field
-
-
-class HarnessConfig(BaseModel):
-    """Configuration for the HYDRA harness behavior."""
-
-    # Inner loop
-    time_budget_seconds: int = Field(
-        default=300, ge=60, description="Training time budget per experiment in seconds"
-    )
-    max_experiments: int = Field(
-        default=1000, ge=0, description="Max experiments before stopping (0=infinite)"
-    )
-
-    # Meta-agent
-    meta_interval: int = Field(
-        default=20, ge=5, description="Run meta-agent every N experiments"
-    )
-    max_meta_changes: int = Field(
-        default=3, ge=1, le=10, description="Max changes per meta-iteration"
-    )
-
-    # Search strategy
-    exploration_mode: Literal["conservative", "balanced", "bold"] = "balanced"
-    exploration_budget: int = Field(
-        default=5, ge=1, description="Consecutive bold experiments when stuck"
-    )
-    stuck_threshold: int = Field(
-        default=10, ge=3, description="No improvement for N experiments = stuck"
-    )
-    crash_threshold: float = Field(
-        default=0.5,
-        ge=0.1,
-        le=1.0,
-        description="Crash rate threshold for BROKEN state",
-    )
-    regression_tolerance: float = Field(
-        default=0.05,
-        ge=0,
-        le=0.2,
-        description="Max val_bpb regression from best (fraction)",
-    )
-    max_regression_pct: float = Field(
-        default=5.0, description="Max % regression from best known val_bpb"
-    )
-
-    # Keep/discard criteria
-    primary_metric: str = "val_bpb"
-    secondary_metrics: dict = Field(
-        default_factory=lambda: {
-            "mhc_spectral_norm": {"max": 2.0},
-            "engram_hit_rate": {"min": 0.1},
-            "hestia_quant_error": {"max": 0.05},
-        }
-    )
-
-    # Experiment execution
-    experiment_timeout: int = Field(
-        default=600, ge=300, description="Kill experiment after N seconds"
-    )
-    warmup_steps: int = Field(
-        default=10, ge=0, description="Steps to exclude from timing"
-    )
-
-    # Git
-    branch_prefix: str = Field(default="autoresearch", description="Branch naming prefix")
-    results_file: str = Field(default="results.tsv", description="Experiment log file")
-
-    # Secondary metric gates (optional keep/discard criteria)
-    gate_mhc_spectral_norm: float | None = Field(
-        default=None, description="Max mhc_spectral_norm for keep (None=disabled)"
-    )
-    gate_engram_hit_rate: float | None = Field(
-        default=None, description="Min engram_hit_rate for keep (None=disabled)"
-    )
+"""Harness configuration for HYDRA's self-evolving outer loop."""
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class HarnessConfig(BaseModel):
+    """Configuration for the HYDRA harness behavior."""
+
+    # Inner loop
+    time_budget_seconds: int = Field(
+        default=300, ge=60, description="Training time budget per experiment in seconds"
+    )
+    max_experiments: int = Field(
+        default=1000, ge=0, description="Max experiments before stopping (0=infinite)"
+    )
+
+    # Meta-agent
+    meta_interval: int = Field(
+        default=20, ge=5, description="Run meta-agent every N experiments"
+    )
+    max_meta_changes: int = Field(
+        default=3, ge=1, le=10, description="Max changes per meta-iteration"
+    )
+
+    # Search strategy
+    exploration_mode: Literal["conservative", "balanced", "bold"] = "balanced"
+    exploration_budget: int = Field(
+        default=5, ge=1, description="Consecutive bold experiments when stuck"
+    )
+    stuck_threshold: int = Field(
+        default=10, ge=3, description="No improvement for N experiments = stuck"
+    )
+    crash_threshold: float = Field(
+        default=0.5,
+        ge=0.1,
+        le=1.0,
+        description="Crash rate threshold for BROKEN state",
+    )
+    regression_tolerance: float = Field(
+        default=0.05,
+        ge=0,
+        le=0.2,
+        description="Max val_bpb regression from best (fraction)",
+    )
+    max_regression_pct: float = Field(
+        default=5.0, description="Max % regression from best known val_bpb"
+    )
+
+    # Keep/discard criteria
+    primary_metric: str = "val_bpb"
+    secondary_metrics: dict = Field(
+        default_factory=lambda: {
+            "mhc_spectral_norm": {"max": 2.0},
+            "engram_hit_rate": {"min": 0.1},
+            "hestia_quant_error": {"max": 0.05},
+        }
+    )
+
+    # Experiment execution
+    experiment_timeout: int = Field(
+        default=600, ge=300, description="Kill experiment after N seconds"
+    )
+    warmup_steps: int = Field(
+        default=10, ge=0, description="Steps to exclude from timing"
+    )
+
+    # Git
+    branch_prefix: str = Field(default="autoresearch", description="Branch naming prefix")
+    results_file: str = Field(default="results.tsv", description="Experiment log file")
+
+    # Secondary metric gates (optional keep/discard criteria)
+    gate_mhc_spectral_norm: float | None = Field(
+        default=None, description="Max mhc_spectral_norm for keep (None=disabled)"
+    )
+    gate_engram_hit_rate: float | None = Field(
+        default=None, description="Min engram_hit_rate for keep (None=disabled)"
+    )

overlay/configs/model_config.py

@@ -1,80 +1,80 @@
-"""Post-SEM-Claw model configuration with Pydantic validation."""
-from pydantic import BaseModel, Field, field_validator
-
-
-class PostSemClawConfig(BaseModel):
-    """Configuration for the Post-SEM-Claw architecture.
-
-    Default values mirror the @dataclass in train.py exactly.
-    train.py is the source of truth — this file must stay in sync with it.
-    """
-
-    # Sequence
-    sequence_len: int = Field(default=2048, description="Context length (from prepare.py MAX_SEQ_LEN)")
-    vocab_size: int = Field(default=8192, description="Vocabulary size (from prepare.py VOCAB_SIZE)")
-
-    # Mamba-3 SSM
-    n_layer: int = Field(default=4, ge=1, le=48, description="Number of Mamba-3 blocks")
-    d_model: int = Field(default=256, ge=64, description="Model embedding dimension")
-    d_state: int = Field(default=64, ge=16, description="SSM state dimension")
-    headdim: int = Field(default=32, ge=16, description="SSM head dimension")
-    n_heads: int = Field(default=8, ge=1, description="Number of SSM heads (d_model // headdim)")
-    expand: int = Field(default=2, ge=1, le=4, description="Inner dim multiplier (inner_dim = expand * d_model)")
-
-    # mHC (Manifold Hyper-Connection)
-    mhc_n_streams: int = Field(default=4, ge=2, le=8, description="Number of residual streams")
-    mhc_sinkhorn_iters: int = Field(default=5, ge=1, le=100, description="Sinkhorn-Knopp iterations")
-
-    # Engram (conditional memory)
-    engram_n_columns: int = Field(default=4096, ge=256, description="Hash table columns")
-    engram_key_dim: int = Field(default=64, ge=16, description="Engram key dimension")
-    engram_layer_idx: int = Field(default=1, ge=0, description="Which layer gets engram (0-indexed)")
-
-    # Hestia QAT (disabled Phase 1, skeleton only)
-    hestia_enabled: bool = Field(default=False, description="Enable Hestia quantization")
-    hestia_bits: float = Field(default=1.58, gt=0, description="Target quantization bits (1.58 = 1.58-bit ternary)")
-
-    # SDR (bypass-only in Phase 1)
-    sdr_enabled: bool = Field(default=False, description="Enable stochastic resonance")
-    sdr_k: int = Field(default=64, ge=1, description="Top-K sparsification")
-    sdr_noise_std: float = Field(default=0.1, ge=0.0, description="SR noise standard deviation")
-
-    @field_validator("n_heads")
-    @classmethod
-    def validate_heads(cls, v: int, info: "FieldValidationInfo") -> int:
-        """Ensure n_heads equals d_model // headdim."""
-        d_model = info.data.get("d_model", 256)
-        headdim = info.data.get("headdim", 32)
-        expected = d_model // headdim
-        if v != expected:
-            raise ValueError(
-                f"n_heads ({v}) must equal d_model // headdim ({expected})"
-            )
-        return v
-
-    def estimate_params(self) -> int:
-        """Rough parameter count estimate based on train.py architecture."""
-        inner = self.expand * self.d_model
-        # in_proj: d_model -> inner + inner + d_state + d_state + n_heads
-        in_proj = self.d_model * (inner + inner + self.d_state + self.d_state + self.n_heads)
-        out_proj = inner * self.d_model
-        # conv1d (kernel=4, groups=inner_dim)
-        conv = inner * 4
-        # A_log, lambda_theta, D: n_heads each (3 vectors)
-        ssm_params = self.n_heads * 3
-        # bc_norm: d_state * 2 (weight + bias)
-        bc_norm = self.d_state * 2
-        per_block = in_proj + out_proj + conv + ssm_params + bc_norm
-        blocks = per_block * self.n_layer
-
-        # Embedding + lm_head (tied or untied)
-        embed = self.vocab_size * self.d_model * 2
-
-        # Engram: one instance at engram_layer_idx
-        # columns * d_model keys + d_model * engram_key_dim projection
-        engram = self.engram_n_columns * self.d_model + self.d_model * self.engram_key_dim
-
-        # mHC mixing matrices: n_layer * mhc_n_streams^2
-        mhc = self.n_layer * self.mhc_n_streams ** 2
-
-        return embed + blocks + engram + mhc
+"""Post-SEM-Claw model configuration with Pydantic validation."""
+from pydantic import BaseModel, Field, field_validator
+
+
+class PostSemClawConfig(BaseModel):
+    """Configuration for the Post-SEM-Claw architecture.
+
+    Default values mirror the @dataclass in train.py exactly.
+    train.py is the source of truth — this file must stay in sync with it.
+    """
+
+    # Sequence
+    sequence_len: int = Field(default=2048, description="Context length (from prepare.py MAX_SEQ_LEN)")
+    vocab_size: int = Field(default=8192, description="Vocabulary size (from prepare.py VOCAB_SIZE)")
+
+    # Mamba-3 SSM
+    n_layer: int = Field(default=4, ge=1, le=48, description="Number of Mamba-3 blocks")
+    d_model: int = Field(default=256, ge=64, description="Model embedding dimension")
+    d_state: int = Field(default=64, ge=16, description="SSM state dimension")
+    headdim: int = Field(default=32, ge=16, description="SSM head dimension")
+    n_heads: int = Field(default=8, ge=1, description="Number of SSM heads (d_model // headdim)")
+    expand: int = Field(default=2, ge=1, le=4, description="Inner dim multiplier (inner_dim = expand * d_model)")
+
+    # mHC (Manifold Hyper-Connection)
+    mhc_n_streams: int = Field(default=4, ge=2, le=8, description="Number of residual streams")
+    mhc_sinkhorn_iters: int = Field(default=5, ge=1, le=100, description="Sinkhorn-Knopp iterations")
+
+    # Engram (conditional memory)
+    engram_n_columns: int = Field(default=4096, ge=256, description="Hash table columns")
+    engram_key_dim: int = Field(default=64, ge=16, description="Engram key dimension")
+    engram_layer_idx: int = Field(default=1, ge=0, description="Which layer gets engram (0-indexed)")
+
+    # Hestia QAT (disabled Phase 1, skeleton only)
+    hestia_enabled: bool = Field(default=False, description="Enable Hestia quantization")
+    hestia_bits: float = Field(default=1.58, gt=0, description="Target quantization bits (1.58 = 1.58-bit ternary)")
+
+    # SDR (bypass-only in Phase 1)
+    sdr_enabled: bool = Field(default=False, description="Enable stochastic resonance")
+    sdr_k: int = Field(default=64, ge=1, description="Top-K sparsification")
+    sdr_noise_std: float = Field(default=0.1, ge=0.0, description="SR noise standard deviation")
+
+    @field_validator("n_heads")
+    @classmethod
+    def validate_heads(cls, v: int, info: "FieldValidationInfo") -> int:
+        """Ensure n_heads equals d_model // headdim."""
+        d_model = info.data.get("d_model", 256)
+        headdim = info.data.get("headdim", 32)
+        expected = d_model // headdim
+        if v != expected:
+            raise ValueError(
+                f"n_heads ({v}) must equal d_model // headdim ({expected})"
+            )
+        return v
+
+    def estimate_params(self) -> int:
+        """Rough parameter count estimate based on train.py architecture."""
+        inner = self.expand * self.d_model
+        # in_proj: d_model -> inner + inner + d_state + d_state + n_heads
+        in_proj = self.d_model * (inner + inner + self.d_state + self.d_state + self.n_heads)
+        out_proj = inner * self.d_model
+        # conv1d (kernel=4, groups=inner_dim)
+        conv = inner * 4
+        # A_log, lambda_theta, D: n_heads each (3 vectors)
+        ssm_params = self.n_heads * 3
+        # bc_norm: d_state * 2 (weight + bias)
+        bc_norm = self.d_state * 2
+        per_block = in_proj + out_proj + conv + ssm_params + bc_norm
+        blocks = per_block * self.n_layer
+
+        # Embedding + lm_head (tied or untied)
+        embed = self.vocab_size * self.d_model * 2
+
+        # Engram: one instance at engram_layer_idx
+        # columns * d_model keys + d_model * engram_key_dim projection
+        engram = self.engram_n_columns * self.d_model + self.d_model * self.engram_key_dim
+
+        # mHC mixing matrices: n_layer * mhc_n_streams^2
+        mhc = self.n_layer * self.mhc_n_streams ** 2
+
+        return embed + blocks + engram + mhc

overlay/htm_rust/Cargo.lock

@@ -1,383 +1,383 @@
-# This file is automatically @generated by Cargo.
-# It is not intended for manual editing.
-version = 4
-
-[[package]]
-name = "autocfg"
-version = "1.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
-
-[[package]]
-name = "cfg-if"
-version = "1.0.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
-
-[[package]]
-name = "cudarc"
-version = "0.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38cd60a9a42ec83a2ed7effb0b1f073270264ea99da7acfc44f7e8d74dee0384"
-dependencies = [
- "libloading",
-]
-
-[[package]]
-name = "getrandom"
-version = "0.2.17"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
-dependencies = [
- "cfg-if",
- "libc",
- "wasi",
-]
-
-[[package]]
-name = "heck"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
-
-[[package]]
-name = "htm_rust"
-version = "0.1.0"
-dependencies = [
- "cudarc",
- "ndarray",
- "numpy",
- "pyo3",
- "rand",
- "rand_xoshiro",
-]
-
-[[package]]
-name = "indoc"
-version = "2.0.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
-dependencies = [
- "rustversion",
-]
-
-[[package]]
-name = "libc"
-version = "0.2.185"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f"
-
-[[package]]
-name = "libloading"
-version = "0.8.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
-dependencies = [
- "cfg-if",
- "windows-link",
-]
-
-[[package]]
-name = "matrixmultiply"
-version = "0.3.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
-dependencies = [
- "autocfg",
- "rawpointer",
-]
-
-[[package]]
-name = "memoffset"
-version = "0.9.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
-dependencies = [
- "autocfg",
-]
-
-[[package]]
-name = "ndarray"
-version = "0.16.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
-dependencies = [
- "matrixmultiply",
- "num-complex",
- "num-integer",
- "num-traits",
- "portable-atomic",
- "portable-atomic-util",
- "rawpointer",
-]
-
-[[package]]
-name = "num-complex"
-version = "0.4.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
-dependencies = [
- "num-traits",
-]
-
-[[package]]
-name = "num-integer"
-version = "0.1.46"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
-dependencies = [
- "num-traits",
-]
-
-[[package]]
-name = "num-traits"
-version = "0.2.19"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
-dependencies = [
- "autocfg",
-]
-
-[[package]]
-name = "numpy"
-version = "0.22.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "edb929bc0da91a4d85ed6c0a84deaa53d411abfb387fc271124f91bf6b89f14e"
-dependencies = [
- "libc",
- "ndarray",
- "num-complex",
- "num-integer",
- "num-traits",
- "pyo3",
- "rustc-hash",
-]
-
-[[package]]
-name = "once_cell"
-version = "1.21.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
-
-[[package]]
-name = "portable-atomic"
-version = "1.13.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
-
-[[package]]
-name = "portable-atomic-util"
-version = "0.2.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3"
-dependencies = [
- "portable-atomic",
-]
-
-[[package]]
-name = "ppv-lite86"
-version = "0.2.21"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
-dependencies = [
- "zerocopy",
-]
-
-[[package]]
-name = "proc-macro2"
-version = "1.0.106"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
-dependencies = [
- "unicode-ident",
-]
-
-[[package]]
-name = "pyo3"
-version = "0.22.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
-dependencies = [
- "cfg-if",
- "indoc",
- "libc",
- "memoffset",
- "once_cell",
- "portable-atomic",
- "pyo3-build-config",
- "pyo3-ffi",
- "pyo3-macros",
- "unindent",
-]
-
-[[package]]
-name = "pyo3-build-config"
-version = "0.22.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
-dependencies = [
- "once_cell",
- "target-lexicon",
-]
-
-[[package]]
-name = "pyo3-ffi"
-version = "0.22.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
-dependencies = [
- "libc",
- "pyo3-build-config",
-]
-
-[[package]]
-name = "pyo3-macros"
-version = "0.22.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
-dependencies = [
- "proc-macro2",
- "pyo3-macros-backend",
- "quote",
- "syn",
-]
-
-[[package]]
-name = "pyo3-macros-backend"
-version = "0.22.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
-dependencies = [
- "heck",
- "proc-macro2",
- "pyo3-build-config",
- "quote",
- "syn",
-]
-
-[[package]]
-name = "quote"
-version = "1.0.45"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
-dependencies = [
- "proc-macro2",
-]
-
-[[package]]
-name = "rand"
-version = "0.8.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
-dependencies = [
- "libc",
- "rand_chacha",
- "rand_core",
-]
-
-[[package]]
-name = "rand_chacha"
-version = "0.3.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
-dependencies = [
- "ppv-lite86",
- "rand_core",
-]
-
-[[package]]
-name = "rand_core"
-version = "0.6.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
-dependencies = [
- "getrandom",
-]
-
-[[package]]
-name = "rand_xoshiro"
-version = "0.6.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
-dependencies = [
- "rand_core",
-]
-
-[[package]]
-name = "rawpointer"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
-
-[[package]]
-name = "rustc-hash"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
-
-[[package]]
-name = "rustversion"
-version = "1.0.22"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
-
-[[package]]
-name = "syn"
-version = "2.0.117"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
-dependencies = [
- "proc-macro2",
- "quote",
- "unicode-ident",
-]
-
-[[package]]
-name = "target-lexicon"
-version = "0.12.16"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
-
-[[package]]
-name = "unicode-ident"
-version = "1.0.24"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
-
-[[package]]
-name = "unindent"
-version = "0.2.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
-
-[[package]]
-name = "wasi"
-version = "0.11.1+wasi-snapshot-preview1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
-
-[[package]]
-name = "windows-link"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
-
-[[package]]
-name = "zerocopy"
-version = "0.8.48"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9"
-dependencies = [
- "zerocopy-derive",
-]
-
-[[package]]
-name = "zerocopy-derive"
-version = "0.8.48"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn",
-]
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
+[[package]]
+name = "cfg-if"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
+
+[[package]]
+name = "cudarc"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38cd60a9a42ec83a2ed7effb0b1f073270264ea99da7acfc44f7e8d74dee0384"
+dependencies = [
+ "libloading",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "htm_rust"
+version = "0.1.0"
+dependencies = [
+ "cudarc",
+ "ndarray",
+ "numpy",
+ "pyo3",
+ "rand",
+ "rand_xoshiro",
+]
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.185"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f"
+
+[[package]]
+name = "libloading"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
+dependencies = [
+ "cfg-if",
+ "windows-link",
+]
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "ndarray"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "numpy"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "edb929bc0da91a4d85ed6c0a84deaa53d411abfb387fc271124f91bf6b89f14e"
+dependencies = [
+ "libc",
+ "ndarray",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "pyo3",
+ "rustc-hash",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "portable-atomic"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "ppv-lite86"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
+dependencies = [
+ "zerocopy",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "rand"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom",
+]
+
+[[package]]
+name = "rand_xoshiro"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
+dependencies = [
+ "rand_core",
+]
+
+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
+[[package]]
+name = "rustc-hash"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
+
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
+[[package]]
+name = "syn"
+version = "2.0.117"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
+[[package]]
+name = "wasi"
+version = "0.11.1+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+
+[[package]]
+name = "windows-link"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
+
+[[package]]
+name = "zerocopy"
+version = "0.8.48"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.48"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]

overlay/htm_rust/Cargo.toml

@@ -1,37 +1,37 @@
-[package]
-name = "htm_rust"
-version = "0.1.0"
-edition = "2021"
-authors = ["Feather/HYDRA"]
-description = "Numenta BAMI-spec Hierarchical Temporal Memory (Spatial Pooler + Temporal Memory) with pyo3 bindings"
-license = "MIT"
-
-[lib]
-name = "htm_rust"
-crate-type = ["cdylib", "rlib"]
-
-[dependencies]
-pyo3 = { version = "0.22", features = ["extension-module"] }
-numpy = "0.22"
-ndarray = "0.16"
-rand = "0.8"
-rand_xoshiro = "0.6"
-# cudarc: CUDA Rust bindings with dynamic-loading (no link-time dep on libcuda).
-# Kernels are embedded as PTX and JIT-compiled at runtime.
-cudarc = { version = "0.12", default-features = false, features = ["dynamic-linking", "driver", "cuda-12010"], optional = true }
-
-[build-dependencies]
-# Only required when building with --features gpu. We shell to nvcc directly
-# so we don't need cc's cuda support (which drags in extra deps).
-
-[features]
-default = []
-# `gpu` adds the HTMRegionGPU class, compiles .cu kernels to PTX at build time,
-# and links cudarc. Without this feature the crate is pure-CPU and has no
-# CUDA dependency at build or run time.
-gpu = ["cudarc"]
-
-[profile.release]
-opt-level = 3
-lto = "thin"
-codegen-units = 1
+[package]
+name = "htm_rust"
+version = "0.1.0"
+edition = "2021"
+authors = ["Feather/HYDRA"]
+description = "Numenta BAMI-spec Hierarchical Temporal Memory (Spatial Pooler + Temporal Memory) with pyo3 bindings"
+license = "MIT"
+
+[lib]
+name = "htm_rust"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+pyo3 = { version = "0.22", features = ["extension-module"] }
+numpy = "0.22"
+ndarray = "0.16"
+rand = "0.8"
+rand_xoshiro = "0.6"
+# cudarc: CUDA Rust bindings with dynamic-loading (no link-time dep on libcuda).
+# Kernels are embedded as PTX and JIT-compiled at runtime.
+cudarc = { version = "0.12", default-features = false, features = ["dynamic-linking", "driver", "cuda-12010"], optional = true }
+
+[build-dependencies]
+# Only required when building with --features gpu. We shell to nvcc directly
+# so we don't need cc's cuda support (which drags in extra deps).
+
+[features]
+default = []
+# `gpu` adds the HTMRegionGPU class, compiles .cu kernels to PTX at build time,
+# and links cudarc. Without this feature the crate is pure-CPU and has no
+# CUDA dependency at build or run time.
+gpu = ["cudarc"]
+
+[profile.release]
+opt-level = 3
+lto = "thin"
+codegen-units = 1

overlay/htm_rust/bench_gpu.py

@@ -1,81 +1,81 @@
-"""Microbenchmark: CPU vs GPU HTMLayer forward at HYDRA training sizes.
-
-Usage:
-    source .venv/bin/activate
-    export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
-    python htm_rust/bench_gpu.py
-"""
-import os
-import sys
-import time
-
-# Ensure /home/mikeb/work/feather is on sys.path so `subsystems` imports.
-_FEATHER = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-if _FEATHER not in sys.path:
-    sys.path.insert(0, _FEATHER)
-
-import numpy as np
-import torch
-
-from subsystems.htm import HTMLayer
-
-
-def bench(layer: HTMLayer, sdr: torch.Tensor, warmup: int = 1, iters: int = 3) -> float:
-    """Return mean ms/forward."""
-    for _ in range(warmup):
-        _ = layer(sdr)
-    if torch.cuda.is_available():
-        torch.cuda.synchronize()
-    t0 = time.perf_counter()
-    for _ in range(iters):
-        _ = layer(sdr)
-    if torch.cuda.is_available():
-        torch.cuda.synchronize()
-    dt = time.perf_counter() - t0
-    return dt * 1000 / iters
-
-
-def main() -> None:
-    # HYDRA training config: B=8, T=2048, bits=16384, cols=2048.
-    B, T, D = int(os.environ.get("B", 8)), int(os.environ.get("T", 2048)), 16384
-    n_cols = 2048
-
-    print(f"config: B={B} T={T} D={D} n_cols={n_cols}")
-    print(f"torch: {torch.__version__} cuda={torch.cuda.is_available()}")
-
-    # Build a fixed sparse SDR once.
-    rng = np.random.default_rng(0)
-    sdr = np.zeros((B, T, D), dtype=bool)
-    on = int(D * 0.02)
-    for b in range(B):
-        for t in range(T):
-            idx = rng.choice(D, size=on, replace=False)
-            sdr[b, t, idx] = True
-    sdr_t = torch.from_numpy(sdr)
-
-    # CPU baseline.
-    print("\n--- CPU ---")
-    cpu_layer = HTMLayer(
-        input_bits=D, n_columns=n_cols, cells_per_column=32,
-        batch_size=B, seed=42, use_gpu=False,
-    )
-    cpu_layer.train()
-    cpu_ms = bench(cpu_layer, sdr_t, warmup=1, iters=2)
-    print(f"CPU: {cpu_ms:.1f} ms/forward  ({cpu_ms/T:.2f} ms/step × T={T})")
-
-    # GPU.
-    print("\n--- GPU ---")
-    gpu_layer = HTMLayer(
-        input_bits=D, n_columns=n_cols, cells_per_column=32,
-        batch_size=B, seed=42, use_gpu=True,
-    )
-    gpu_layer.train()
-    sdr_cuda = sdr_t.cuda()
-    gpu_ms = bench(gpu_layer, sdr_cuda, warmup=1, iters=2)
-    print(f"GPU: {gpu_ms:.1f} ms/forward  ({gpu_ms/T:.2f} ms/step × T={T})")
-
-    print(f"\nSpeedup: {cpu_ms / gpu_ms:.2f}x")
-
-
-if __name__ == "__main__":
-    main()
+"""Microbenchmark: CPU vs GPU HTMLayer forward at HYDRA training sizes.
+
+Usage:
+    source .venv/bin/activate
+    export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+    python htm_rust/bench_gpu.py
+"""
+import os
+import sys
+import time
+
+# Ensure /home/mikeb/work/feather is on sys.path so `subsystems` imports.
+_FEATHER = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _FEATHER not in sys.path:
+    sys.path.insert(0, _FEATHER)
+
+import numpy as np
+import torch
+
+from subsystems.htm import HTMLayer
+
+
+def bench(layer: HTMLayer, sdr: torch.Tensor, warmup: int = 1, iters: int = 3) -> float:
+    """Return mean ms/forward."""
+    for _ in range(warmup):
+        _ = layer(sdr)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        _ = layer(sdr)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    dt = time.perf_counter() - t0
+    return dt * 1000 / iters
+
+
+def main() -> None:
+    # HYDRA training config: B=8, T=2048, bits=16384, cols=2048.
+    B, T, D = int(os.environ.get("B", 8)), int(os.environ.get("T", 2048)), 16384
+    n_cols = 2048
+
+    print(f"config: B={B} T={T} D={D} n_cols={n_cols}")
+    print(f"torch: {torch.__version__} cuda={torch.cuda.is_available()}")
+
+    # Build a fixed sparse SDR once.
+    rng = np.random.default_rng(0)
+    sdr = np.zeros((B, T, D), dtype=bool)
+    on = int(D * 0.02)
+    for b in range(B):
+        for t in range(T):
+            idx = rng.choice(D, size=on, replace=False)
+            sdr[b, t, idx] = True
+    sdr_t = torch.from_numpy(sdr)
+
+    # CPU baseline.
+    print("\n--- CPU ---")
+    cpu_layer = HTMLayer(
+        input_bits=D, n_columns=n_cols, cells_per_column=32,
+        batch_size=B, seed=42, use_gpu=False,
+    )
+    cpu_layer.train()
+    cpu_ms = bench(cpu_layer, sdr_t, warmup=1, iters=2)
+    print(f"CPU: {cpu_ms:.1f} ms/forward  ({cpu_ms/T:.2f} ms/step × T={T})")
+
+    # GPU.
+    print("\n--- GPU ---")
+    gpu_layer = HTMLayer(
+        input_bits=D, n_columns=n_cols, cells_per_column=32,
+        batch_size=B, seed=42, use_gpu=True,
+    )
+    gpu_layer.train()
+    sdr_cuda = sdr_t.cuda()
+    gpu_ms = bench(gpu_layer, sdr_cuda, warmup=1, iters=2)
+    print(f"GPU: {gpu_ms:.1f} ms/forward  ({gpu_ms/T:.2f} ms/step × T={T})")
+
+    print(f"\nSpeedup: {cpu_ms / gpu_ms:.2f}x")
+
+
+if __name__ == "__main__":
+    main()

overlay/htm_rust/build.rs

@@ -1,162 +1,162 @@
-//! Build script: compiles `.cu` kernel files to PTX when the `gpu` feature
-//! is enabled. PTX files are embedded into the final Rust binary via
-//! `include_str!` / `OUT_DIR` constants and JIT-loaded at runtime by cudarc.
-//!
-//! No-op when `gpu` feature is off — CPU-only builds have zero CUDA
-//! toolchain dependency.
-//!
-//! nvcc lookup order:
-//!   1. $NVCC env var
-//!   2. `nvcc` on PATH
-//!   3. `/usr/local/cuda-12.1/bin/nvcc`
-//!   4. `/usr/local/cuda/bin/nvcc`
-//!
-//! Default target: sm_86 (Ampere A10G / RTX 30xx). Override with $HTM_CUDA_ARCH (e.g. sm_90a for H200).
-
-use std::env;
-use std::path::PathBuf;
-use std::process::Command;
-
-fn main() {
-    // Re-run whenever we edit the build script or any kernel source.
-    println!("cargo:rerun-if-changed=build.rs");
-
-    let gpu = env::var_os("CARGO_FEATURE_GPU").is_some();
-    if !gpu {
-        return;
-    }
-
-    // Kernels to compile. Each .cu file → one .ptx file, embedded by name.
-    let kernels: &[&str] = &[
-        "sp_overlap",
-        "sp_topk",
-        "sp_learn",
-        "sp_duty",
-        "sp_boost_fused",
-        "tm_predict",
-        "tm_activate",
-        "tm_learn",
-        "tm_punish",
-        "tm_grow",
-        "tm_anomaly",
-        "tm_reset",
-        "htm_fused_step",
-    ];
-
-    let kernels_dir = PathBuf::from("src/gpu/kernels");
-    for k in kernels {
-        let src = kernels_dir.join(format!("{k}.cu"));
-        println!("cargo:rerun-if-changed={}", src.display());
-    }
-
-    let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR"));
-    let arch = env::var("HTM_CUDA_ARCH").unwrap_or_else(|_| "sm_86".into());
-
-    let nvcc = find_nvcc();
-    println!("cargo:warning=htm_rust: nvcc = {nvcc}");
-    println!("cargo:warning=htm_rust: target arch = {arch}");
-
-    // Prefer gcc-12 if present (CUDA 12.1 doesn't support gcc-13+ headers).
-    let host_compiler = env::var("HTM_CUDA_CCBIN")
-        .ok()
-        .or_else(|| {
-            for cand in ["/usr/bin/gcc-12", "/usr/bin/gcc-11"] {
-                if std::path::Path::new(cand).exists() {
-                    return Some(cand.to_string());
-                }
-            }
-            None
-        });
-
-    // Optionally patch the emitted PTX `.version` header down to match an
-    // older driver. Useful when the system driver (e.g. on WSL2) is older
-    // than the nvcc toolchain. Set HTM_PTX_VERSION to e.g. "7.8" or "8.0".
-    let ptx_version_override = env::var("HTM_PTX_VERSION").ok();
-
-    for k in kernels {
-        let src = kernels_dir.join(format!("{k}.cu"));
-        let ptx = out_dir.join(format!("{k}.ptx"));
-        if !src.exists() {
-            panic!("missing kernel source: {}", src.display());
-        }
-        let mut cmd = Command::new(&nvcc);
-        // Note: `--use_fast_math` breaks bit-parity with host `expf`, which
-        // in turn flips boost tie-breaks in SP learning. We accept the tiny
-        // perf loss for correctness; the hot overlap kernel has no transcendentals.
-        cmd.args([
-            "--ptx",
-            "-O3",
-            "-rdc=true",
-            "-arch",
-            &arch,
-        ]);
-        // `cooperative_groups::this_cluster()` is not declared for Ampere
-        // device compiles in CUDA 12.x, even if guarded by __CUDA_ARCH__ in
-        // some nvcc front-end phases. Define an explicit build-time kill
-        // switch for all non-Hopper targets so sm_86/A10G only sees the
-        // cooperative-grid path.
-        if !arch.starts_with("sm_90") {
-            cmd.arg("-DHTM_DISABLE_CLUSTER=1");
-        }
-        if let Some(cc) = &host_compiler {
-            cmd.args(["-ccbin", cc]);
-        }
-        cmd.arg("-o").arg(&ptx).arg(&src);
-        let status = cmd
-            .status()
-            .unwrap_or_else(|e| panic!("failed to spawn nvcc: {e}"));
-        if !status.success() {
-            panic!("nvcc failed for {}", src.display());
-        }
-
-        if let Some(ver) = &ptx_version_override {
-            // Read, patch, write.
-            let text = std::fs::read_to_string(&ptx)
-                .unwrap_or_else(|e| panic!("read {} failed: {e}", ptx.display()));
-            // Match `.version X.Y` where X and Y are digits. Replace whole line.
-            let patched: String = text
-                .lines()
-                .map(|line| {
-                    let t = line.trim_start();
-                    if t.starts_with(".version ") {
-                        format!(".version {ver}")
-                    } else {
-                        line.to_string()
-                    }
-                })
-                .collect::<Vec<_>>()
-                .join("\n");
-            std::fs::write(&ptx, patched)
-                .unwrap_or_else(|e| panic!("write {} failed: {e}", ptx.display()));
-        }
-    }
-
-    // Export OUT_DIR for include_str! in Rust.
-    println!(
-        "cargo:rustc-env=HTM_GPU_PTX_DIR={}",
-        out_dir.display()
-    );
-}
-
-fn find_nvcc() -> String {
-    if let Ok(n) = env::var("NVCC") {
-        return n;
-    }
-    // Try PATH.
-    if Command::new("nvcc").arg("--version").output().is_ok() {
-        return "nvcc".into();
-    }
-    for cand in [
-        "/usr/local/cuda-12.1/bin/nvcc",
-        "/usr/local/cuda/bin/nvcc",
-        "/usr/local/cuda-12/bin/nvcc",
-    ] {
-        if std::path::Path::new(cand).exists() {
-            return cand.into();
-        }
-    }
-    panic!(
-        "nvcc not found. Set $NVCC or install CUDA toolkit. \
-         Tried PATH, /usr/local/cuda-12.1, /usr/local/cuda."
-    );
-}
+//! Build script: compiles `.cu` kernel files to PTX when the `gpu` feature
+//! is enabled. PTX files are embedded into the final Rust binary via
+//! `include_str!` / `OUT_DIR` constants and JIT-loaded at runtime by cudarc.
+//!
+//! No-op when `gpu` feature is off — CPU-only builds have zero CUDA
+//! toolchain dependency.
+//!
+//! nvcc lookup order:
+//!   1. $NVCC env var
+//!   2. `nvcc` on PATH
+//!   3. `/usr/local/cuda-12.1/bin/nvcc`
+//!   4. `/usr/local/cuda/bin/nvcc`
+//!
+//! Default target: sm_86 (Ampere A10G / RTX 30xx). Override with $HTM_CUDA_ARCH (e.g. sm_90a for H200).
+
+use std::env;
+use std::path::PathBuf;
+use std::process::Command;
+
+fn main() {
+    // Re-run whenever we edit the build script or any kernel source.
+    println!("cargo:rerun-if-changed=build.rs");
+
+    let gpu = env::var_os("CARGO_FEATURE_GPU").is_some();
+    if !gpu {
+        return;
+    }
+
+    // Kernels to compile. Each .cu file → one .ptx file, embedded by name.
+    let kernels: &[&str] = &[
+        "sp_overlap",
+        "sp_topk",
+        "sp_learn",
+        "sp_duty",
+        "sp_boost_fused",
+        "tm_predict",
+        "tm_activate",
+        "tm_learn",
+        "tm_punish",
+        "tm_grow",
+        "tm_anomaly",
+        "tm_reset",
+        "htm_fused_step",
+    ];
+
+    let kernels_dir = PathBuf::from("src/gpu/kernels");
+    for k in kernels {
+        let src = kernels_dir.join(format!("{k}.cu"));
+        println!("cargo:rerun-if-changed={}", src.display());
+    }
+
+    let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR"));
+    let arch = env::var("HTM_CUDA_ARCH").unwrap_or_else(|_| "sm_86".into());
+
+    let nvcc = find_nvcc();
+    println!("cargo:warning=htm_rust: nvcc = {nvcc}");
+    println!("cargo:warning=htm_rust: target arch = {arch}");
+
+    // Prefer gcc-12 if present (CUDA 12.1 doesn't support gcc-13+ headers).
+    let host_compiler = env::var("HTM_CUDA_CCBIN")
+        .ok()
+        .or_else(|| {
+            for cand in ["/usr/bin/gcc-12", "/usr/bin/gcc-11"] {
+                if std::path::Path::new(cand).exists() {
+                    return Some(cand.to_string());
+                }
+            }
+            None
+        });
+
+    // Optionally patch the emitted PTX `.version` header down to match an
+    // older driver. Useful when the system driver (e.g. on WSL2) is older
+    // than the nvcc toolchain. Set HTM_PTX_VERSION to e.g. "7.8" or "8.0".
+    let ptx_version_override = env::var("HTM_PTX_VERSION").ok();
+
+    for k in kernels {
+        let src = kernels_dir.join(format!("{k}.cu"));
+        let ptx = out_dir.join(format!("{k}.ptx"));
+        if !src.exists() {
+            panic!("missing kernel source: {}", src.display());
+        }
+        let mut cmd = Command::new(&nvcc);
+        // Note: `--use_fast_math` breaks bit-parity with host `expf`, which
+        // in turn flips boost tie-breaks in SP learning. We accept the tiny
+        // perf loss for correctness; the hot overlap kernel has no transcendentals.
+        cmd.args([
+            "--ptx",
+            "-O3",
+            "-rdc=true",
+            "-arch",
+            &arch,
+        ]);
+        // `cooperative_groups::this_cluster()` is not declared for Ampere
+        // device compiles in CUDA 12.x, even if guarded by __CUDA_ARCH__ in
+        // some nvcc front-end phases. Define an explicit build-time kill
+        // switch for all non-Hopper targets so sm_86/A10G only sees the
+        // cooperative-grid path.
+        if !arch.starts_with("sm_90") {
+            cmd.arg("-DHTM_DISABLE_CLUSTER=1");
+        }
+        if let Some(cc) = &host_compiler {
+            cmd.args(["-ccbin", cc]);
+        }
+        cmd.arg("-o").arg(&ptx).arg(&src);
+        let status = cmd
+            .status()
+            .unwrap_or_else(|e| panic!("failed to spawn nvcc: {e}"));
+        if !status.success() {
+            panic!("nvcc failed for {}", src.display());
+        }
+
+        if let Some(ver) = &ptx_version_override {
+            // Read, patch, write.
+            let text = std::fs::read_to_string(&ptx)
+                .unwrap_or_else(|e| panic!("read {} failed: {e}", ptx.display()));
+            // Match `.version X.Y` where X and Y are digits. Replace whole line.
+            let patched: String = text
+                .lines()
+                .map(|line| {
+                    let t = line.trim_start();
+                    if t.starts_with(".version ") {
+                        format!(".version {ver}")
+                    } else {
+                        line.to_string()
+                    }
+                })
+                .collect::<Vec<_>>()
+                .join("\n");
+            std::fs::write(&ptx, patched)
+                .unwrap_or_else(|e| panic!("write {} failed: {e}", ptx.display()));
+        }
+    }
+
+    // Export OUT_DIR for include_str! in Rust.
+    println!(
+        "cargo:rustc-env=HTM_GPU_PTX_DIR={}",
+        out_dir.display()
+    );
+}
+
+fn find_nvcc() -> String {
+    if let Ok(n) = env::var("NVCC") {
+        return n;
+    }
+    // Try PATH.
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return "nvcc".into();
+    }
+    for cand in [
+        "/usr/local/cuda-12.1/bin/nvcc",
+        "/usr/local/cuda/bin/nvcc",
+        "/usr/local/cuda-12/bin/nvcc",
+    ] {
+        if std::path::Path::new(cand).exists() {
+            return cand.into();
+        }
+    }
+    panic!(
+        "nvcc not found. Set $NVCC or install CUDA toolkit. \
+         Tried PATH, /usr/local/cuda-12.1, /usr/local/cuda."
+    );
+}

overlay/htm_rust/docs/GPU_HTM.md

@@ -1,302 +1,302 @@
-# GPU HTM Backend
-
-## Status
-
-**FUSED MEGAKERNEL: entire T-timestep SP+TM forward collapsed into a single
-CUDA launch per forward pass.**
-
-* Legacy path: 12 kernels × T=2048 timesteps = 24K launches per forward.
-* Fused path: **1 launch per forward** (24000× launch-overhead reduction).
-* End-to-end training throughput: **~2.7k → ~60k tok/sec** (~22x speedup).
-* Fused path uses per-column threshold inhibition instead of global top-K
-  (see §Fused Kernel below — this is a real architectural change).
-
-## Fused Kernel
-
-### Why
-
-Global top-K column selection requires cross-block synchronization at every
-timestep. On WSL2/sm_86 without `-rdc=true`, `cooperative_groups::grid_sync()`
-is unreliable. Without a grid sync, collapsing the T-loop into one kernel is
-impossible, so every forward pays 12×T kernel launches and 90%+ of runtime is
-CUDA launch overhead + small-kernel tails.
-
-### How
-
-Replace global top-K with **per-column threshold activation**:
-
-    is_active[c] = (overlap[c] * boost[c]) > inhibition_threshold[c]
-
-`inhibition_threshold[c]` is a per-column scalar, learned via EMA update:
-
-    err = active_duty[c] - sparsity_target
-    new_thr = clamp(thr + thr_adapt_rate * err * 100, 0.1, 1000)
-
-This is biologically grounded (GABAergic local lateral inhibition in
-neocortical columns) and supported by HTM theory. The duty-cycle-driven
-feedback loop was already present; we simply redirect its output to drive
-activation threshold instead of multiplicative boost. The global top-K,
-which had no biological basis, is removed.
-
-### Cross-block coherence
-
-- **Ping-pong bitsets** for `cell_active_bits` and `cell_winner_bits`: at
-  even t write to `_a`, read from `_b`; at odd t reversed. This eliminates
-  the need for an in-place snapshot kernel between timesteps.
-- **Primary path: cooperative launch + hardware grid sync**. Host code probes
-  `CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH`, computes the cooperative whole-grid
-  residency limit from occupancy, and launches the fused megakernel with
-  `cuLaunchCooperativeKernel`. In-kernel barriers use
-  `cooperative_groups::this_grid().sync()`.
-- **Fallback path: software grid barrier** via a 3-slot atomic counter array
-  (`barrier_counters`). This remains as a compatibility fallback when
-  cooperative launch is unavailable.
-- **Launch invariant**: cooperative launch is capped to the hardware residency
-  limit for `blockDim.x = 1024`; software fallback remains capped conservatively
-  (`HTM_FUSED_GRID_CAP`, default 8) to avoid whole-grid spin deadlock.
-
-### Kernel structure
-
-```
-for t in 0..T:
-    # Phase 0: clear curr_active/curr_winner for my column range
-    grid_barrier()
-    # Phase A: SP overlap → boost → threshold → SP learn → duty + threshold EMA
-    grid_barrier()
-    # Phase B: TM predict (per cell, per seg) → TM learn (reinforce on match)
-    #                   → burst if none predicted → segment grow/reinforce
-    grid_barrier()
-    # Phase C: block 0 writes anomaly[t]
-```
-
-Each warp owns a contiguous slice of columns. At grid=24 blocks × 32 warps =
-768 warps, n_columns=2048 → 2-3 columns per warp.
-
-### Parity with legacy GPU path
-
-**Semantics diverge**. Legacy: exactly `k = round(sparsity * n_cols)` columns
-active per step. Fused: variable, converging to `sparsity * n_cols` on
-average via the per-column EMA. Anomaly decay on repeating sequences is
-preserved (see `gpu_fused_tm_anomaly_decays_on_repeating_sequence` test).
-
-This is an intentional architectural change committed under
-`no-bypass/full-architecture` per program.md rules. The legacy top-K path
-(`step_many_cuda`) remains available for reference and can be re-enabled via
-`HYDRA_HTM_FUSED=0`.
-
-### Tests
-
-- `gpu_threshold_converges_to_sparsity` (tests.rs): 1000-step warmup on
-  random SDRs, then measure mean active cols/step on next 200 steps. Must
-  land within [0.25×, 4×] of `sparsity_target * n_cols`.
-- `gpu_fused_tm_anomaly_decays_on_repeating_sequence`: feed A,B,C repeating
-  for 300 steps. Late anomaly must be < early anomaly AND < 0.5.
-
-## Legacy Pipeline (kept for fallback)
-
-* SP: 5 kernels, bit-identical parity with CPU under strict-parity mode.
-* TM: 7 kernels, relaxed-parity with CPU.
-* Speedup at training size (B=8, T=2048, bits=16384): **3.83x** vs CPU.
-
-## Building
-
-CPU-only (default, zero CUDA dep):
-```bash
-cargo build --release
-```
-
-GPU-enabled:
-```bash
-export PATH=/usr/local/cuda-12.1/bin:$PATH
-export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
-export HTM_PTX_VERSION=7.8   # lower if driver older than nvcc
-cargo build --release --features gpu
-cargo test  --release --features gpu --lib   # fused path includes cooperative launch + grid-sync tests
-
-# Python wheel:
-maturin develop --release --features gpu --manifest-path htm_rust/Cargo.toml
-```
-
-## Architecture
-
-### Module layout
-```
-src/gpu/
-  mod.rs            # HTMRegionGpu pyclass + step_many_gpu (full pipeline)
-  sp_gpu.rs         # Persistent SP device buffers + step_batch_with_tm
-  tm_gpu.rs         # Persistent TM device buffers + step (predict→activate→learn)
-  tests.rs          # CPU-vs-GPU SP parity + end-to-end TM anomaly decay
-  kernels/
-    sp_overlap.cu       # per-column overlap reduction
-    sp_topk.cu          # k-WTA top-K winner selection
-    sp_learn.cu         # Hebbian +inc/-dec on proximal synapses
-    sp_duty.cu          # EMA duty-cycle update
-    sp_boost_fused.cu   # fused mean + exp boost (GPU-side)
-    tm_reset.cu         # per-step: snapshot active→prev, clear buffers
-    tm_predict.cu       # per-cell: score owned segments vs prev_active_bits
-    tm_activate.cu      # per-col: activate predicted cells OR burst
-    tm_learn.cu         # per-cell: reinforce correctly-predicted segments
-    tm_punish.cu        # per-cell: decay matching segs on inactive cols
-    tm_grow.cu          # per-bursting-col: reuse matching seg OR create new,
-                        #                    grow synapses to prev_winners
-    tm_anomaly.cu       # per-step: unpredicted/active ratio
-```
-
-### Persistent SP state (per region, unchanged from Phase 1)
-At n_cols=2048, S=40, bits=16384: ~355 KB persistent + ~90 KB transient.
-
-### Persistent TM state (per region)
-
-Capacity knobs (configured in `tm_gpu.rs`):
-- `MAX_SEGMENTS_PER_CELL = 4`
-- `MAX_SYN_PER_SEGMENT   = 20`
-
-At cells_per_col=32, n_cols=2048:
-- `n_cells          = 65_536`
-- `n_segments_max   = 262_144`   (~262K)
-- `n_synapses_max   = 5_242_880` (~5.2M)
-
-| Buffer                | Shape / type         | Notes                                  |
-|-----------------------|----------------------|----------------------------------------|
-| `seg_cell_id`         | (n_segs,) u32        | owning cell; U32_MAX = unused          |
-| `seg_syn_count`       | (n_segs,) u32        | #active synapses in slot               |
-| `syn_presyn`          | (n_segs × S,) u32    | presynaptic cell indices               |
-| `syn_perm`            | (n_segs × S,) i16    | permanence scaled 0..32767 (0.0..1.0)  |
-| `cell_seg_count`      | (n_cells,) u32       | segments allocated on each cell        |
-| `cell_active_bits`    | (n_cells/32,) u32    | packed bitset, current step            |
-| `cell_winner_bits`    | (n_cells/32,) u32    | packed bitset, current step            |
-| `cell_predictive_bits`| (n_cells/32,) u32    | set by predict, read by activate       |
-| `prev_active_bits`    | (n_cells/32,) u32    | snapshot at step start                 |
-| `prev_winner_bits`    | (n_cells/32,) u32    | snapshot at step start                 |
-| `col_predicted`       | (n_cols,) u8         | set if any cell in col is predictive   |
-| `col_best_match`      | (n_cols,) u32        | packed (pot<<21 | seg_id), atomicMax  |
-| `seg_num_active_conn` | (n_segs,) u32        | output of predict                      |
-| `seg_num_active_pot`  | (n_segs,) u32        | output of predict                      |
-| `unpredicted_count`   | (1,) u32             | atomic counter for anomaly             |
-| `burst_cols_flat`     | (n_cols,) u32        | list of bursting cols                  |
-| `burst_cols_count`    | (1,) u32             | length of above list                   |
-
-**Total per TM region: ~42 MB.** Batch of 8 regions: ~340 MB. Fits 6 GB RTX 3060.
-
-### Per-step pipeline (single iteration of `step_batch_with_tm`)
-
-```
-  SP side                            TM side
-  ---------                          ---------
-  1. D2D input slice → inp_dev
-  2. sp_overlap (n_cols blocks)
-  3. sp_topk    (1 block)
-  4. sp_learn   (n_cols blocks)
-  5. sp_duty    (n_cols/256 blocks)
-  6. sp_boost_fused (1 block)
-  7. D2D active_mask → cols_dev[ti]
-                                     8. tm_reset_step   (ceil(n_cells/32/256))
-                                     9. tm_predict      (n_cells blocks × 32 thr)
-                                    10. tm_activate     (n_cols/256 blocks)
-                                    11. tm_anomaly      (1 block)
-                                    if learn:
-                                    12. tm_learn        (n_cells blocks)
-                                    13. tm_punish       (n_cells blocks)
-                                    14. tm_grow         (n_cols blocks — early-exits)
-```
-
-No host sync in the T-step loop. At the end one `dtoh_sync_copy` each for
-`cols_dev` (T × n_cols bytes) and `anom_dev` (T × f32).
-
-## Parity
-
-### SP: strict bit-identical
-See Phase 1 docs — `gpu_sp_matches_cpu_with_learn` over 50 steps passes exact.
-
-### TM: relaxed-parity
-The GPU TM has known, deliberate deviations from CPU to admit massive parallelism:
-
-1. **Bursting winner cell**: CPU picks the least-used cell (fewest segments) with
-   random tiebreak. GPU picks cell 0 of the column (deterministic, branch-free).
-   Learning dynamics are preserved because segment creation/reinforcement is
-   the dominant effect, not which specific cell in a bursting column wins.
-
-2. **Permanence storage**: i16 fixed-point (scale 32767) vs f32. Rounding
-   differs by <=1 ULP of the scale (~3.0e-5), below any meaningful learning
-   quantum (inc=0.10, dec=0.10, predicted_segment_dec=0.10).
-
-3. **Grown synapse candidate order**: CPU randomly samples from prev_winner_cells.
-   GPU iterates prev_winner_bits words in a pseudo-random rotated order keyed
-   by (bursting_col_idx, iter_seed). Output is a different subset but same size.
-
-4. **Segment LRU eviction**: CPU tracks `last_used_iteration` per segment.
-   GPU wraps around (slot = count % max_segments_per_cell). In the autoresearch
-   loop where TM resets every forward, eviction rarely triggers.
-
-The GPU parity test (`gpu_tm_anomaly_decays_on_repeating_sequence`) feeds a
-repeating A,B,C sequence and asserts anomaly decays: **1.000 early → 0.000 late**.
-
-## Bottleneck Analysis
-
-| Source                           | Cost/step (B=8 T=2048)   |
-|----------------------------------|-------------------------:|
-| 14 kernel launches               | ~70 μs                   |
-| ~262K predict/learn/punish blocks| ~2.5 ms                  |
-| No D2H until end-of-batch        | 0 μs                     |
-| Final D2H (T × n_cols + T × f32) | ~200 μs per region       |
-
-Per-step wall time at B=8 T=2048:
-- CPU (reference): **~11.4 ms / step**
-- GPU (current):   **~2.98 ms / step**
-- **Speedup: 3.83x**
-
-## End-to-End Training Benchmark
-
-**Config**: B=8, T=2048, vocab=8192, 60-second time budget, full HYDRA stack
-(SDR Semantic + HTM + Mamba-3 + Engram + mHC + Hestia QAT).
-
-**Results**:
-- GPU util: **97-98% sustained**
-- VRAM: **5.4 GB / 6.0 GB** (90% utilisation)
-- Steps completed: 16
-- tok/sec: **~2,200-2,500** (stable post-warmup)
-- Final val_bpb: **2.249** (from ~3.1 initial)
-- Factual eval: 1/9 hits
-
-Compared to previous CPU-HTM baseline (~100 tok/s), the full-GPU HTM delivers
-**~22x end-to-end throughput** — far above the 3-10x target.
-
-## Bench Commands
-
-```bash
-source .venv/bin/activate
-export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
-
-# Microbench
-B=8 T=2048 python htm_rust/bench_gpu.py
-
-# Full training
-HYDRA_TIME_BUDGET=60 HYDRA_BATCH_SIZE=8 HYDRA_TOTAL_BATCH=32768 python -u train.py
-```
-
-## Known Limitations / Future Work
-
-- **Segment-compacted launches**: predict/learn/punish iterate all n_cells
-  blocks, using `cell_seg_count` to skip empty cells. A compacted live-cell
-  list would shave another ~40% of launch overhead.
-- **Winner selection**: currently cell 0 of bursting col. Proper least-used
-  selection would help stability of cross-column patterns.
-- **Single CUDA stream per region**: with B=8 regions we serialise on stream 0.
-  Multi-stream would lift the ~20% launch overhead at small batch sizes.
-- **Permanence bump on chronically under-stimulated columns**: SP's strict-parity
-  bump is not mirrored on GPU fast path. Effect on long runs needs measurement.
-- **`seg_num_active_conn` output is reused across reinforce + punish**: the two
-  kernels each launch n_cells blocks. They could be fused into one for one fewer
-  kernel launch per step.
-
-## Files
-
-- `htm_rust/build.rs` — nvcc-driven PTX compilation, 12 kernels.
-- `htm_rust/Cargo.toml` — `gpu` feature flag, cudarc dep.
-- `htm_rust/src/gpu/mod.rs` — `HTMRegionGpu` pyclass + `step_many_gpu`.
-- `htm_rust/src/gpu/sp_gpu.rs` — SP state + `step_batch_with_tm`.
-- `htm_rust/src/gpu/tm_gpu.rs` — TM state + `step`.
-- `htm_rust/src/gpu/tests.rs` — parity + correctness tests.
-- `htm_rust/src/gpu/kernels/*.cu` — 5 SP + 7 TM kernels.
-- `htm_rust/bench_gpu.py` — CPU-vs-GPU microbench.
-- `subsystems/htm.py` — transparent GPU/CPU backend selection in `HTMLayer`.
+# GPU HTM Backend
+
+## Status
+
+**FUSED MEGAKERNEL: entire T-timestep SP+TM forward collapsed into a single
+CUDA launch per forward pass.**
+
+* Legacy path: 12 kernels × T=2048 timesteps = 24K launches per forward.
+* Fused path: **1 launch per forward** (24000× launch-overhead reduction).
+* End-to-end training throughput: **~2.7k → ~60k tok/sec** (~22x speedup).
+* Fused path uses per-column threshold inhibition instead of global top-K
+  (see §Fused Kernel below — this is a real architectural change).
+
+## Fused Kernel
+
+### Why
+
+Global top-K column selection requires cross-block synchronization at every
+timestep. On WSL2/sm_86 without `-rdc=true`, `cooperative_groups::grid_sync()`
+is unreliable. Without a grid sync, collapsing the T-loop into one kernel is
+impossible, so every forward pays 12×T kernel launches and 90%+ of runtime is
+CUDA launch overhead + small-kernel tails.
+
+### How
+
+Replace global top-K with **per-column threshold activation**:
+
+    is_active[c] = (overlap[c] * boost[c]) > inhibition_threshold[c]
+
+`inhibition_threshold[c]` is a per-column scalar, learned via EMA update:
+
+    err = active_duty[c] - sparsity_target
+    new_thr = clamp(thr + thr_adapt_rate * err * 100, 0.1, 1000)
+
+This is biologically grounded (GABAergic local lateral inhibition in
+neocortical columns) and supported by HTM theory. The duty-cycle-driven
+feedback loop was already present; we simply redirect its output to drive
+activation threshold instead of multiplicative boost. The global top-K,
+which had no biological basis, is removed.
+
+### Cross-block coherence
+
+- **Ping-pong bitsets** for `cell_active_bits` and `cell_winner_bits`: at
+  even t write to `_a`, read from `_b`; at odd t reversed. This eliminates
+  the need for an in-place snapshot kernel between timesteps.
+- **Primary path: cooperative launch + hardware grid sync**. Host code probes
+  `CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH`, computes the cooperative whole-grid
+  residency limit from occupancy, and launches the fused megakernel with
+  `cuLaunchCooperativeKernel`. In-kernel barriers use
+  `cooperative_groups::this_grid().sync()`.
+- **Fallback path: software grid barrier** via a 3-slot atomic counter array
+  (`barrier_counters`). This remains as a compatibility fallback when
+  cooperative launch is unavailable.
+- **Launch invariant**: cooperative launch is capped to the hardware residency
+  limit for `blockDim.x = 1024`; software fallback remains capped conservatively
+  (`HTM_FUSED_GRID_CAP`, default 8) to avoid whole-grid spin deadlock.
+
+### Kernel structure
+
+```
+for t in 0..T:
+    # Phase 0: clear curr_active/curr_winner for my column range
+    grid_barrier()
+    # Phase A: SP overlap → boost → threshold → SP learn → duty + threshold EMA
+    grid_barrier()
+    # Phase B: TM predict (per cell, per seg) → TM learn (reinforce on match)
+    #                   → burst if none predicted → segment grow/reinforce
+    grid_barrier()
+    # Phase C: block 0 writes anomaly[t]
+```
+
+Each warp owns a contiguous slice of columns. At grid=24 blocks × 32 warps =
+768 warps, n_columns=2048 → 2-3 columns per warp.
+
+### Parity with legacy GPU path
+
+**Semantics diverge**. Legacy: exactly `k = round(sparsity * n_cols)` columns
+active per step. Fused: variable, converging to `sparsity * n_cols` on
+average via the per-column EMA. Anomaly decay on repeating sequences is
+preserved (see `gpu_fused_tm_anomaly_decays_on_repeating_sequence` test).
+
+This is an intentional architectural change committed under
+`no-bypass/full-architecture` per program.md rules. The legacy top-K path
+(`step_many_cuda`) remains available for reference and can be re-enabled via
+`HYDRA_HTM_FUSED=0`.
+
+### Tests
+
+- `gpu_threshold_converges_to_sparsity` (tests.rs): 1000-step warmup on
+  random SDRs, then measure mean active cols/step on next 200 steps. Must
+  land within [0.25×, 4×] of `sparsity_target * n_cols`.
+- `gpu_fused_tm_anomaly_decays_on_repeating_sequence`: feed A,B,C repeating
+  for 300 steps. Late anomaly must be < early anomaly AND < 0.5.
+
+## Legacy Pipeline (kept for fallback)
+
+* SP: 5 kernels, bit-identical parity with CPU under strict-parity mode.
+* TM: 7 kernels, relaxed-parity with CPU.
+* Speedup at training size (B=8, T=2048, bits=16384): **3.83x** vs CPU.
+
+## Building
+
+CPU-only (default, zero CUDA dep):
+```bash
+cargo build --release
+```
+
+GPU-enabled:
+```bash
+export PATH=/usr/local/cuda-12.1/bin:$PATH
+export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+export HTM_PTX_VERSION=7.8   # lower if driver older than nvcc
+cargo build --release --features gpu
+cargo test  --release --features gpu --lib   # fused path includes cooperative launch + grid-sync tests
+
+# Python wheel:
+maturin develop --release --features gpu --manifest-path htm_rust/Cargo.toml
+```
+
+## Architecture
+
+### Module layout
+```
+src/gpu/
+  mod.rs            # HTMRegionGpu pyclass + step_many_gpu (full pipeline)
+  sp_gpu.rs         # Persistent SP device buffers + step_batch_with_tm
+  tm_gpu.rs         # Persistent TM device buffers + step (predict→activate→learn)
+  tests.rs          # CPU-vs-GPU SP parity + end-to-end TM anomaly decay
+  kernels/
+    sp_overlap.cu       # per-column overlap reduction
+    sp_topk.cu          # k-WTA top-K winner selection
+    sp_learn.cu         # Hebbian +inc/-dec on proximal synapses
+    sp_duty.cu          # EMA duty-cycle update
+    sp_boost_fused.cu   # fused mean + exp boost (GPU-side)
+    tm_reset.cu         # per-step: snapshot active→prev, clear buffers
+    tm_predict.cu       # per-cell: score owned segments vs prev_active_bits
+    tm_activate.cu      # per-col: activate predicted cells OR burst
+    tm_learn.cu         # per-cell: reinforce correctly-predicted segments
+    tm_punish.cu        # per-cell: decay matching segs on inactive cols
+    tm_grow.cu          # per-bursting-col: reuse matching seg OR create new,
+                        #                    grow synapses to prev_winners
+    tm_anomaly.cu       # per-step: unpredicted/active ratio
+```
+
+### Persistent SP state (per region, unchanged from Phase 1)
+At n_cols=2048, S=40, bits=16384: ~355 KB persistent + ~90 KB transient.
+
+### Persistent TM state (per region)
+
+Capacity knobs (configured in `tm_gpu.rs`):
+- `MAX_SEGMENTS_PER_CELL = 4`
+- `MAX_SYN_PER_SEGMENT   = 20`
+
+At cells_per_col=32, n_cols=2048:
+- `n_cells          = 65_536`
+- `n_segments_max   = 262_144`   (~262K)
+- `n_synapses_max   = 5_242_880` (~5.2M)
+
+| Buffer                | Shape / type         | Notes                                  |
+|-----------------------|----------------------|----------------------------------------|
+| `seg_cell_id`         | (n_segs,) u32        | owning cell; U32_MAX = unused          |
+| `seg_syn_count`       | (n_segs,) u32        | #active synapses in slot               |
+| `syn_presyn`          | (n_segs × S,) u32    | presynaptic cell indices               |
+| `syn_perm`            | (n_segs × S,) i16    | permanence scaled 0..32767 (0.0..1.0)  |
+| `cell_seg_count`      | (n_cells,) u32       | segments allocated on each cell        |
+| `cell_active_bits`    | (n_cells/32,) u32    | packed bitset, current step            |
+| `cell_winner_bits`    | (n_cells/32,) u32    | packed bitset, current step            |
+| `cell_predictive_bits`| (n_cells/32,) u32    | set by predict, read by activate       |
+| `prev_active_bits`    | (n_cells/32,) u32    | snapshot at step start                 |
+| `prev_winner_bits`    | (n_cells/32,) u32    | snapshot at step start                 |
+| `col_predicted`       | (n_cols,) u8         | set if any cell in col is predictive   |
+| `col_best_match`      | (n_cols,) u32        | packed (pot<<21 | seg_id), atomicMax  |
+| `seg_num_active_conn` | (n_segs,) u32        | output of predict                      |
+| `seg_num_active_pot`  | (n_segs,) u32        | output of predict                      |
+| `unpredicted_count`   | (1,) u32             | atomic counter for anomaly             |
+| `burst_cols_flat`     | (n_cols,) u32        | list of bursting cols                  |
+| `burst_cols_count`    | (1,) u32             | length of above list                   |
+
+**Total per TM region: ~42 MB.** Batch of 8 regions: ~340 MB. Fits 6 GB RTX 3060.
+
+### Per-step pipeline (single iteration of `step_batch_with_tm`)
+
+```
+  SP side                            TM side
+  ---------                          ---------
+  1. D2D input slice → inp_dev
+  2. sp_overlap (n_cols blocks)
+  3. sp_topk    (1 block)
+  4. sp_learn   (n_cols blocks)
+  5. sp_duty    (n_cols/256 blocks)
+  6. sp_boost_fused (1 block)
+  7. D2D active_mask → cols_dev[ti]
+                                     8. tm_reset_step   (ceil(n_cells/32/256))
+                                     9. tm_predict      (n_cells blocks × 32 thr)
+                                    10. tm_activate     (n_cols/256 blocks)
+                                    11. tm_anomaly      (1 block)
+                                    if learn:
+                                    12. tm_learn        (n_cells blocks)
+                                    13. tm_punish       (n_cells blocks)
+                                    14. tm_grow         (n_cols blocks — early-exits)
+```
+
+No host sync in the T-step loop. At the end one `dtoh_sync_copy` each for
+`cols_dev` (T × n_cols bytes) and `anom_dev` (T × f32).
+
+## Parity
+
+### SP: strict bit-identical
+See Phase 1 docs — `gpu_sp_matches_cpu_with_learn` over 50 steps passes exact.
+
+### TM: relaxed-parity
+The GPU TM has known, deliberate deviations from CPU to admit massive parallelism:
+
+1. **Bursting winner cell**: CPU picks the least-used cell (fewest segments) with
+   random tiebreak. GPU picks cell 0 of the column (deterministic, branch-free).
+   Learning dynamics are preserved because segment creation/reinforcement is
+   the dominant effect, not which specific cell in a bursting column wins.
+
+2. **Permanence storage**: i16 fixed-point (scale 32767) vs f32. Rounding
+   differs by <=1 ULP of the scale (~3.0e-5), below any meaningful learning
+   quantum (inc=0.10, dec=0.10, predicted_segment_dec=0.10).
+
+3. **Grown synapse candidate order**: CPU randomly samples from prev_winner_cells.
+   GPU iterates prev_winner_bits words in a pseudo-random rotated order keyed
+   by (bursting_col_idx, iter_seed). Output is a different subset but same size.
+
+4. **Segment LRU eviction**: CPU tracks `last_used_iteration` per segment.
+   GPU wraps around (slot = count % max_segments_per_cell). In the autoresearch
+   loop where TM resets every forward, eviction rarely triggers.
+
+The GPU parity test (`gpu_tm_anomaly_decays_on_repeating_sequence`) feeds a
+repeating A,B,C sequence and asserts anomaly decays: **1.000 early → 0.000 late**.
+
+## Bottleneck Analysis
+
+| Source                           | Cost/step (B=8 T=2048)   |
+|----------------------------------|-------------------------:|
+| 14 kernel launches               | ~70 μs                   |
+| ~262K predict/learn/punish blocks| ~2.5 ms                  |
+| No D2H until end-of-batch        | 0 μs                     |
+| Final D2H (T × n_cols + T × f32) | ~200 μs per region       |
+
+Per-step wall time at B=8 T=2048:
+- CPU (reference): **~11.4 ms / step**
+- GPU (current):   **~2.98 ms / step**
+- **Speedup: 3.83x**
+
+## End-to-End Training Benchmark
+
+**Config**: B=8, T=2048, vocab=8192, 60-second time budget, full HYDRA stack
+(SDR Semantic + HTM + Mamba-3 + Engram + mHC + Hestia QAT).
+
+**Results**:
+- GPU util: **97-98% sustained**
+- VRAM: **5.4 GB / 6.0 GB** (90% utilisation)
+- Steps completed: 16
+- tok/sec: **~2,200-2,500** (stable post-warmup)
+- Final val_bpb: **2.249** (from ~3.1 initial)
+- Factual eval: 1/9 hits
+
+Compared to previous CPU-HTM baseline (~100 tok/s), the full-GPU HTM delivers
+**~22x end-to-end throughput** — far above the 3-10x target.
+
+## Bench Commands
+
+```bash
+source .venv/bin/activate
+export LD_LIBRARY_PATH=/usr/lib/wsl/lib:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+
+# Microbench
+B=8 T=2048 python htm_rust/bench_gpu.py
+
+# Full training
+HYDRA_TIME_BUDGET=60 HYDRA_BATCH_SIZE=8 HYDRA_TOTAL_BATCH=32768 python -u train.py
+```
+
+## Known Limitations / Future Work
+
+- **Segment-compacted launches**: predict/learn/punish iterate all n_cells
+  blocks, using `cell_seg_count` to skip empty cells. A compacted live-cell
+  list would shave another ~40% of launch overhead.
+- **Winner selection**: currently cell 0 of bursting col. Proper least-used
+  selection would help stability of cross-column patterns.
+- **Single CUDA stream per region**: with B=8 regions we serialise on stream 0.
+  Multi-stream would lift the ~20% launch overhead at small batch sizes.
+- **Permanence bump on chronically under-stimulated columns**: SP's strict-parity
+  bump is not mirrored on GPU fast path. Effect on long runs needs measurement.
+- **`seg_num_active_conn` output is reused across reinforce + punish**: the two
+  kernels each launch n_cells blocks. They could be fused into one for one fewer
+  kernel launch per step.
+
+## Files
+
+- `htm_rust/build.rs` — nvcc-driven PTX compilation, 12 kernels.
+- `htm_rust/Cargo.toml` — `gpu` feature flag, cudarc dep.
+- `htm_rust/src/gpu/mod.rs` — `HTMRegionGpu` pyclass + `step_many_gpu`.
+- `htm_rust/src/gpu/sp_gpu.rs` — SP state + `step_batch_with_tm`.
+- `htm_rust/src/gpu/tm_gpu.rs` — TM state + `step`.
+- `htm_rust/src/gpu/tests.rs` — parity + correctness tests.
+- `htm_rust/src/gpu/kernels/*.cu` — 5 SP + 7 TM kernels.
+- `htm_rust/bench_gpu.py` — CPU-vs-GPU microbench.
+- `subsystems/htm.py` — transparent GPU/CPU backend selection in `HTMLayer`.

overlay/htm_rust/pyproject.toml

@@ -1,17 +1,17 @@
-[build-system]
-requires = ["maturin>=1.4,<2.0"]
-build-backend = "maturin"
-
-[project]
-name = "htm_rust"
-version = "0.1.0"
-description = "Numenta BAMI-spec HTM (Spatial Pooler + Temporal Memory) in Rust with pyo3 bindings"
-requires-python = ">=3.11"
-classifiers = [
-    "Programming Language :: Rust",
-    "Programming Language :: Python :: Implementation :: CPython",
-]
-
-[tool.maturin]
-features = ["pyo3/extension-module"]
-module-name = "htm_rust"
+[build-system]
+requires = ["maturin>=1.4,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "htm_rust"
+version = "0.1.0"
+description = "Numenta BAMI-spec HTM (Spatial Pooler + Temporal Memory) in Rust with pyo3 bindings"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Rust",
+    "Programming Language :: Python :: Implementation :: CPython",
+]
+
+[tool.maturin]
+features = ["pyo3/extension-module"]
+module-name = "htm_rust"

overlay/htm_rust/src/gpu/fused.rs

@@ -1,719 +1,719 @@
-//! Fused HTM megakernel launcher.
-//!
-//! Collapses the 12-kernel per-timestep pipeline (and the outer T-loop) into
-//! a single kernel launch per forward. See `kernels/htm_fused_step.cu` for
-//! the kernel design and the cross-block coherence strategy (grid barrier
-//! via device counter with all blocks concurrently resident).
-//!
-//! Launch invariant: `grid_dim.x <= concurrent-block capacity`. Host code
-//! probes the device SM count at construction and caps grid_dim.x
-//! accordingly — otherwise the grid barrier deadlocks.
-//!
-//! Semantic change from the top-K pipeline: activation is per-column
-//! threshold-based (local lateral inhibition) instead of global top-K.
-//! A per-column `inhibition_threshold` is tracked and EMA-steered to hit
-//! the sparsity target. This is a real architectural change and is
-//! documented in `docs/GPU_HTM.md`.
-
-#![cfg(feature = "gpu")]
-
-use std::ffi::CString;
-use std::sync::Arc;
-
-use cudarc::driver::{
-    result, sys, CudaDevice, CudaSlice, DevicePtr, DeviceRepr, DriverError, LaunchConfig,
-};
-use cudarc::nvrtc::Ptx;
-
-use super::sp_gpu::SpatialPoolerGpu;
-use super::tm_gpu::{TemporalMemoryGpu, MAX_SEGMENTS_PER_CELL, MAX_SYN_PER_SEGMENT};
-
-const PTX_HTM_FUSED: &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/htm_fused_step.ptx"));
-
-/// Struct-by-value pointer pack — matches C-side `FusedPtrs`.
-///
-/// NOTE: `barrier_counters` is kept as an ABI-compat dummy (always 0). The
-/// C-side `FusedPtrs` still has the field at the same byte offset; removing
-/// it here would shift all subsequent fields and break the layout. Worker A
-/// will eventually delete the field from both sides once the kernel is
-/// updated; until then we zero it.
-#[repr(C)]
-#[derive(Clone, Copy)]
-pub struct FusedPtrs {
-    pub syn_bit: u64,
-    pub syn_perm: u64,
-    pub boost: u64,
-    pub active_duty: u64,
-    pub inhibition_threshold: u64,
-    pub seg_cell_id: u64,
-    pub seg_syn_count: u64,
-    pub syn_presyn: u64,
-    pub tm_syn_perm: u64,
-    pub cell_seg_count: u64,
-    pub cell_active_a: u64,
-    pub cell_active_b: u64,
-    pub cell_winner_a: u64,
-    pub cell_winner_b: u64,
-    pub inputs: u64,
-    pub cols_out: u64,
-    pub anom_out: u64,
-    /// ABI-compat dummy — always 0. No device memory is allocated for this
-    /// field; the cluster barrier replaces the old software DLB barrier.
-    pub barrier_counters: u64,
-    pub step_scratch: u64,
-}
-
-unsafe impl DeviceRepr for FusedPtrs {}
-
-/// Launch-time config — matches C-side `FusedConfig` 1:1.
-#[repr(C)]
-#[derive(Clone, Copy)]
-pub struct FusedConfig {
-    pub input_bits: u32,
-    pub n_columns: u32,
-    pub synapses_per_col: u32,
-    pub conn_thr: f32,
-    pub sp_inc: f32,
-    pub sp_dec: f32,
-    pub sparsity_target: f32,
-    pub duty_alpha: f32,
-    pub thr_adapt_rate: f32,
-    pub cells_per_column: u32,
-    pub n_cells: u32,
-    pub bits_words: u32,
-    pub max_segments_per_cell: u32,
-    pub synapses_per_segment: u32,
-    pub activation_threshold: u32,
-    pub learning_threshold: u32,
-    pub max_new_synapses: u32,
-    pub conn_thr_i16: i32,
-    pub perm_inc_i16: i32,
-    pub perm_dec_i16: i32,
-    pub predicted_seg_dec_i16: i32,
-    pub initial_perm_i16: i32,
-    pub t: u32,
-    pub learn: u32,
-    pub iter_seed: u32,
-    pub cooperative_grid_sync: u32,
-}
-
-unsafe impl DeviceRepr for FusedConfig {}
-
-/// Cluster launch parameters probed at construction time.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub(crate) struct ClusterInfo {
-    /// Maximum cluster size supported by this device (0 = cluster unsupported).
-    pub max_cluster_size: u32,
-}
-
-// There is only ONE launch mode: non-cooperative launch with Hopper Thread
-// Block Cluster attribute (`CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION`). The old
-// software DLB barrier and the cooperative-launch path are both removed.
-// Cluster barriers replace both.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub(crate) struct FusedLaunchPlan {
-    pub grid_dim_x: u32,
-    pub block_dim_x: u32,
-    pub cooperative_grid_limit: u32,
-    pub sm_count: u32,
-}
-
-fn fused_grid_cap_override() -> Option<u32> {
-    std::env::var("HTM_FUSED_GRID_CAP")
-        .ok()
-        .and_then(|s| s.parse::<u32>().ok())
-        .map(|v| v.max(1))
-}
-
-pub(crate) fn plan_fused_launch(
-    sm_count: u32,
-    cooperative_supported: bool,
-    cooperative_grid_limit: u32,
-    grid_cap_override: Option<u32>,
-) -> Result<FusedLaunchPlan, String> {
-    let sm_count = sm_count.max(1);
-    // 1024 threads/block exceeds the register file on Ampere and makes the
-    // cooperative-grid residency probe lie when the launch uses a different
-    // block size. Keep the planned block size identical to the occupancy probe.
-    let block_dim_x = 256u32;
-
-    // Cluster launch path: cooperative launch is not required. Keep the probe
-    // result for residency estimation only.
-    if !cooperative_supported {
-        eprintln!("[htm_rust] INFO: cooperative launch unsupported; cluster path only.");
-    }
-
-    // Cluster constraint: grid_dim_x must equal the cluster size (16) so that
-    // each region maps to exactly one cluster. `HTM_FUSED_GRID_CAP` can lower
-    // this for debugging but should not exceed 16 for cluster correctness.
-    let default_grid_cap = 16u32;
-    let grid_cap = grid_cap_override.unwrap_or(default_grid_cap).min(16);
-    let resident_bound = if cooperative_grid_limit > 0 {
-        cooperative_grid_limit.max(sm_count * 2)
-    } else {
-        sm_count * 2
-    };
-    Ok(FusedLaunchPlan {
-        grid_dim_x: resident_bound.min(grid_cap).max(1),
-        block_dim_x,
-        cooperative_grid_limit: resident_bound,
-        sm_count,
-    })
-}
-
-pub(crate) fn plan_batched_grid_dim(
-    grid_dim_x: u32,
-    cooperative_grid_limit: u32,
-    batch_regions: usize,
-    use_cluster: bool,
-) -> Result<u32, String> {
-    if use_cluster {
-        return Ok(grid_dim_x.max(1));
-    }
-
-    let batch_regions = batch_regions.max(1) as u32;
-    if cooperative_grid_limit == 0 {
-        return Err("COOPERATIVE_LAUNCH_TOO_LARGE: cooperative launch limit unavailable".into());
-    }
-
-    let max_grid_x = cooperative_grid_limit / batch_regions;
-    if max_grid_x == 0 {
-        return Err(format!(
-            "COOPERATIVE_LAUNCH_TOO_LARGE: batch_regions={batch_regions} exceeds cooperative_grid_limit={cooperative_grid_limit}"
-        ));
-    }
-
-    Ok(grid_dim_x.min(max_grid_x).max(1))
-}
-
-pub(super) struct RawFusedKernel {
-    module: sys::CUmodule,
-    pub(super) function: sys::CUfunction,
-    pub(super) function_batched: sys::CUfunction,
-}
-
-unsafe impl Send for RawFusedKernel {}
-unsafe impl Sync for RawFusedKernel {}
-
-impl Drop for RawFusedKernel {
-    fn drop(&mut self) {
-        unsafe {
-            let _ = result::module::unload(self.module);
-        }
-    }
-}
-
-/// Owns fused-path-only device state:
-///   - per-column inhibition threshold (replaces global top-K)
-///   - ping-pong cell_active/cell_winner bitsets
-///   - step_scratch (n_active, n_unpred per timestep)
-///   - cluster launch capability info
-pub struct FusedState {
-    dev: Arc<CudaDevice>,
-    pub(super) raw_kernel: RawFusedKernel,
-
-    pub inhibition_threshold: CudaSlice<f32>,
-    pub cell_active_bits_a: CudaSlice<u32>,
-    pub cell_active_bits_b: CudaSlice<u32>,
-    pub cell_winner_bits_a: CudaSlice<u32>,
-    pub cell_winner_bits_b: CudaSlice<u32>,
-    pub step_scratch: CudaSlice<u32>, // length 6
-
-    pub grid_dim_x: u32,
-    pub block_dim_x: u32,
-    pub cooperative_grid_limit: u32,
-    pub iter_counter: u32,
-
-    /// Hopper cluster launch capability (0 = unsupported).
-    pub cluster_info: ClusterInfo,
-
-    // Config mirror (read-only after init).
-    #[allow(dead_code)]
-    pub initial_threshold: f32,
-}
-
-impl FusedState {
-    pub fn new(
-        dev: Arc<CudaDevice>,
-        n_columns: usize,
-        cells_per_column: usize,
-        initial_threshold: f32,
-    ) -> Result<Self, DriverError> {
-        let n_cells = n_columns * cells_per_column;
-        assert!(
-            n_cells % 32 == 0,
-            "n_cells must be divisible by 32 for bitsets"
-        );
-        let bits_words = n_cells / 32;
-
-        let mut inhibition_threshold = dev.alloc_zeros::<f32>(n_columns)?;
-        let init_vec = vec![initial_threshold; n_columns];
-        dev.htod_sync_copy_into(&init_vec, &mut inhibition_threshold)?;
-
-        let cell_active_bits_a = dev.alloc_zeros::<u32>(bits_words)?;
-        let cell_active_bits_b = dev.alloc_zeros::<u32>(bits_words)?;
-        let cell_winner_bits_a = dev.alloc_zeros::<u32>(bits_words)?;
-        let cell_winner_bits_b = dev.alloc_zeros::<u32>(bits_words)?;
-        let step_scratch = dev.alloc_zeros::<u32>(6)?;
-
-        unsafe {
-            result::ctx::set_current(*dev.cu_primary_ctx())?;
-        }
-        if dev.get_func("htm_fused", "htm_fused_step").is_none() {
-            dev.load_ptx(
-                Ptx::from_src(PTX_HTM_FUSED),
-                "htm_fused",
-                &["htm_fused_step", "htm_fused_step_batched"],
-            )?;
-        }
-        let ptx = CString::new(PTX_HTM_FUSED).expect("PTX contains no interior nul bytes");
-        let module = unsafe { result::module::load_data(ptx.as_ptr().cast()) }?;
-        let function = unsafe {
-            result::module::get_function(module, CString::new("htm_fused_step").unwrap())
-        }?;
-        let function_batched = unsafe {
-            result::module::get_function(module, CString::new("htm_fused_step_batched").unwrap())
-        }?;
-
-        // Cluster size 16 on Hopper is "non-portable" (> 8 requires opt-in).
-        // Must set CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED=1 on
-        // every launched kernel function, otherwise cuLaunchKernelEx rejects
-        // the cluster dim with CUDA_ERROR_INVALID_CLUSTER_SIZE.
-        unsafe {
-            let attr =
-                sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED;
-            // Ignore errors: older CUDA may lack the attribute, in which case
-            // only portable sizes (<= 8) work — plan_fused_launch caps at 8.
-            let _ = sys::lib().cuFuncSetAttribute(function, attr, 1);
-            let _ = sys::lib().cuFuncSetAttribute(function_batched, attr, 1);
-        }
-
-        // Probe SM count.
-        let sm_count = match dev.attribute(
-            cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
-        ) {
-            Ok(v) => v as u32,
-            Err(_) => 16u32,
-        };
-
-        // T1: Probe Hopper cluster launch capability.
-        let max_cluster_size = match dev
-            .attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_CLUSTER_LAUNCH)
-        {
-            Ok(v) if v > 0 => {
-                // H200/sm_90a supports up to 16 blocks per cluster.
-                // There is no MAX_CLUSTER_SIZE attribute in CUDA 12.4; hard-code the
-                // Hopper maximum which is 16 (8 SMs × 2 blocks/SM = 16 blocks/cluster).
-                16u32
-            }
-            _ => 0u32,
-        };
-        eprintln!("[htm_rust] cluster: max_cluster_size={}", max_cluster_size);
-        let cluster_info = ClusterInfo { max_cluster_size };
-
-        let cooperative_supported = matches!(
-            dev.attribute(sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH),
-            Ok(v) if v > 0
-        );
-        let cooperative_grid_limit = if cooperative_supported {
-            let blocks_per_sm = unsafe {
-                // Must match plan_fused_launch(): the A10G/Ampere-safe fused
-                // kernel launch uses 256 threads/block, not the historical
-                // 1024-thread Hopper occupancy probe.
-                result::occupancy::max_active_block_per_multiprocessor(function, 256, 0)
-            }
-            .ok()
-            .map(|v| v.max(0) as u32)
-            .unwrap_or(0);
-            sm_count.saturating_mul(blocks_per_sm)
-        } else {
-            0
-        };
-        let launch_plan = plan_fused_launch(
-            sm_count,
-            cooperative_supported,
-            cooperative_grid_limit,
-            fused_grid_cap_override(),
-        )
-        .map_err(|msg| {
-            // Surface as a CUDA-ish error so callers can propagate.
-            eprintln!("[htm_rust] FATAL: {msg}");
-            DriverError(cudarc::driver::sys::CUresult::CUDA_ERROR_NOT_SUPPORTED)
-        })?;
-
-        eprintln!(
-            "[htm_rust] fused kernel: sm_count={} grid_dim_x={} cooperative_grid_limit={} cluster_max={}",
-            launch_plan.sm_count, launch_plan.grid_dim_x, launch_plan.cooperative_grid_limit,
-            cluster_info.max_cluster_size,
-        );
-
-        Ok(Self {
-            dev,
-            raw_kernel: RawFusedKernel {
-                module,
-                function,
-                function_batched,
-            },
-            inhibition_threshold,
-            cell_active_bits_a,
-            cell_active_bits_b,
-            cell_winner_bits_a,
-            cell_winner_bits_b,
-            step_scratch,
-            grid_dim_x: launch_plan.grid_dim_x,
-            block_dim_x: launch_plan.block_dim_x,
-            cooperative_grid_limit: launch_plan.cooperative_grid_limit,
-            iter_counter: 0,
-            cluster_info,
-            initial_threshold,
-        })
-    }
-
-    /// Reset fused state. Called at region.reset().
-    pub fn reset(&mut self) -> Result<(), DriverError> {
-        self.dev.memset_zeros(&mut self.cell_active_bits_a)?;
-        self.dev.memset_zeros(&mut self.cell_active_bits_b)?;
-        self.dev.memset_zeros(&mut self.cell_winner_bits_a)?;
-        self.dev.memset_zeros(&mut self.cell_winner_bits_b)?;
-        self.dev.memset_zeros(&mut self.step_scratch)?;
-        // Do NOT reset inhibition_threshold — it's learned state. A hard
-        // reset of TM state should NOT forget the sparsity calibration.
-        Ok(())
-    }
-}
-
-/// Launch the fused megakernel. Processes all T timesteps in one kernel.
-///
-/// Uses `cuLaunchKernelEx` with `CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION=(16,1,1)`
-/// when the device supports cluster launch, otherwise falls back to a plain
-/// `launch_kernel`. For single-region launches, grid_dim_x <= 16 ensures the
-/// entire grid fits in one cluster.
-#[allow(clippy::too_many_arguments)]
-pub fn launch_fused(
-    sp: &mut SpatialPoolerGpu,
-    tm: &mut TemporalMemoryGpu,
-    fused: &mut FusedState,
-    inputs_flat: &CudaSlice<u8>,
-    cols_out: &mut CudaSlice<u8>,
-    anom_out: &mut CudaSlice<f32>,
-    t: usize,
-    input_bits: usize,
-    learn: bool,
-) -> Result<(), DriverError> {
-    // Reset step_scratch before each launch (safe re-entry).
-    sp.dev_ref().memset_zeros(&mut fused.step_scratch)?;
-
-    fused.iter_counter = fused.iter_counter.wrapping_add(1);
-
-    let cfg = FusedConfig {
-        input_bits: input_bits as u32,
-        n_columns: sp.n_columns_accessor() as u32,
-        synapses_per_col: sp.synapses_per_col_accessor() as u32,
-        conn_thr: sp.conn_thr_accessor(),
-        sp_inc: sp.inc_accessor(),
-        sp_dec: sp.dec_accessor(),
-        sparsity_target: sp.sparsity_accessor(),
-        duty_alpha: 1.0f32 / sp.duty_period_accessor().max(1.0),
-        thr_adapt_rate: 0.001f32,
-        cells_per_column: tm.cells_per_column as u32,
-        n_cells: tm.n_cells as u32,
-        bits_words: tm.bits_words as u32,
-        max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
-        synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
-        activation_threshold: tm.activation_threshold,
-        learning_threshold: tm.learning_threshold,
-        max_new_synapses: tm.max_new_synapse_count,
-        conn_thr_i16: tm.conn_thr_i16 as i32,
-        perm_inc_i16: tm.perm_inc_i16 as i32,
-        perm_dec_i16: tm.perm_dec_i16 as i32,
-        predicted_seg_dec_i16: tm.predicted_seg_dec_i16 as i32,
-        initial_perm_i16: tm.initial_perm_i16 as i32,
-        t: t as u32,
-        learn: if learn { 1 } else { 0 },
-        iter_seed: fused.iter_counter,
-        cooperative_grid_sync: 1,
-    };
-
-    let ptrs = FusedPtrs {
-        syn_bit: *sp.syn_bit_accessor().device_ptr(),
-        syn_perm: *sp.syn_perm_accessor().device_ptr(),
-        boost: *sp.boost_accessor().device_ptr(),
-        active_duty: *sp.active_duty_accessor().device_ptr(),
-        inhibition_threshold: *fused.inhibition_threshold.device_ptr(),
-        seg_cell_id: *tm.seg_cell_id_accessor().device_ptr(),
-        seg_syn_count: *tm.seg_syn_count_accessor().device_ptr(),
-        syn_presyn: *tm.syn_presyn_accessor().device_ptr(),
-        tm_syn_perm: *tm.syn_perm_accessor().device_ptr(),
-        cell_seg_count: *tm.cell_seg_count_accessor().device_ptr(),
-        cell_active_a: *fused.cell_active_bits_a.device_ptr(),
-        cell_active_b: *fused.cell_active_bits_b.device_ptr(),
-        cell_winner_a: *fused.cell_winner_bits_a.device_ptr(),
-        cell_winner_b: *fused.cell_winner_bits_b.device_ptr(),
-        inputs: *inputs_flat.device_ptr(),
-        cols_out: *cols_out.device_ptr(),
-        anom_out: *anom_out.device_ptr(),
-        barrier_counters: 0u64, // ABI-compat dummy; cluster barrier replaces DLB.
-        step_scratch: *fused.step_scratch.device_ptr(),
-    };
-
-    let grid_x = fused.grid_dim_x;
-    let block_x = fused.block_dim_x;
-    let cu_stream = *sp.dev_ref().cu_stream();
-    let use_cluster = fused.cluster_info.max_cluster_size > 0;
-
-    unsafe {
-        result::ctx::set_current(*sp.dev_ref().cu_primary_ctx())?;
-        let mut kernel_params: [*mut std::ffi::c_void; 2] = [
-            (&ptrs as *const FusedPtrs).cast_mut().cast(),
-            (&cfg as *const FusedConfig).cast_mut().cast(),
-        ];
-
-        if use_cluster {
-            // T10: Hopper cluster launch with CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION.
-            // cluster_dim=(16,1,1) maps the entire single-region grid into one cluster.
-            let mut attr: sys::CUlaunchAttribute = std::mem::zeroed();
-            attr.id = sys::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
-            attr.value.clusterDim.x = 16;
-            attr.value.clusterDim.y = 1;
-            attr.value.clusterDim.z = 1;
-
-            let mut launch_cfg: sys::CUlaunchConfig = std::mem::zeroed();
-            launch_cfg.gridDimX = grid_x;
-            launch_cfg.gridDimY = 1;
-            launch_cfg.gridDimZ = 1;
-            launch_cfg.blockDimX = block_x;
-            launch_cfg.blockDimY = 1;
-            launch_cfg.blockDimZ = 1;
-            launch_cfg.sharedMemBytes = 0;
-            launch_cfg.hStream = cu_stream;
-            launch_cfg.numAttrs = 1;
-            launch_cfg.attrs = &mut attr as *mut sys::CUlaunchAttribute;
-
-            let ret = sys::lib().cuLaunchKernelEx(
-                &launch_cfg as *const sys::CUlaunchConfig,
-                fused.raw_kernel.function,
-                kernel_params.as_mut_ptr(),
-                std::ptr::null_mut(),
-            );
-            if ret != sys::CUresult::CUDA_SUCCESS {
-                return Err(DriverError(ret));
-            }
-        } else {
-            // Pre-Hopper: cooperative kernel launch. The fused kernel uses
-            // cg::this_grid().sync(); normal launches poison the CUDA context
-            // with an asynchronous unspecified launch failure.
-            let ret = sys::lib().cuLaunchCooperativeKernel(
-                fused.raw_kernel.function,
-                grid_x,
-                1,
-                1,
-                block_x,
-                1,
-                1,
-                0,
-                cu_stream,
-                kernel_params.as_mut_ptr(),
-            );
-            if ret != sys::CUresult::CUDA_SUCCESS {
-                return Err(DriverError(ret));
-            }
-        }
-    }
-
-    Ok(())
-}
-
-/// Single batched non-cooperative launch for B regions with DLB sync. Uses the same kernel
-/// body; each block reads its region's FusedPtrs from a device-side array
-/// indexed by blockIdx.y. All regions share the same config (same
-/// input_bits/n_columns/etc.) so we pass one FusedConfig.
-///
-/// This breaks through the CUDA cooperative-kernel device-level
-/// serialization: multiple cooperative launches are serialized regardless
-/// of stream, but one cooperative launch with grid.y=B processes all
-/// regions in a single invocation — ~B× speedup vs B sequential launches.
-#[allow(clippy::too_many_arguments)]
-/// Low-level raw-pointer entry, called by PyO3 binding which holds the
-/// mutable borrows. Safety: each `*mut HTMRegionGpu` must point to a live,
-/// uniquely-borrowed region. All regions must be distinct.
-pub(super) fn launch_fused_batched_raw(
-    region_ptrs: &[*mut super::HTMRegionGpu],
-    inputs_per_region: &[u64],
-    cols_per_region: &[u64],
-    anom_per_region: &[u64],
-    t: usize,
-    input_bits: usize,
-    learn: bool,
-) -> Result<(), DriverError> {
-    let b = region_ptrs.len();
-    assert_eq!(inputs_per_region.len(), b);
-    assert_eq!(cols_per_region.len(), b);
-    assert_eq!(anom_per_region.len(), b);
-    assert!(b >= 1, "need at least one region");
-
-    // Reset per-region step_scratch before each launch.
-    for &rp in region_ptrs.iter() {
-        let r = unsafe { &mut *rp };
-        let dev = r.sp_gpu.dev_ref().clone();
-        dev.memset_zeros(&mut r.fused_state.step_scratch)?;
-        r.fused_state.iter_counter = r.fused_state.iter_counter.wrapping_add(1);
-    }
-
-    // Shared config — all regions use identical sp/tm parameters.
-    let (grid_x, block_x, cooperative_grid_limit, function_batched, cu_stream, cu_ctx) = {
-        let r0 = unsafe { &*region_ptrs[0] };
-        (
-            r0.fused_state.grid_dim_x,
-            r0.fused_state.block_dim_x,
-            r0.fused_state.cooperative_grid_limit,
-            r0.fused_state.raw_kernel.function_batched,
-            *r0.sp_gpu.dev_ref().cu_stream(),
-            *r0.sp_gpu.dev_ref().cu_primary_ctx(),
-        )
-    };
-
-    let cfg = {
-        let r = unsafe { &*region_ptrs[0] };
-        FusedConfig {
-            input_bits: input_bits as u32,
-            n_columns: r.sp_gpu.n_columns_accessor() as u32,
-            synapses_per_col: r.sp_gpu.synapses_per_col_accessor() as u32,
-            conn_thr: r.sp_gpu.conn_thr_accessor(),
-            sp_inc: r.sp_gpu.inc_accessor(),
-            sp_dec: r.sp_gpu.dec_accessor(),
-            sparsity_target: r.sp_gpu.sparsity_accessor(),
-            duty_alpha: 1.0f32 / r.sp_gpu.duty_period_accessor().max(1.0),
-            thr_adapt_rate: 0.001f32,
-            cells_per_column: r.tm_gpu.cells_per_column as u32,
-            n_cells: r.tm_gpu.n_cells as u32,
-            bits_words: r.tm_gpu.bits_words as u32,
-            max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
-            synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
-            activation_threshold: r.tm_gpu.activation_threshold,
-            learning_threshold: r.tm_gpu.learning_threshold,
-            max_new_synapses: r.tm_gpu.max_new_synapse_count,
-            conn_thr_i16: r.tm_gpu.conn_thr_i16 as i32,
-            perm_inc_i16: r.tm_gpu.perm_inc_i16 as i32,
-            perm_dec_i16: r.tm_gpu.perm_dec_i16 as i32,
-            predicted_seg_dec_i16: r.tm_gpu.predicted_seg_dec_i16 as i32,
-            initial_perm_i16: r.tm_gpu.initial_perm_i16 as i32,
-            t: t as u32,
-            learn: if learn { 1 } else { 0 },
-            iter_seed: r.fused_state.iter_counter,
-            cooperative_grid_sync: 1,
-        }
-    };
-
-    // Build B FusedPtrs per-region.
-    let ptrs_vec: Vec<FusedPtrs> = (0..b)
-        .map(|i| {
-            let r = unsafe { &*region_ptrs[i] };
-            FusedPtrs {
-                syn_bit: *r.sp_gpu.syn_bit_accessor().device_ptr(),
-                syn_perm: *r.sp_gpu.syn_perm_accessor().device_ptr(),
-                boost: *r.sp_gpu.boost_accessor().device_ptr(),
-                active_duty: *r.sp_gpu.active_duty_accessor().device_ptr(),
-                inhibition_threshold: *r.fused_state.inhibition_threshold.device_ptr(),
-                seg_cell_id: *r.tm_gpu.seg_cell_id_accessor().device_ptr(),
-                seg_syn_count: *r.tm_gpu.seg_syn_count_accessor().device_ptr(),
-                syn_presyn: *r.tm_gpu.syn_presyn_accessor().device_ptr(),
-                tm_syn_perm: *r.tm_gpu.syn_perm_accessor().device_ptr(),
-                cell_seg_count: *r.tm_gpu.cell_seg_count_accessor().device_ptr(),
-                cell_active_a: *r.fused_state.cell_active_bits_a.device_ptr(),
-                cell_active_b: *r.fused_state.cell_active_bits_b.device_ptr(),
-                cell_winner_a: *r.fused_state.cell_winner_bits_a.device_ptr(),
-                cell_winner_b: *r.fused_state.cell_winner_bits_b.device_ptr(),
-                inputs: inputs_per_region[i],
-                cols_out: cols_per_region[i],
-                anom_out: anom_per_region[i],
-                barrier_counters: 0u64, // ABI-compat dummy; cluster barrier replaces DLB.
-                step_scratch: *r.fused_state.step_scratch.device_ptr(),
-            }
-        })
-        .collect();
-
-    // Upload FusedPtrs array to device (B * sizeof(FusedPtrs) bytes).
-    // FusedPtrs is repr(C) + DeviceRepr so htod_sync_copy handles it.
-    let dev = unsafe { &*region_ptrs[0] }.sp_gpu.dev_ref().clone();
-    let ptrs_dev: CudaSlice<FusedPtrs> = dev.htod_sync_copy(&ptrs_vec)?;
-    let ptrs_dev_ptr: u64 = *ptrs_dev.device_ptr();
-
-    // T10: Cluster launch for batched regions.
-    // Grid = (grid_x, B, 1) with cluster_dim=(16,1,1): each region (Y slice)
-    // occupies exactly one cluster of 16 blocks. All 8 clusters run concurrently
-    // on the H200's 132 SMs (8 × 16 = 128 blocks ≤ 132 SMs).
-    let use_cluster = {
-        let r0 = unsafe { &*region_ptrs[0] };
-        r0.fused_state.cluster_info.max_cluster_size > 0
-    };
-    let grid_x =
-        plan_batched_grid_dim(grid_x, cooperative_grid_limit, b, use_cluster).map_err(|msg| {
-            eprintln!("[htm_rust] FATAL: {msg}");
-            DriverError(cudarc::driver::sys::CUresult::CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE)
-        })?;
-
-    unsafe {
-        result::ctx::set_current(cu_ctx)?;
-        let mut kernel_params: [*mut std::ffi::c_void; 2] = [
-            (&ptrs_dev_ptr as *const u64).cast_mut().cast(),
-            (&cfg as *const FusedConfig).cast_mut().cast(),
-        ];
-
-        if use_cluster {
-            let mut attr: sys::CUlaunchAttribute = std::mem::zeroed();
-            attr.id = sys::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
-            attr.value.clusterDim.x = 16;
-            attr.value.clusterDim.y = 1;
-            attr.value.clusterDim.z = 1;
-
-            let mut launch_cfg: sys::CUlaunchConfig = std::mem::zeroed();
-            launch_cfg.gridDimX = grid_x;
-            launch_cfg.gridDimY = b as u32;
-            launch_cfg.gridDimZ = 1;
-            launch_cfg.blockDimX = block_x;
-            launch_cfg.blockDimY = 1;
-            launch_cfg.blockDimZ = 1;
-            launch_cfg.sharedMemBytes = 0;
-            launch_cfg.hStream = cu_stream;
-            launch_cfg.numAttrs = 1;
-            launch_cfg.attrs = &mut attr as *mut sys::CUlaunchAttribute;
-
-            let ret = sys::lib().cuLaunchKernelEx(
-                &launch_cfg as *const sys::CUlaunchConfig,
-                function_batched,
-                kernel_params.as_mut_ptr(),
-                std::ptr::null_mut(),
-            );
-            if ret != sys::CUresult::CUDA_SUCCESS {
-                return Err(DriverError(ret));
-            }
-        } else {
-            // Pre-Hopper: cooperative kernel launch. The fused kernel uses
-            // cg::this_grid().sync(), which is only valid under cooperative
-            // launch. A normal launch can run until the first grid.sync() and
-            // then poison the CUDA context with an unspecified launch failure.
-            let ret = sys::lib().cuLaunchCooperativeKernel(
-                function_batched,
-                grid_x,
-                b as u32,
-                1,
-                block_x,
-                1,
-                1,
-                0,
-                cu_stream,
-                kernel_params.as_mut_ptr(),
-            );
-            if ret != sys::CUresult::CUDA_SUCCESS {
-                return Err(DriverError(ret));
-            }
-        }
-    }
-
-    // `ptrs_dev` is a per-call device array consumed by the async kernel.
-    // Keep it alive until the kernel has read it; otherwise dropping/freeing
-    // it immediately after launch can surface as a later unrelated CUDA error.
-    dev.synchronize()?;
-
-    Ok(())
-}
+//! Fused HTM megakernel launcher.
+//!
+//! Collapses the 12-kernel per-timestep pipeline (and the outer T-loop) into
+//! a single kernel launch per forward. See `kernels/htm_fused_step.cu` for
+//! the kernel design and the cross-block coherence strategy (grid barrier
+//! via device counter with all blocks concurrently resident).
+//!
+//! Launch invariant: `grid_dim.x <= concurrent-block capacity`. Host code
+//! probes the device SM count at construction and caps grid_dim.x
+//! accordingly — otherwise the grid barrier deadlocks.
+//!
+//! Semantic change from the top-K pipeline: activation is per-column
+//! threshold-based (local lateral inhibition) instead of global top-K.
+//! A per-column `inhibition_threshold` is tracked and EMA-steered to hit
+//! the sparsity target. This is a real architectural change and is
+//! documented in `docs/GPU_HTM.md`.
+
+#![cfg(feature = "gpu")]
+
+use std::ffi::CString;
+use std::sync::Arc;
+
+use cudarc::driver::{
+    result, sys, CudaDevice, CudaSlice, DevicePtr, DeviceRepr, DriverError, LaunchConfig,
+};
+use cudarc::nvrtc::Ptx;
+
+use super::sp_gpu::SpatialPoolerGpu;
+use super::tm_gpu::{TemporalMemoryGpu, MAX_SEGMENTS_PER_CELL, MAX_SYN_PER_SEGMENT};
+
+const PTX_HTM_FUSED: &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/htm_fused_step.ptx"));
+
+/// Struct-by-value pointer pack — matches C-side `FusedPtrs`.
+///
+/// NOTE: `barrier_counters` is kept as an ABI-compat dummy (always 0). The
+/// C-side `FusedPtrs` still has the field at the same byte offset; removing
+/// it here would shift all subsequent fields and break the layout. Worker A
+/// will eventually delete the field from both sides once the kernel is
+/// updated; until then we zero it.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct FusedPtrs {
+    pub syn_bit: u64,
+    pub syn_perm: u64,
+    pub boost: u64,
+    pub active_duty: u64,
+    pub inhibition_threshold: u64,
+    pub seg_cell_id: u64,
+    pub seg_syn_count: u64,
+    pub syn_presyn: u64,
+    pub tm_syn_perm: u64,
+    pub cell_seg_count: u64,
+    pub cell_active_a: u64,
+    pub cell_active_b: u64,
+    pub cell_winner_a: u64,
+    pub cell_winner_b: u64,
+    pub inputs: u64,
+    pub cols_out: u64,
+    pub anom_out: u64,
+    /// ABI-compat dummy — always 0. No device memory is allocated for this
+    /// field; the cluster barrier replaces the old software DLB barrier.
+    pub barrier_counters: u64,
+    pub step_scratch: u64,
+}
+
+unsafe impl DeviceRepr for FusedPtrs {}
+
+/// Launch-time config — matches C-side `FusedConfig` 1:1.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct FusedConfig {
+    pub input_bits: u32,
+    pub n_columns: u32,
+    pub synapses_per_col: u32,
+    pub conn_thr: f32,
+    pub sp_inc: f32,
+    pub sp_dec: f32,
+    pub sparsity_target: f32,
+    pub duty_alpha: f32,
+    pub thr_adapt_rate: f32,
+    pub cells_per_column: u32,
+    pub n_cells: u32,
+    pub bits_words: u32,
+    pub max_segments_per_cell: u32,
+    pub synapses_per_segment: u32,
+    pub activation_threshold: u32,
+    pub learning_threshold: u32,
+    pub max_new_synapses: u32,
+    pub conn_thr_i16: i32,
+    pub perm_inc_i16: i32,
+    pub perm_dec_i16: i32,
+    pub predicted_seg_dec_i16: i32,
+    pub initial_perm_i16: i32,
+    pub t: u32,
+    pub learn: u32,
+    pub iter_seed: u32,
+    pub cooperative_grid_sync: u32,
+}
+
+unsafe impl DeviceRepr for FusedConfig {}
+
+/// Cluster launch parameters probed at construction time.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct ClusterInfo {
+    /// Maximum cluster size supported by this device (0 = cluster unsupported).
+    pub max_cluster_size: u32,
+}
+
+// There is only ONE launch mode: non-cooperative launch with Hopper Thread
+// Block Cluster attribute (`CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION`). The old
+// software DLB barrier and the cooperative-launch path are both removed.
+// Cluster barriers replace both.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct FusedLaunchPlan {
+    pub grid_dim_x: u32,
+    pub block_dim_x: u32,
+    pub cooperative_grid_limit: u32,
+    pub sm_count: u32,
+}
+
+fn fused_grid_cap_override() -> Option<u32> {
+    std::env::var("HTM_FUSED_GRID_CAP")
+        .ok()
+        .and_then(|s| s.parse::<u32>().ok())
+        .map(|v| v.max(1))
+}
+
+pub(crate) fn plan_fused_launch(
+    sm_count: u32,
+    cooperative_supported: bool,
+    cooperative_grid_limit: u32,
+    grid_cap_override: Option<u32>,
+) -> Result<FusedLaunchPlan, String> {
+    let sm_count = sm_count.max(1);
+    // 1024 threads/block exceeds the register file on Ampere and makes the
+    // cooperative-grid residency probe lie when the launch uses a different
+    // block size. Keep the planned block size identical to the occupancy probe.
+    let block_dim_x = 256u32;
+
+    // Cluster launch path: cooperative launch is not required. Keep the probe
+    // result for residency estimation only.
+    if !cooperative_supported {
+        eprintln!("[htm_rust] INFO: cooperative launch unsupported; cluster path only.");
+    }
+
+    // Cluster constraint: grid_dim_x must equal the cluster size (16) so that
+    // each region maps to exactly one cluster. `HTM_FUSED_GRID_CAP` can lower
+    // this for debugging but should not exceed 16 for cluster correctness.
+    let default_grid_cap = 16u32;
+    let grid_cap = grid_cap_override.unwrap_or(default_grid_cap).min(16);
+    let resident_bound = if cooperative_grid_limit > 0 {
+        cooperative_grid_limit.max(sm_count * 2)
+    } else {
+        sm_count * 2
+    };
+    Ok(FusedLaunchPlan {
+        grid_dim_x: resident_bound.min(grid_cap).max(1),
+        block_dim_x,
+        cooperative_grid_limit: resident_bound,
+        sm_count,
+    })
+}
+
+pub(crate) fn plan_batched_grid_dim(
+    grid_dim_x: u32,
+    cooperative_grid_limit: u32,
+    batch_regions: usize,
+    use_cluster: bool,
+) -> Result<u32, String> {
+    if use_cluster {
+        return Ok(grid_dim_x.max(1));
+    }
+
+    let batch_regions = batch_regions.max(1) as u32;
+    if cooperative_grid_limit == 0 {
+        return Err("COOPERATIVE_LAUNCH_TOO_LARGE: cooperative launch limit unavailable".into());
+    }
+
+    let max_grid_x = cooperative_grid_limit / batch_regions;
+    if max_grid_x == 0 {
+        return Err(format!(
+            "COOPERATIVE_LAUNCH_TOO_LARGE: batch_regions={batch_regions} exceeds cooperative_grid_limit={cooperative_grid_limit}"
+        ));
+    }
+
+    Ok(grid_dim_x.min(max_grid_x).max(1))
+}
+
+pub(super) struct RawFusedKernel {
+    module: sys::CUmodule,
+    pub(super) function: sys::CUfunction,
+    pub(super) function_batched: sys::CUfunction,
+}
+
+unsafe impl Send for RawFusedKernel {}
+unsafe impl Sync for RawFusedKernel {}
+
+impl Drop for RawFusedKernel {
+    fn drop(&mut self) {
+        unsafe {
+            let _ = result::module::unload(self.module);
+        }
+    }
+}
+
+/// Owns fused-path-only device state:
+///   - per-column inhibition threshold (replaces global top-K)
+///   - ping-pong cell_active/cell_winner bitsets
+///   - step_scratch (n_active, n_unpred per timestep)
+///   - cluster launch capability info
+pub struct FusedState {
+    dev: Arc<CudaDevice>,
+    pub(super) raw_kernel: RawFusedKernel,
+
+    pub inhibition_threshold: CudaSlice<f32>,
+    pub cell_active_bits_a: CudaSlice<u32>,
+    pub cell_active_bits_b: CudaSlice<u32>,
+    pub cell_winner_bits_a: CudaSlice<u32>,
+    pub cell_winner_bits_b: CudaSlice<u32>,
+    pub step_scratch: CudaSlice<u32>, // length 6
+
+    pub grid_dim_x: u32,
+    pub block_dim_x: u32,
+    pub cooperative_grid_limit: u32,
+    pub iter_counter: u32,
+
+    /// Hopper cluster launch capability (0 = unsupported).
+    pub cluster_info: ClusterInfo,
+
+    // Config mirror (read-only after init).
+    #[allow(dead_code)]
+    pub initial_threshold: f32,
+}
+
+impl FusedState {
+    pub fn new(
+        dev: Arc<CudaDevice>,
+        n_columns: usize,
+        cells_per_column: usize,
+        initial_threshold: f32,
+    ) -> Result<Self, DriverError> {
+        let n_cells = n_columns * cells_per_column;
+        assert!(
+            n_cells % 32 == 0,
+            "n_cells must be divisible by 32 for bitsets"
+        );
+        let bits_words = n_cells / 32;
+
+        let mut inhibition_threshold = dev.alloc_zeros::<f32>(n_columns)?;
+        let init_vec = vec![initial_threshold; n_columns];
+        dev.htod_sync_copy_into(&init_vec, &mut inhibition_threshold)?;
+
+        let cell_active_bits_a = dev.alloc_zeros::<u32>(bits_words)?;
+        let cell_active_bits_b = dev.alloc_zeros::<u32>(bits_words)?;
+        let cell_winner_bits_a = dev.alloc_zeros::<u32>(bits_words)?;
+        let cell_winner_bits_b = dev.alloc_zeros::<u32>(bits_words)?;
+        let step_scratch = dev.alloc_zeros::<u32>(6)?;
+
+        unsafe {
+            result::ctx::set_current(*dev.cu_primary_ctx())?;
+        }
+        if dev.get_func("htm_fused", "htm_fused_step").is_none() {
+            dev.load_ptx(
+                Ptx::from_src(PTX_HTM_FUSED),
+                "htm_fused",
+                &["htm_fused_step", "htm_fused_step_batched"],
+            )?;
+        }
+        let ptx = CString::new(PTX_HTM_FUSED).expect("PTX contains no interior nul bytes");
+        let module = unsafe { result::module::load_data(ptx.as_ptr().cast()) }?;
+        let function = unsafe {
+            result::module::get_function(module, CString::new("htm_fused_step").unwrap())
+        }?;
+        let function_batched = unsafe {
+            result::module::get_function(module, CString::new("htm_fused_step_batched").unwrap())
+        }?;
+
+        // Cluster size 16 on Hopper is "non-portable" (> 8 requires opt-in).
+        // Must set CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED=1 on
+        // every launched kernel function, otherwise cuLaunchKernelEx rejects
+        // the cluster dim with CUDA_ERROR_INVALID_CLUSTER_SIZE.
+        unsafe {
+            let attr =
+                sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED;
+            // Ignore errors: older CUDA may lack the attribute, in which case
+            // only portable sizes (<= 8) work — plan_fused_launch caps at 8.
+            let _ = sys::lib().cuFuncSetAttribute(function, attr, 1);
+            let _ = sys::lib().cuFuncSetAttribute(function_batched, attr, 1);
+        }
+
+        // Probe SM count.
+        let sm_count = match dev.attribute(
+            cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
+        ) {
+            Ok(v) => v as u32,
+            Err(_) => 16u32,
+        };
+
+        // T1: Probe Hopper cluster launch capability.
+        let max_cluster_size = match dev
+            .attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_CLUSTER_LAUNCH)
+        {
+            Ok(v) if v > 0 => {
+                // H200/sm_90a supports up to 16 blocks per cluster.
+                // There is no MAX_CLUSTER_SIZE attribute in CUDA 12.4; hard-code the
+                // Hopper maximum which is 16 (8 SMs × 2 blocks/SM = 16 blocks/cluster).
+                16u32
+            }
+            _ => 0u32,
+        };
+        eprintln!("[htm_rust] cluster: max_cluster_size={}", max_cluster_size);
+        let cluster_info = ClusterInfo { max_cluster_size };
+
+        let cooperative_supported = matches!(
+            dev.attribute(sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH),
+            Ok(v) if v > 0
+        );
+        let cooperative_grid_limit = if cooperative_supported {
+            let blocks_per_sm = unsafe {
+                // Must match plan_fused_launch(): the A10G/Ampere-safe fused
+                // kernel launch uses 256 threads/block, not the historical
+                // 1024-thread Hopper occupancy probe.
+                result::occupancy::max_active_block_per_multiprocessor(function, 256, 0)
+            }
+            .ok()
+            .map(|v| v.max(0) as u32)
+            .unwrap_or(0);
+            sm_count.saturating_mul(blocks_per_sm)
+        } else {
+            0
+        };
+        let launch_plan = plan_fused_launch(
+            sm_count,
+            cooperative_supported,
+            cooperative_grid_limit,
+            fused_grid_cap_override(),
+        )
+        .map_err(|msg| {
+            // Surface as a CUDA-ish error so callers can propagate.
+            eprintln!("[htm_rust] FATAL: {msg}");
+            DriverError(cudarc::driver::sys::CUresult::CUDA_ERROR_NOT_SUPPORTED)
+        })?;
+
+        eprintln!(
+            "[htm_rust] fused kernel: sm_count={} grid_dim_x={} cooperative_grid_limit={} cluster_max={}",
+            launch_plan.sm_count, launch_plan.grid_dim_x, launch_plan.cooperative_grid_limit,
+            cluster_info.max_cluster_size,
+        );
+
+        Ok(Self {
+            dev,
+            raw_kernel: RawFusedKernel {
+                module,
+                function,
+                function_batched,
+            },
+            inhibition_threshold,
+            cell_active_bits_a,
+            cell_active_bits_b,
+            cell_winner_bits_a,
+            cell_winner_bits_b,
+            step_scratch,
+            grid_dim_x: launch_plan.grid_dim_x,
+            block_dim_x: launch_plan.block_dim_x,
+            cooperative_grid_limit: launch_plan.cooperative_grid_limit,
+            iter_counter: 0,
+            cluster_info,
+            initial_threshold,
+        })
+    }
+
+    /// Reset fused state. Called at region.reset().
+    pub fn reset(&mut self) -> Result<(), DriverError> {
+        self.dev.memset_zeros(&mut self.cell_active_bits_a)?;
+        self.dev.memset_zeros(&mut self.cell_active_bits_b)?;
+        self.dev.memset_zeros(&mut self.cell_winner_bits_a)?;
+        self.dev.memset_zeros(&mut self.cell_winner_bits_b)?;
+        self.dev.memset_zeros(&mut self.step_scratch)?;
+        // Do NOT reset inhibition_threshold — it's learned state. A hard
+        // reset of TM state should NOT forget the sparsity calibration.
+        Ok(())
+    }
+}
+
+/// Launch the fused megakernel. Processes all T timesteps in one kernel.
+///
+/// Uses `cuLaunchKernelEx` with `CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION=(16,1,1)`
+/// when the device supports cluster launch, otherwise falls back to a plain
+/// `launch_kernel`. For single-region launches, grid_dim_x <= 16 ensures the
+/// entire grid fits in one cluster.
+#[allow(clippy::too_many_arguments)]
+pub fn launch_fused(
+    sp: &mut SpatialPoolerGpu,
+    tm: &mut TemporalMemoryGpu,
+    fused: &mut FusedState,
+    inputs_flat: &CudaSlice<u8>,
+    cols_out: &mut CudaSlice<u8>,
+    anom_out: &mut CudaSlice<f32>,
+    t: usize,
+    input_bits: usize,
+    learn: bool,
+) -> Result<(), DriverError> {
+    // Reset step_scratch before each launch (safe re-entry).
+    sp.dev_ref().memset_zeros(&mut fused.step_scratch)?;
+
+    fused.iter_counter = fused.iter_counter.wrapping_add(1);
+
+    let cfg = FusedConfig {
+        input_bits: input_bits as u32,
+        n_columns: sp.n_columns_accessor() as u32,
+        synapses_per_col: sp.synapses_per_col_accessor() as u32,
+        conn_thr: sp.conn_thr_accessor(),
+        sp_inc: sp.inc_accessor(),
+        sp_dec: sp.dec_accessor(),
+        sparsity_target: sp.sparsity_accessor(),
+        duty_alpha: 1.0f32 / sp.duty_period_accessor().max(1.0),
+        thr_adapt_rate: 0.001f32,
+        cells_per_column: tm.cells_per_column as u32,
+        n_cells: tm.n_cells as u32,
+        bits_words: tm.bits_words as u32,
+        max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
+        synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
+        activation_threshold: tm.activation_threshold,
+        learning_threshold: tm.learning_threshold,
+        max_new_synapses: tm.max_new_synapse_count,
+        conn_thr_i16: tm.conn_thr_i16 as i32,
+        perm_inc_i16: tm.perm_inc_i16 as i32,
+        perm_dec_i16: tm.perm_dec_i16 as i32,
+        predicted_seg_dec_i16: tm.predicted_seg_dec_i16 as i32,
+        initial_perm_i16: tm.initial_perm_i16 as i32,
+        t: t as u32,
+        learn: if learn { 1 } else { 0 },
+        iter_seed: fused.iter_counter,
+        cooperative_grid_sync: 1,
+    };
+
+    let ptrs = FusedPtrs {
+        syn_bit: *sp.syn_bit_accessor().device_ptr(),
+        syn_perm: *sp.syn_perm_accessor().device_ptr(),
+        boost: *sp.boost_accessor().device_ptr(),
+        active_duty: *sp.active_duty_accessor().device_ptr(),
+        inhibition_threshold: *fused.inhibition_threshold.device_ptr(),
+        seg_cell_id: *tm.seg_cell_id_accessor().device_ptr(),
+        seg_syn_count: *tm.seg_syn_count_accessor().device_ptr(),
+        syn_presyn: *tm.syn_presyn_accessor().device_ptr(),
+        tm_syn_perm: *tm.syn_perm_accessor().device_ptr(),
+        cell_seg_count: *tm.cell_seg_count_accessor().device_ptr(),
+        cell_active_a: *fused.cell_active_bits_a.device_ptr(),
+        cell_active_b: *fused.cell_active_bits_b.device_ptr(),
+        cell_winner_a: *fused.cell_winner_bits_a.device_ptr(),
+        cell_winner_b: *fused.cell_winner_bits_b.device_ptr(),
+        inputs: *inputs_flat.device_ptr(),
+        cols_out: *cols_out.device_ptr(),
+        anom_out: *anom_out.device_ptr(),
+        barrier_counters: 0u64, // ABI-compat dummy; cluster barrier replaces DLB.
+        step_scratch: *fused.step_scratch.device_ptr(),
+    };
+
+    let grid_x = fused.grid_dim_x;
+    let block_x = fused.block_dim_x;
+    let cu_stream = *sp.dev_ref().cu_stream();
+    let use_cluster = fused.cluster_info.max_cluster_size > 0;
+
+    unsafe {
+        result::ctx::set_current(*sp.dev_ref().cu_primary_ctx())?;
+        let mut kernel_params: [*mut std::ffi::c_void; 2] = [
+            (&ptrs as *const FusedPtrs).cast_mut().cast(),
+            (&cfg as *const FusedConfig).cast_mut().cast(),
+        ];
+
+        if use_cluster {
+            // T10: Hopper cluster launch with CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION.
+            // cluster_dim=(16,1,1) maps the entire single-region grid into one cluster.
+            let mut attr: sys::CUlaunchAttribute = std::mem::zeroed();
+            attr.id = sys::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
+            attr.value.clusterDim.x = 16;
+            attr.value.clusterDim.y = 1;
+            attr.value.clusterDim.z = 1;
+
+            let mut launch_cfg: sys::CUlaunchConfig = std::mem::zeroed();
+            launch_cfg.gridDimX = grid_x;
+            launch_cfg.gridDimY = 1;
+            launch_cfg.gridDimZ = 1;
+            launch_cfg.blockDimX = block_x;
+            launch_cfg.blockDimY = 1;
+            launch_cfg.blockDimZ = 1;
+            launch_cfg.sharedMemBytes = 0;
+            launch_cfg.hStream = cu_stream;
+            launch_cfg.numAttrs = 1;
+            launch_cfg.attrs = &mut attr as *mut sys::CUlaunchAttribute;
+
+            let ret = sys::lib().cuLaunchKernelEx(
+                &launch_cfg as *const sys::CUlaunchConfig,
+                fused.raw_kernel.function,
+                kernel_params.as_mut_ptr(),
+                std::ptr::null_mut(),
+            );
+            if ret != sys::CUresult::CUDA_SUCCESS {
+                return Err(DriverError(ret));
+            }
+        } else {
+            // Pre-Hopper: cooperative kernel launch. The fused kernel uses
+            // cg::this_grid().sync(); normal launches poison the CUDA context
+            // with an asynchronous unspecified launch failure.
+            let ret = sys::lib().cuLaunchCooperativeKernel(
+                fused.raw_kernel.function,
+                grid_x,
+                1,
+                1,
+                block_x,
+                1,
+                1,
+                0,
+                cu_stream,
+                kernel_params.as_mut_ptr(),
+            );
+            if ret != sys::CUresult::CUDA_SUCCESS {
+                return Err(DriverError(ret));
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Single batched non-cooperative launch for B regions with DLB sync. Uses the same kernel
+/// body; each block reads its region's FusedPtrs from a device-side array
+/// indexed by blockIdx.y. All regions share the same config (same
+/// input_bits/n_columns/etc.) so we pass one FusedConfig.
+///
+/// This breaks through the CUDA cooperative-kernel device-level
+/// serialization: multiple cooperative launches are serialized regardless
+/// of stream, but one cooperative launch with grid.y=B processes all
+/// regions in a single invocation — ~B× speedup vs B sequential launches.
+#[allow(clippy::too_many_arguments)]
+/// Low-level raw-pointer entry, called by PyO3 binding which holds the
+/// mutable borrows. Safety: each `*mut HTMRegionGpu` must point to a live,
+/// uniquely-borrowed region. All regions must be distinct.
+pub(super) fn launch_fused_batched_raw(
+    region_ptrs: &[*mut super::HTMRegionGpu],
+    inputs_per_region: &[u64],
+    cols_per_region: &[u64],
+    anom_per_region: &[u64],
+    t: usize,
+    input_bits: usize,
+    learn: bool,
+) -> Result<(), DriverError> {
+    let b = region_ptrs.len();
+    assert_eq!(inputs_per_region.len(), b);
+    assert_eq!(cols_per_region.len(), b);
+    assert_eq!(anom_per_region.len(), b);
+    assert!(b >= 1, "need at least one region");
+
+    // Reset per-region step_scratch before each launch.
+    for &rp in region_ptrs.iter() {
+        let r = unsafe { &mut *rp };
+        let dev = r.sp_gpu.dev_ref().clone();
+        dev.memset_zeros(&mut r.fused_state.step_scratch)?;
+        r.fused_state.iter_counter = r.fused_state.iter_counter.wrapping_add(1);
+    }
+
+    // Shared config — all regions use identical sp/tm parameters.
+    let (grid_x, block_x, cooperative_grid_limit, function_batched, cu_stream, cu_ctx) = {
+        let r0 = unsafe { &*region_ptrs[0] };
+        (
+            r0.fused_state.grid_dim_x,
+            r0.fused_state.block_dim_x,
+            r0.fused_state.cooperative_grid_limit,
+            r0.fused_state.raw_kernel.function_batched,
+            *r0.sp_gpu.dev_ref().cu_stream(),
+            *r0.sp_gpu.dev_ref().cu_primary_ctx(),
+        )
+    };
+
+    let cfg = {
+        let r = unsafe { &*region_ptrs[0] };
+        FusedConfig {
+            input_bits: input_bits as u32,
+            n_columns: r.sp_gpu.n_columns_accessor() as u32,
+            synapses_per_col: r.sp_gpu.synapses_per_col_accessor() as u32,
+            conn_thr: r.sp_gpu.conn_thr_accessor(),
+            sp_inc: r.sp_gpu.inc_accessor(),
+            sp_dec: r.sp_gpu.dec_accessor(),
+            sparsity_target: r.sp_gpu.sparsity_accessor(),
+            duty_alpha: 1.0f32 / r.sp_gpu.duty_period_accessor().max(1.0),
+            thr_adapt_rate: 0.001f32,
+            cells_per_column: r.tm_gpu.cells_per_column as u32,
+            n_cells: r.tm_gpu.n_cells as u32,
+            bits_words: r.tm_gpu.bits_words as u32,
+            max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
+            synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
+            activation_threshold: r.tm_gpu.activation_threshold,
+            learning_threshold: r.tm_gpu.learning_threshold,
+            max_new_synapses: r.tm_gpu.max_new_synapse_count,
+            conn_thr_i16: r.tm_gpu.conn_thr_i16 as i32,
+            perm_inc_i16: r.tm_gpu.perm_inc_i16 as i32,
+            perm_dec_i16: r.tm_gpu.perm_dec_i16 as i32,
+            predicted_seg_dec_i16: r.tm_gpu.predicted_seg_dec_i16 as i32,
+            initial_perm_i16: r.tm_gpu.initial_perm_i16 as i32,
+            t: t as u32,
+            learn: if learn { 1 } else { 0 },
+            iter_seed: r.fused_state.iter_counter,
+            cooperative_grid_sync: 1,
+        }
+    };
+
+    // Build B FusedPtrs per-region.
+    let ptrs_vec: Vec<FusedPtrs> = (0..b)
+        .map(|i| {
+            let r = unsafe { &*region_ptrs[i] };
+            FusedPtrs {
+                syn_bit: *r.sp_gpu.syn_bit_accessor().device_ptr(),
+                syn_perm: *r.sp_gpu.syn_perm_accessor().device_ptr(),
+                boost: *r.sp_gpu.boost_accessor().device_ptr(),
+                active_duty: *r.sp_gpu.active_duty_accessor().device_ptr(),
+                inhibition_threshold: *r.fused_state.inhibition_threshold.device_ptr(),
+                seg_cell_id: *r.tm_gpu.seg_cell_id_accessor().device_ptr(),
+                seg_syn_count: *r.tm_gpu.seg_syn_count_accessor().device_ptr(),
+                syn_presyn: *r.tm_gpu.syn_presyn_accessor().device_ptr(),
+                tm_syn_perm: *r.tm_gpu.syn_perm_accessor().device_ptr(),
+                cell_seg_count: *r.tm_gpu.cell_seg_count_accessor().device_ptr(),
+                cell_active_a: *r.fused_state.cell_active_bits_a.device_ptr(),
+                cell_active_b: *r.fused_state.cell_active_bits_b.device_ptr(),
+                cell_winner_a: *r.fused_state.cell_winner_bits_a.device_ptr(),
+                cell_winner_b: *r.fused_state.cell_winner_bits_b.device_ptr(),
+                inputs: inputs_per_region[i],
+                cols_out: cols_per_region[i],
+                anom_out: anom_per_region[i],
+                barrier_counters: 0u64, // ABI-compat dummy; cluster barrier replaces DLB.
+                step_scratch: *r.fused_state.step_scratch.device_ptr(),
+            }
+        })
+        .collect();
+
+    // Upload FusedPtrs array to device (B * sizeof(FusedPtrs) bytes).
+    // FusedPtrs is repr(C) + DeviceRepr so htod_sync_copy handles it.
+    let dev = unsafe { &*region_ptrs[0] }.sp_gpu.dev_ref().clone();
+    let ptrs_dev: CudaSlice<FusedPtrs> = dev.htod_sync_copy(&ptrs_vec)?;
+    let ptrs_dev_ptr: u64 = *ptrs_dev.device_ptr();
+
+    // T10: Cluster launch for batched regions.
+    // Grid = (grid_x, B, 1) with cluster_dim=(16,1,1): each region (Y slice)
+    // occupies exactly one cluster of 16 blocks. All 8 clusters run concurrently
+    // on the H200's 132 SMs (8 × 16 = 128 blocks ≤ 132 SMs).
+    let use_cluster = {
+        let r0 = unsafe { &*region_ptrs[0] };
+        r0.fused_state.cluster_info.max_cluster_size > 0
+    };
+    let grid_x =
+        plan_batched_grid_dim(grid_x, cooperative_grid_limit, b, use_cluster).map_err(|msg| {
+            eprintln!("[htm_rust] FATAL: {msg}");
+            DriverError(cudarc::driver::sys::CUresult::CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE)
+        })?;
+
+    unsafe {
+        result::ctx::set_current(cu_ctx)?;
+        let mut kernel_params: [*mut std::ffi::c_void; 2] = [
+            (&ptrs_dev_ptr as *const u64).cast_mut().cast(),
+            (&cfg as *const FusedConfig).cast_mut().cast(),
+        ];
+
+        if use_cluster {
+            let mut attr: sys::CUlaunchAttribute = std::mem::zeroed();
+            attr.id = sys::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
+            attr.value.clusterDim.x = 16;
+            attr.value.clusterDim.y = 1;
+            attr.value.clusterDim.z = 1;
+
+            let mut launch_cfg: sys::CUlaunchConfig = std::mem::zeroed();
+            launch_cfg.gridDimX = grid_x;
+            launch_cfg.gridDimY = b as u32;
+            launch_cfg.gridDimZ = 1;
+            launch_cfg.blockDimX = block_x;
+            launch_cfg.blockDimY = 1;
+            launch_cfg.blockDimZ = 1;
+            launch_cfg.sharedMemBytes = 0;
+            launch_cfg.hStream = cu_stream;
+            launch_cfg.numAttrs = 1;
+            launch_cfg.attrs = &mut attr as *mut sys::CUlaunchAttribute;
+
+            let ret = sys::lib().cuLaunchKernelEx(
+                &launch_cfg as *const sys::CUlaunchConfig,
+                function_batched,
+                kernel_params.as_mut_ptr(),
+                std::ptr::null_mut(),
+            );
+            if ret != sys::CUresult::CUDA_SUCCESS {
+                return Err(DriverError(ret));
+            }
+        } else {
+            // Pre-Hopper: cooperative kernel launch. The fused kernel uses
+            // cg::this_grid().sync(), which is only valid under cooperative
+            // launch. A normal launch can run until the first grid.sync() and
+            // then poison the CUDA context with an unspecified launch failure.
+            let ret = sys::lib().cuLaunchCooperativeKernel(
+                function_batched,
+                grid_x,
+                b as u32,
+                1,
+                block_x,
+                1,
+                1,
+                0,
+                cu_stream,
+                kernel_params.as_mut_ptr(),
+            );
+            if ret != sys::CUresult::CUDA_SUCCESS {
+                return Err(DriverError(ret));
+            }
+        }
+    }
+
+    // `ptrs_dev` is a per-call device array consumed by the async kernel.
+    // Keep it alive until the kernel has read it; otherwise dropping/freeing
+    // it immediately after launch can surface as a later unrelated CUDA error.
+    dev.synchronize()?;
+
+    Ok(())
+}

overlay/htm_rust/src/gpu/kernels/htm_fused_step.cu

@@ -1,677 +1,677 @@
-// Fused HTM megakernel — SP + TM, all T timesteps in a single launch.
-//
-// Design rationale:
-//   - Global top-K column selection requires cross-block synchronization at
-//     every timestep (grid.sync is unreliable on WSL2/sm_86 without rdc=true).
-//   - Replace with per-column threshold activation using local lateral
-//     inhibition: column c activates if overlap[c]*boost[c] > threshold[c].
-//     Threshold is a per-column running-EMA learned scalar that steers the
-//     column's long-run activation rate toward the global sparsity target.
-//   - This is biologically grounded (GABAergic local inhibition) and supported
-//     by HTM theory (duty-cycle boost already drives this loop; we just
-//     change which lever the EMA pulls).
-//
-// Launch shape:
-//   grid  = min(device SM count, 16)  // hard cap — see below
-//   block = 1024 threads = 32 warps
-//   Each warp of 32 owns a contiguous column slice (n_columns / total_warps).
-//
-// Cross-block coherence:
-//   - Ping-pong buffers for cell_active/cell_winner: write _a at even t,
-//     read _b; reversed at odd t.
-//   - Preferred path: cooperative launch + hardware whole-grid sync.
-//   - Fallback path: software 3-slot rotating grid barrier for devices/drivers
-//     that cannot do cooperative launch.
-//
-// 2026-04-16: grid_dim reduced from 28 to 16 after deadlock RCA. The previous
-// cap of 28 relied on all blocks being concurrently resident on a 30-SM RTX
-// 3060 Laptop. Under thermal throttling effective residency dropped to ~20-24,
-// leaving scheduled blocks spinning on the software grid barrier waiting for
-// peer blocks that would never run. 16 blocks is below any realistic residency
-// floor and preserves enough warp parallelism (16*32 = 512 warps) to saturate
-// memory bandwidth on the spatial-pooler stage.
-//
-// Kernel signature uses struct-by-value for pointers and config to stay
-// inside cudarc's launch-arg count limit.
-
-#include <cooperative_groups.h>
-#include <cooperative_groups/memcpy_async.h>
-
-namespace cg = cooperative_groups;
-
-// Maximum columns owned per cluster-block in DSMEM.
-// Supports n_columns up to COLS_PER_CLUSTER_BLOCK_MAX * cluster_size.
-// At cluster_size=16: supports up to 256*16=4096 columns.
-// Each array costs 256*4 = 1024 bytes; three arrays = 3072 bytes per SM —
-// well under the 228 KB H200 shared-memory cap.
-#define COLS_PER_CLUSTER_BLOCK_MAX 256u
-
-// Maximum input_bits supported by the TMA-multicast staging tile.
-// At 32 KB this covers the production SDR width (16384 bits) with 2× headroom.
-// Total shared per SM: 32768 (tile) + 3072 (DSMEM float arrays) = ~35 KB —
-// well under the 228 KB H200 limit.
-//
-// Expected speedup from TMA multicast input staging (T9/T11):
-//   - Without staging: 16 SMs × T × (input_bits GMEM reads per timestep)
-//   - With staging:    1 TMA DMA per timestep, shared reads from L1 thereafter
-//   - Theoretical DRAM bandwidth reduction: ~16× on input reads
-//   - Wall-clock reduction estimate: -20 to -40 ms from reduced input fetch latency
-#define INPUT_BITS_MAX 32768u
-
-extern "C" {
-
-struct FusedPtrs {
-    unsigned long long syn_bit;
-    unsigned long long syn_perm;
-    unsigned long long boost;
-    unsigned long long active_duty;
-    unsigned long long inhibition_threshold;
-    unsigned long long seg_cell_id;
-    unsigned long long seg_syn_count;
-    unsigned long long syn_presyn;
-    unsigned long long tm_syn_perm;
-    unsigned long long cell_seg_count;
-    unsigned long long cell_active_a;
-    unsigned long long cell_active_b;
-    unsigned long long cell_winner_a;
-    unsigned long long cell_winner_b;
-    unsigned long long inputs;
-    unsigned long long cols_out;
-    unsigned long long anom_out;
-    unsigned long long barrier_counters;
-    unsigned long long step_scratch;
-};
-
-struct FusedConfig {
-    // SP constants
-    unsigned int input_bits;
-    unsigned int n_columns;
-    unsigned int synapses_per_col;
-    float        conn_thr;
-    float        sp_inc;
-    float        sp_dec;
-    float        sparsity_target;
-    float        duty_alpha;
-    float        thr_adapt_rate;
-    // TM constants
-    unsigned int cells_per_column;
-    unsigned int n_cells;
-    unsigned int bits_words;
-    unsigned int max_segments_per_cell;
-    unsigned int synapses_per_segment;
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int max_new_synapses;
-    int          conn_thr_i16;
-    int          perm_inc_i16;
-    int          perm_dec_i16;
-    int          predicted_seg_dec_i16;
-    int          initial_perm_i16;
-    // Loop constants
-    unsigned int T;
-    unsigned int learn;
-    unsigned int iter_seed;
-    unsigned int cooperative_grid_sync;
-};
-
-// Hardware cluster barrier using Hopper sm_90a cooperative_groups::this_cluster().sync().
-// Replaces the former software Decoupled Look-Back (DLB) atomic-spin barrier.
-//
-// cluster::sync() is a single PTX instruction (barrier.cluster) that resolves
-// in ~10-40 ns inside the cluster, with no device-level serialization.
-// Multiple clusters (one per HTM region) run fully concurrently — bounded
-// only by SM count (8 clusters × 16 SMs = 128 ≤ 132 on H200).
-//
-// The flags / expected / phase / cooperative_grid_sync parameters are kept
-// in the signature for call-site compatibility but are unused.
-__device__ static inline void fused_grid_barrier(cg::grid_group grid,
-                                                 unsigned int * /* flags — unused */,
-                                                 unsigned int /* expected — unused */,
-                                                 unsigned int /* phase — unused */,
-                                                 unsigned int /* cooperative_grid_sync — unused */) {
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    // Hopper+ : hardware cluster barrier (~10-40 ns)
-    auto cluster = cg::this_cluster();
-    cluster.sync();
-#else
-    // Pre-Hopper (sm_80, sm_86, sm_89): grid-level cooperative sync.
-    // Requires cooperative kernel launch. ~us-ms range, adequate for HTM
-    // workload (kernel launch frequency is low).
-    grid.sync();
-#endif
-}
-
-__device__ static inline unsigned int warp_sum_u32(unsigned int v) {
-    for (int off = 16; off > 0; off >>= 1) {
-        v += __shfl_down_sync(0xffffffffu, v, off);
-    }
-    return v;
-}
-
-// Core kernel body — works for both single-region and batched launches.
-// Single-region: caller passes the one FusedPtrs struct.
-// Batched: each block reads its region's FusedPtrs via blockIdx.y before
-// calling this. State is independent per region (each region owns its own
-// GPU buffers); grid.sync() is the only cross-block primitive and it
-// spans ALL blocks in the grid (harmless over-sync across regions).
-__device__ static inline
-void htm_fused_step_body(const FusedPtrs& P, const FusedConfig& cfg) {
-    cg::grid_group grid = cg::this_grid();
-    // Cast pointers.
-    const unsigned int  * __restrict__ syn_bit               = (const unsigned int*)P.syn_bit;
-    float               * __restrict__ syn_perm              = (float*)P.syn_perm;
-    float               * __restrict__ boost                 = (float*)P.boost;
-    float               * __restrict__ active_duty           = (float*)P.active_duty;
-    float               * __restrict__ inhibition_threshold  = (float*)P.inhibition_threshold;
-    unsigned int        * __restrict__ seg_cell_id           = (unsigned int*)P.seg_cell_id;
-    unsigned int        * __restrict__ seg_syn_count         = (unsigned int*)P.seg_syn_count;
-    unsigned int        * __restrict__ syn_presyn            = (unsigned int*)P.syn_presyn;
-    short               * __restrict__ tm_syn_perm           = (short*)P.tm_syn_perm;
-    unsigned int        * __restrict__ cell_seg_count        = (unsigned int*)P.cell_seg_count;
-    unsigned int        * __restrict__ cell_active_a         = (unsigned int*)P.cell_active_a;
-    unsigned int        * __restrict__ cell_active_b         = (unsigned int*)P.cell_active_b;
-    unsigned int        * __restrict__ cell_winner_a         = (unsigned int*)P.cell_winner_a;
-    unsigned int        * __restrict__ cell_winner_b         = (unsigned int*)P.cell_winner_b;
-    const unsigned char * __restrict__ inputs                = (const unsigned char*)P.inputs;
-    unsigned char       * __restrict__ cols_out              = (unsigned char*)P.cols_out;
-    float               * __restrict__ anom_out              = (float*)P.anom_out;
-    unsigned int        * __restrict__ barrier_counters      = (unsigned int*)P.barrier_counters;
-    unsigned int        * __restrict__ step_scratch          = (unsigned int*)P.step_scratch;
-
-    const unsigned int tid     = threadIdx.x;
-    const unsigned int lane    = tid & 31u;
-    const unsigned int warp    = tid >> 5;
-    const unsigned int warps_per_block = blockDim.x >> 5;
-    const unsigned int gwarp   = blockIdx.x * warps_per_block + warp;
-    const unsigned int n_warps = gridDim.x * warps_per_block;
-
-    const unsigned int n_cols  = cfg.n_columns;
-    const unsigned int col_lo  = (gwarp * n_cols) / n_warps;
-    const unsigned int col_hi  = ((gwarp + 1) * n_cols) / n_warps;
-
-    unsigned int phase = 0u;
-
-    // =========================================================
-    // DSMEM: Cluster-distributed shared memory for hot per-column
-    // state (inhibition_threshold, boost, active_duty).
-    //
-    // On Hopper (sm_90+): Each block in the cluster owns a contiguous
-    // slice of columns in its own __shared__ arrays. Any block can
-    // peer-read another block's slice via cluster.map_shared_rank().
-    //
-    // On Ampere (sm_86) and other pre-Hopper: No cluster support.
-    // Read/write directly from/to global memory (inhibition_threshold,
-    // boost, active_duty device pointers). Slightly higher latency but
-    // functionally correct.
-    // =========================================================
-
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    // Hopper+ cluster path
-    auto cluster = cg::this_cluster();
-    const unsigned int cluster_block_rank = cluster.block_rank();  // 0..cluster_size-1
-    const unsigned int cluster_sz         = cluster.num_blocks();  // == gridDim.x (≤16)
-#else
-    // Pre-Hopper: no cluster, each block is independent.
-    const unsigned int cluster_block_rank = blockIdx.x;
-    const unsigned int cluster_sz         = gridDim.x;
-#endif
-
-    // Partition n_cols evenly across cluster blocks.
-    // Each block owns cols_per_block columns starting at my_col_start.
-    const unsigned int cols_per_block =
-        (n_cols + cluster_sz - 1u) / cluster_sz;               // ceil div
-    const unsigned int my_col_start =
-        cluster_block_rank * cols_per_block;
-    const unsigned int my_col_end =
-        (my_col_start + cols_per_block < n_cols)
-            ? (my_col_start + cols_per_block) : n_cols;        // clamp
-
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    // Cluster-distributed shared memory arrays.
-    // Each block holds at most COLS_PER_CLUSTER_BLOCK_MAX floats per array.
-    // Peer blocks address into each other's smem via map_shared_rank.
-    __shared__ float s_inhib_thr [COLS_PER_CLUSTER_BLOCK_MAX];
-    __shared__ float s_boost     [COLS_PER_CLUSTER_BLOCK_MAX];
-    __shared__ float s_active_duty[COLS_PER_CLUSTER_BLOCK_MAX];
-#endif
-
-    // TMA multicast input staging tile (T9) — HOPPER ONLY.
-    //
-    // On Hopper: cg::memcpy_async with cluster scope multicasts input to all
-    // 16 SMs, reducing DRAM traffic by ~16×.
-    // On Ampere: 32 KB smem allocation exceeds per-block budget when
-    // cooperatively launched (48 KB total, registers eat the rest). Skip the
-    // tile entirely — Stage A reads from GMEM directly (original path).
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    __shared__ __align__(16) unsigned char s_input_tile[INPUT_BITS_MAX];
-#endif
-
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    // Initial GMEM → smem load (reads state from previous forward call).
-    // Each block loads only its own slice; tid strides across the slice.
-    for (unsigned int c = my_col_start + tid; c < my_col_end; c += blockDim.x) {
-        const unsigned int off = c - my_col_start;
-        s_inhib_thr [off] = inhibition_threshold[c];
-        s_boost     [off] = boost[c];
-        s_active_duty[off] = active_duty[c];
-    }
-
-    // All blocks in the cluster must finish loading before any block
-    // starts reading peer smem inside the T-loop.
-    cluster.sync();
-#else
-    // Pre-Hopper: no smem caching needed — reads go directly to GMEM.
-    // Grid sync ensures all blocks have completed Phase 0 init before T-loop.
-    grid.sync();
-#endif
-
-    const unsigned int S   = cfg.synapses_per_col;
-    const unsigned int cpc = cfg.cells_per_column;
-    const unsigned int SPS = cfg.synapses_per_segment;
-    const unsigned int MSC = cfg.max_segments_per_cell;
-
-    // Main timestep loop.
-    for (unsigned int t = 0u; t < cfg.T; t++) {
-        const unsigned int inp_off      = t * cfg.input_bits;
-        const unsigned int col_base_out = t * n_cols;
-
-        unsigned int * curr_active = (t & 1u) ? cell_active_b : cell_active_a;
-        unsigned int * prev_active = (t & 1u) ? cell_active_a : cell_active_b;
-        unsigned int * curr_winner = (t & 1u) ? cell_winner_b : cell_winner_a;
-        unsigned int * prev_winner = (t & 1u) ? cell_winner_a : cell_winner_b;
-
-        // ---- Phase 0: clear curr bitsets for my cell range ----
-        const unsigned int my_cell_lo = col_lo * cpc;
-        const unsigned int my_cell_hi = col_hi * cpc;
-        if (cpc == 32u) {
-            // Fast path: one word per column.
-            for (unsigned int c = col_lo + lane; c < col_hi; c += 32u) {
-                curr_active[c] = 0u;
-                curr_winner[c] = 0u;
-            }
-        } else {
-            for (unsigned int cell = my_cell_lo + lane; cell < my_cell_hi; cell += 32u) {
-                unsigned int w = cell >> 5;
-                unsigned int m = 1u << (cell & 31u);
-                atomicAnd(&curr_active[w], ~m);
-                atomicAnd(&curr_winner[w], ~m);
-            }
-        }
-
-        // Block 0, lane 0, warp 0 resets step-scratch counters.
-        if (blockIdx.x == 0u && tid == 0u) {
-            step_scratch[0] = 0u;
-            step_scratch[1] = 0u;
-        }
-
-        // ---- BARRIER 1 ----
-        // Fence: make the above clear-bitsets + scratch writes globally
-        // visible before peer blocks observe "barrier arrived".
-        __threadfence();
-        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
-
-        // =========================================================
-        // T9: TMA MULTICAST INPUT STAGING
-        //
-        // Issue a single cluster-scope async DMA to broadcast this
-        // timestep's input slice into s_input_tile across all 16 SMs
-        // in the cluster simultaneously.  On Hopper sm_90a,
-        // cg::memcpy_async with cluster scope maps to the TMA
-        // hardware unit (cp.async.bulk.tensor multicast), reducing
-        // DRAM input traffic by ~16× vs each block fetching its own
-        // copy from GMEM.
-        //
-        // The staging is gated on cfg.input_bits <= INPUT_BITS_MAX.
-        // If the tile is too small (custom large input_bits), we fall
-        // back to per-thread GMEM reads in Stage A (identical to the
-        // original path; use_input_tile==false).
-        //
-        // Ordering: BARRIER 1 completes before we issue the DMA.
-        // The DMA completes before Stage A reads s_input_tile.
-        // =========================================================
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-        const bool use_input_tile = (cfg.input_bits <= INPUT_BITS_MAX);
-        if (use_input_tile) {
-            auto tb = cg::this_thread_block();
-            cg::memcpy_async(tb, s_input_tile,
-                             inputs + inp_off,
-                             cfg.input_bits);
-            cg::wait(tb);
-            cluster.sync();
-        }
-#else
-        const bool use_input_tile = false;
-#endif
-
-        // =========================================================
-        // STAGE A: Spatial Pooler
-        //
-        // Hot per-column state (boost, inhibition_threshold,
-        // active_duty) is served from cluster DSMEM rather than
-        // GMEM for each of the T timesteps.  GMEM is written on
-        // update so state persists across forward calls.
-        // =========================================================
-        for (unsigned int c = col_lo; c < col_hi; c++) {
-            unsigned int base = c * S;
-            unsigned int local = 0u;
-            for (unsigned int s = lane; s < S; s += 32u) {
-                unsigned int b = syn_bit[base + s];
-                float p = syn_perm[base + s];
-                // T9: read from cluster-broadcast tile when available;
-                // fall back to direct GMEM when input_bits > INPUT_BITS_MAX.
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-                unsigned int inp_byte = use_input_tile
-                    ? (unsigned int)s_input_tile[b]
-                    : (unsigned int)inputs[inp_off + b];
-#else
-                unsigned int inp_byte = (unsigned int)inputs[inp_off + b];
-#endif
-                unsigned int hit = ((inp_byte != 0u) && (p >= cfg.conn_thr)) ? 1u : 0u;
-                local += hit;
-            }
-            unsigned int overlap = warp_sum_u32(local);
-            overlap = __shfl_sync(0xffffffffu, overlap, 0);
-
-            // Read boost + threshold for column c.
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-            // Hopper: read from cluster-distributed shared memory.
-            const unsigned int owner_block  = c / cols_per_block;
-            const unsigned int owner_offset = c - owner_block * cols_per_block;
-            float boost_val = cluster.map_shared_rank(s_boost,      owner_block)[owner_offset];
-            float thr       = cluster.map_shared_rank(s_inhib_thr,  owner_block)[owner_offset];
-#else
-            // Pre-Hopper: read directly from global memory.
-            float boost_val = boost[c];
-            float thr       = inhibition_threshold[c];
-#endif
-
-            float boosted = (float)overlap * boost_val;
-            unsigned int is_active = (boosted > thr) ? 1u : 0u;
-
-            if (lane == 0) {
-                cols_out[col_base_out + c] = (unsigned char)is_active;
-                if (is_active) {
-                    atomicAdd(&step_scratch[0], 1u);
-                }
-            }
-
-            // SP learn (Hebbian) on active columns.
-            // T9: use tile for input reads here too.
-            if (cfg.learn && is_active) {
-                for (unsigned int s = lane; s < S; s += 32u) {
-                    unsigned int b = syn_bit[base + s];
-                    float p = syn_perm[base + s];
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-                    unsigned int inp_byte = use_input_tile
-                        ? (unsigned int)s_input_tile[b]
-                        : (unsigned int)inputs[inp_off + b];
-#else
-                    unsigned int inp_byte = (unsigned int)inputs[inp_off + b];
-#endif
-                    if (inp_byte != 0u) {
-                        p += cfg.sp_inc;
-                        if (p > 1.0f) p = 1.0f;
-                    } else {
-                        p -= cfg.sp_dec;
-                        if (p < 0.0f) p = 0.0f;
-                    }
-                    syn_perm[base + s] = p;
-                }
-            }
-
-            // active_duty EMA + threshold adaptation.
-            // Writes go to both DSMEM (hot path, Hopper only) and GMEM (persistence).
-            if (lane == 0) {
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-                float ad = cluster.map_shared_rank(s_active_duty, owner_block)[owner_offset];
-#else
-                float ad = active_duty[c];
-#endif
-                float sample = is_active ? 1.0f : 0.0f;
-                ad = (1.0f - cfg.duty_alpha) * ad + cfg.duty_alpha * sample;
-
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-                // Writeback: peer smem (for next timestep read) + GMEM (persistence).
-                cluster.map_shared_rank(s_active_duty, owner_block)[owner_offset] = ad;
-#endif
-                active_duty[c] = ad;
-
-                // Threshold steers toward target sparsity.
-                float err = ad - cfg.sparsity_target;
-                float new_thr = thr + cfg.thr_adapt_rate * err * 100.0f;
-                if (new_thr < 0.1f) new_thr = 0.1f;
-                if (new_thr > 1000.0f) new_thr = 1000.0f;
-
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-                // Writeback: peer smem (for next timestep read) + GMEM (persistence).
-                cluster.map_shared_rank(s_inhib_thr, owner_block)[owner_offset] = new_thr;
-#endif
-                inhibition_threshold[c] = new_thr;
-            }
-        }
-
-        // ---- DSMEM WRITEBACK SYNC: peer-smem writes must be visible cluster-wide ----
-        //
-        // On Hopper: cluster.sync() ensures all peer smem writes from this
-        // timestep are visible to all blocks before Stage B / next t.
-        // On pre-Hopper: no smem peer writes occur (all state in GMEM),
-        // so no extra sync needed here — the grid barrier below suffices.
-#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-        cluster.sync();
-#endif
-
-        // ---- BARRIER 2: SP active_mask must be visible before TM reads ----
-        // Fence: flush cols_out + active_duty + inhibition_threshold + step_scratch
-        // writes to global memory before peers advance past this barrier.
-        __threadfence();
-        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
-
-        // =========================================================
-        // STAGE B: Temporal Memory
-        // =========================================================
-        for (unsigned int c = col_lo; c < col_hi; c++) {
-            unsigned int col_active = cols_out[col_base_out + c];
-            if (col_active == 0u) continue;
-
-            unsigned int base_cell = c * cpc;
-            unsigned int any_predicted = 0u;
-            unsigned int best_seg_id_for_grow = 0xFFFFFFFFu;
-            unsigned int best_pot_count = 0u;
-
-            for (unsigned int k = 0u; k < cpc; k++) {
-                unsigned int cell = base_cell + k;
-                unsigned int n_segs_here = cell_seg_count[cell];
-                if (n_segs_here > MSC) n_segs_here = MSC;
-                if (n_segs_here == 0u) continue;
-
-                unsigned int seg_base_id = cell * MSC;
-                unsigned int cell_is_predictive = 0u;
-
-                for (unsigned int ls = 0u; ls < n_segs_here; ls++) {
-                    unsigned int seg = seg_base_id + ls;
-                    unsigned int n_syn = seg_syn_count[seg];
-                    if (n_syn == 0u) continue;
-                    unsigned int syn_base = seg * SPS;
-
-                    unsigned int l_conn = 0u;
-                    unsigned int l_pot  = 0u;
-                    for (unsigned int s = lane; s < n_syn; s += 32u) {
-                        unsigned int presyn = syn_presyn[syn_base + s];
-                        unsigned int w = prev_active[presyn >> 5];
-                        unsigned int bit = (w >> (presyn & 31u)) & 1u;
-                        if (bit) {
-                            l_pot += 1u;
-                            int p = (int)tm_syn_perm[syn_base + s];
-                            if (p >= cfg.conn_thr_i16) l_conn += 1u;
-                        }
-                    }
-                    unsigned int tot_conn = warp_sum_u32(l_conn);
-                    unsigned int tot_pot  = warp_sum_u32(l_pot);
-                    tot_conn = __shfl_sync(0xffffffffu, tot_conn, 0);
-                    tot_pot  = __shfl_sync(0xffffffffu, tot_pot, 0);
-
-                    if (tot_conn >= cfg.activation_threshold) cell_is_predictive = 1u;
-                    if (tot_pot >= cfg.learning_threshold && tot_pot > best_pot_count) {
-                        best_pot_count = tot_pot;
-                        best_seg_id_for_grow = seg;
-                    }
-
-                    // Reinforce predicted-and-correct segment.
-                    if (cfg.learn && tot_conn >= cfg.activation_threshold) {
-                        for (unsigned int s = lane; s < n_syn; s += 32u) {
-                            unsigned int presyn = syn_presyn[syn_base + s];
-                            unsigned int w = prev_active[presyn >> 5];
-                            unsigned int bit = (w >> (presyn & 31u)) & 1u;
-                            int p = (int)tm_syn_perm[syn_base + s];
-                            if (bit) {
-                                int np = p + cfg.perm_inc_i16;
-                                if (np > 32767) np = 32767;
-                                tm_syn_perm[syn_base + s] = (short)np;
-                            } else {
-                                int np = p - cfg.perm_dec_i16;
-                                if (np < 0) np = 0;
-                                tm_syn_perm[syn_base + s] = (short)np;
-                            }
-                        }
-                    }
-                }
-
-                if (cell_is_predictive) {
-                    any_predicted = 1u;
-                    if (lane == 0) {
-                        unsigned int w = cell >> 5;
-                        unsigned int m = 1u << (cell & 31u);
-                        atomicOr(&curr_active[w], m);
-                        atomicOr(&curr_winner[w], m);
-                    }
-                }
-            }
-
-            // BURST if no predicted.
-            if (!any_predicted) {
-                if (lane == 0) {
-                    for (unsigned int k = 0u; k < cpc; k++) {
-                        unsigned int cell = base_cell + k;
-                        unsigned int w = cell >> 5;
-                        unsigned int m = 1u << (cell & 31u);
-                        atomicOr(&curr_active[w], m);
-                    }
-                    unsigned int win = base_cell;
-                    unsigned int ww = win >> 5;
-                    unsigned int wm = 1u << (win & 31u);
-                    atomicOr(&curr_winner[ww], wm);
-                    atomicAdd(&step_scratch[1], 1u);
-                }
-
-                if (cfg.learn) {
-                    unsigned int target_seg;
-                    unsigned int existing_syn;
-                    if (best_seg_id_for_grow != 0xFFFFFFFFu) {
-                        // Reuse best matching segment.
-                        target_seg = best_seg_id_for_grow;
-                        existing_syn = seg_syn_count[target_seg];
-                        target_seg = __shfl_sync(0xffffffffu, target_seg, 0);
-                        existing_syn = __shfl_sync(0xffffffffu, existing_syn, 0);
-
-                        // Reinforce its existing synapses.
-                        unsigned int syn_base = target_seg * SPS;
-                        for (unsigned int s = lane; s < existing_syn; s += 32u) {
-                            unsigned int presyn = syn_presyn[syn_base + s];
-                            unsigned int w = prev_active[presyn >> 5];
-                            unsigned int bit = (w >> (presyn & 31u)) & 1u;
-                            int p = (int)tm_syn_perm[syn_base + s];
-                            if (bit) {
-                                int np = p + cfg.perm_inc_i16;
-                                if (np > 32767) np = 32767;
-                                tm_syn_perm[syn_base + s] = (short)np;
-                            } else {
-                                int np = p - cfg.perm_dec_i16;
-                                if (np < 0) np = 0;
-                                tm_syn_perm[syn_base + s] = (short)np;
-                            }
-                        }
-                    } else {
-                        // Allocate new segment on winner cell (cell 0 of col).
-                        unsigned int new_seg = 0u;
-                        if (lane == 0) {
-                            unsigned int winner_cell = base_cell;
-                            unsigned int slot = atomicAdd(&cell_seg_count[winner_cell], 1u);
-                            if (slot >= MSC) slot = slot % MSC;
-                            new_seg = winner_cell * MSC + slot;
-                            seg_cell_id[new_seg] = winner_cell;
-                            seg_syn_count[new_seg] = 0u;
-                        }
-                        target_seg = __shfl_sync(0xffffffffu, new_seg, 0);
-                        existing_syn = 0u;
-                    }
-
-                    // Grow synapses to prev_winner cells — lane 0 serialized.
-                    unsigned int room = (SPS > existing_syn) ? (SPS - existing_syn) : 0u;
-                    unsigned int max_grow = (cfg.max_new_synapses < room) ? cfg.max_new_synapses : room;
-                    if (lane == 0 && max_grow > 0u) {
-                        unsigned int syn_base = target_seg * SPS;
-                        unsigned int grown = 0u;
-                        unsigned int start_off = (c * 2654435761u + cfg.iter_seed + t) % cfg.bits_words;
-                        for (unsigned int w_off = 0u;
-                             w_off < cfg.bits_words && grown < max_grow;
-                             w_off++) {
-                            unsigned int widx = (start_off + w_off) % cfg.bits_words;
-                            unsigned int word = prev_winner[widx];
-                            while (word != 0u && grown < max_grow) {
-                                unsigned int bit_pos = __ffs(word) - 1u;
-                                word &= ~(1u << bit_pos);
-                                unsigned int cell_id = widx * 32u + bit_pos;
-                                if (cell_id >= cfg.n_cells) continue;
-                                bool exists = false;
-                                for (unsigned int es = 0u; es < existing_syn + grown; es++) {
-                                    if (syn_presyn[syn_base + es] == cell_id) { exists = true; break; }
-                                }
-                                if (exists) continue;
-                                unsigned int write_idx = existing_syn + grown;
-                                if (write_idx >= SPS) break;
-                                syn_presyn[syn_base + write_idx] = cell_id;
-                                tm_syn_perm[syn_base + write_idx] = (short)cfg.initial_perm_i16;
-                                grown++;
-                            }
-                        }
-                        if (grown > 0u) {
-                            seg_syn_count[target_seg] = existing_syn + grown;
-                        }
-                    }
-                }
-            }
-        }
-
-        // ---- BARRIER 3: TM writes complete before anomaly + next-step read ----
-        // Fence: flush curr_active/curr_winner bitsets + tm_syn_perm +
-        // seg_syn_count + syn_presyn before peers advance and consume them as
-        // prev_active/prev_winner at t+1.
-        __threadfence();
-        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
-
-        // Write anomaly for step t.
-        if (blockIdx.x == 0u && tid == 0u) {
-            unsigned int total = step_scratch[0];
-            unsigned int bad   = step_scratch[1];
-            float anom = (total > 0u) ? ((float)bad / (float)total) : 0.0f;
-            anom_out[t] = anom;
-        }
-    }
-}
-
-// Single-region kernel (legacy call site).
-__global__ __launch_bounds__(256, 2)
-void htm_fused_step(FusedPtrs P, FusedConfig cfg) {
-    htm_fused_step_body(P, cfg);
-}
-
-// Batched kernel: one cooperative launch for B regions. grid.y = B,
-// grid.x = per-region block count. Each block reads its region's
-// FusedPtrs from the device array via blockIdx.y.
-__global__ __launch_bounds__(256, 2)
-void htm_fused_step_batched(const FusedPtrs* __restrict__ P_arr, FusedConfig cfg) {
-    const FusedPtrs P = P_arr[blockIdx.y];
-    htm_fused_step_body(P, cfg);
-}
-
-} // extern "C"
+// Fused HTM megakernel — SP + TM, all T timesteps in a single launch.
+//
+// Design rationale:
+//   - Global top-K column selection requires cross-block synchronization at
+//     every timestep (grid.sync is unreliable on WSL2/sm_86 without rdc=true).
+//   - Replace with per-column threshold activation using local lateral
+//     inhibition: column c activates if overlap[c]*boost[c] > threshold[c].
+//     Threshold is a per-column running-EMA learned scalar that steers the
+//     column's long-run activation rate toward the global sparsity target.
+//   - This is biologically grounded (GABAergic local inhibition) and supported
+//     by HTM theory (duty-cycle boost already drives this loop; we just
+//     change which lever the EMA pulls).
+//
+// Launch shape:
+//   grid  = min(device SM count, 16)  // hard cap — see below
+//   block = 1024 threads = 32 warps
+//   Each warp of 32 owns a contiguous column slice (n_columns / total_warps).
+//
+// Cross-block coherence:
+//   - Ping-pong buffers for cell_active/cell_winner: write _a at even t,
+//     read _b; reversed at odd t.
+//   - Preferred path: cooperative launch + hardware whole-grid sync.
+//   - Fallback path: software 3-slot rotating grid barrier for devices/drivers
+//     that cannot do cooperative launch.
+//
+// 2026-04-16: grid_dim reduced from 28 to 16 after deadlock RCA. The previous
+// cap of 28 relied on all blocks being concurrently resident on a 30-SM RTX
+// 3060 Laptop. Under thermal throttling effective residency dropped to ~20-24,
+// leaving scheduled blocks spinning on the software grid barrier waiting for
+// peer blocks that would never run. 16 blocks is below any realistic residency
+// floor and preserves enough warp parallelism (16*32 = 512 warps) to saturate
+// memory bandwidth on the spatial-pooler stage.
+//
+// Kernel signature uses struct-by-value for pointers and config to stay
+// inside cudarc's launch-arg count limit.
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+
+namespace cg = cooperative_groups;
+
+// Maximum columns owned per cluster-block in DSMEM.
+// Supports n_columns up to COLS_PER_CLUSTER_BLOCK_MAX * cluster_size.
+// At cluster_size=16: supports up to 256*16=4096 columns.
+// Each array costs 256*4 = 1024 bytes; three arrays = 3072 bytes per SM —
+// well under the 228 KB H200 shared-memory cap.
+#define COLS_PER_CLUSTER_BLOCK_MAX 256u
+
+// Maximum input_bits supported by the TMA-multicast staging tile.
+// At 32 KB this covers the production SDR width (16384 bits) with 2× headroom.
+// Total shared per SM: 32768 (tile) + 3072 (DSMEM float arrays) = ~35 KB —
+// well under the 228 KB H200 limit.
+//
+// Expected speedup from TMA multicast input staging (T9/T11):
+//   - Without staging: 16 SMs × T × (input_bits GMEM reads per timestep)
+//   - With staging:    1 TMA DMA per timestep, shared reads from L1 thereafter
+//   - Theoretical DRAM bandwidth reduction: ~16× on input reads
+//   - Wall-clock reduction estimate: -20 to -40 ms from reduced input fetch latency
+#define INPUT_BITS_MAX 32768u
+
+extern "C" {
+
+struct FusedPtrs {
+    unsigned long long syn_bit;
+    unsigned long long syn_perm;
+    unsigned long long boost;
+    unsigned long long active_duty;
+    unsigned long long inhibition_threshold;
+    unsigned long long seg_cell_id;
+    unsigned long long seg_syn_count;
+    unsigned long long syn_presyn;
+    unsigned long long tm_syn_perm;
+    unsigned long long cell_seg_count;
+    unsigned long long cell_active_a;
+    unsigned long long cell_active_b;
+    unsigned long long cell_winner_a;
+    unsigned long long cell_winner_b;
+    unsigned long long inputs;
+    unsigned long long cols_out;
+    unsigned long long anom_out;
+    unsigned long long barrier_counters;
+    unsigned long long step_scratch;
+};
+
+struct FusedConfig {
+    // SP constants
+    unsigned int input_bits;
+    unsigned int n_columns;
+    unsigned int synapses_per_col;
+    float        conn_thr;
+    float        sp_inc;
+    float        sp_dec;
+    float        sparsity_target;
+    float        duty_alpha;
+    float        thr_adapt_rate;
+    // TM constants
+    unsigned int cells_per_column;
+    unsigned int n_cells;
+    unsigned int bits_words;
+    unsigned int max_segments_per_cell;
+    unsigned int synapses_per_segment;
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int max_new_synapses;
+    int          conn_thr_i16;
+    int          perm_inc_i16;
+    int          perm_dec_i16;
+    int          predicted_seg_dec_i16;
+    int          initial_perm_i16;
+    // Loop constants
+    unsigned int T;
+    unsigned int learn;
+    unsigned int iter_seed;
+    unsigned int cooperative_grid_sync;
+};
+
+// Hardware cluster barrier using Hopper sm_90a cooperative_groups::this_cluster().sync().
+// Replaces the former software Decoupled Look-Back (DLB) atomic-spin barrier.
+//
+// cluster::sync() is a single PTX instruction (barrier.cluster) that resolves
+// in ~10-40 ns inside the cluster, with no device-level serialization.
+// Multiple clusters (one per HTM region) run fully concurrently — bounded
+// only by SM count (8 clusters × 16 SMs = 128 ≤ 132 on H200).
+//
+// The flags / expected / phase / cooperative_grid_sync parameters are kept
+// in the signature for call-site compatibility but are unused.
+__device__ static inline void fused_grid_barrier(cg::grid_group grid,
+                                                 unsigned int * /* flags — unused */,
+                                                 unsigned int /* expected — unused */,
+                                                 unsigned int /* phase — unused */,
+                                                 unsigned int /* cooperative_grid_sync — unused */) {
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    // Hopper+ : hardware cluster barrier (~10-40 ns)
+    auto cluster = cg::this_cluster();
+    cluster.sync();
+#else
+    // Pre-Hopper (sm_80, sm_86, sm_89): grid-level cooperative sync.
+    // Requires cooperative kernel launch. ~us-ms range, adequate for HTM
+    // workload (kernel launch frequency is low).
+    grid.sync();
+#endif
+}
+
+__device__ static inline unsigned int warp_sum_u32(unsigned int v) {
+    for (int off = 16; off > 0; off >>= 1) {
+        v += __shfl_down_sync(0xffffffffu, v, off);
+    }
+    return v;
+}
+
+// Core kernel body — works for both single-region and batched launches.
+// Single-region: caller passes the one FusedPtrs struct.
+// Batched: each block reads its region's FusedPtrs via blockIdx.y before
+// calling this. State is independent per region (each region owns its own
+// GPU buffers); grid.sync() is the only cross-block primitive and it
+// spans ALL blocks in the grid (harmless over-sync across regions).
+__device__ static inline
+void htm_fused_step_body(const FusedPtrs& P, const FusedConfig& cfg) {
+    cg::grid_group grid = cg::this_grid();
+    // Cast pointers.
+    const unsigned int  * __restrict__ syn_bit               = (const unsigned int*)P.syn_bit;
+    float               * __restrict__ syn_perm              = (float*)P.syn_perm;
+    float               * __restrict__ boost                 = (float*)P.boost;
+    float               * __restrict__ active_duty           = (float*)P.active_duty;
+    float               * __restrict__ inhibition_threshold  = (float*)P.inhibition_threshold;
+    unsigned int        * __restrict__ seg_cell_id           = (unsigned int*)P.seg_cell_id;
+    unsigned int        * __restrict__ seg_syn_count         = (unsigned int*)P.seg_syn_count;
+    unsigned int        * __restrict__ syn_presyn            = (unsigned int*)P.syn_presyn;
+    short               * __restrict__ tm_syn_perm           = (short*)P.tm_syn_perm;
+    unsigned int        * __restrict__ cell_seg_count        = (unsigned int*)P.cell_seg_count;
+    unsigned int        * __restrict__ cell_active_a         = (unsigned int*)P.cell_active_a;
+    unsigned int        * __restrict__ cell_active_b         = (unsigned int*)P.cell_active_b;
+    unsigned int        * __restrict__ cell_winner_a         = (unsigned int*)P.cell_winner_a;
+    unsigned int        * __restrict__ cell_winner_b         = (unsigned int*)P.cell_winner_b;
+    const unsigned char * __restrict__ inputs                = (const unsigned char*)P.inputs;
+    unsigned char       * __restrict__ cols_out              = (unsigned char*)P.cols_out;
+    float               * __restrict__ anom_out              = (float*)P.anom_out;
+    unsigned int        * __restrict__ barrier_counters      = (unsigned int*)P.barrier_counters;
+    unsigned int        * __restrict__ step_scratch          = (unsigned int*)P.step_scratch;
+
+    const unsigned int tid     = threadIdx.x;
+    const unsigned int lane    = tid & 31u;
+    const unsigned int warp    = tid >> 5;
+    const unsigned int warps_per_block = blockDim.x >> 5;
+    const unsigned int gwarp   = blockIdx.x * warps_per_block + warp;
+    const unsigned int n_warps = gridDim.x * warps_per_block;
+
+    const unsigned int n_cols  = cfg.n_columns;
+    const unsigned int col_lo  = (gwarp * n_cols) / n_warps;
+    const unsigned int col_hi  = ((gwarp + 1) * n_cols) / n_warps;
+
+    unsigned int phase = 0u;
+
+    // =========================================================
+    // DSMEM: Cluster-distributed shared memory for hot per-column
+    // state (inhibition_threshold, boost, active_duty).
+    //
+    // On Hopper (sm_90+): Each block in the cluster owns a contiguous
+    // slice of columns in its own __shared__ arrays. Any block can
+    // peer-read another block's slice via cluster.map_shared_rank().
+    //
+    // On Ampere (sm_86) and other pre-Hopper: No cluster support.
+    // Read/write directly from/to global memory (inhibition_threshold,
+    // boost, active_duty device pointers). Slightly higher latency but
+    // functionally correct.
+    // =========================================================
+
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    // Hopper+ cluster path
+    auto cluster = cg::this_cluster();
+    const unsigned int cluster_block_rank = cluster.block_rank();  // 0..cluster_size-1
+    const unsigned int cluster_sz         = cluster.num_blocks();  // == gridDim.x (≤16)
+#else
+    // Pre-Hopper: no cluster, each block is independent.
+    const unsigned int cluster_block_rank = blockIdx.x;
+    const unsigned int cluster_sz         = gridDim.x;
+#endif
+
+    // Partition n_cols evenly across cluster blocks.
+    // Each block owns cols_per_block columns starting at my_col_start.
+    const unsigned int cols_per_block =
+        (n_cols + cluster_sz - 1u) / cluster_sz;               // ceil div
+    const unsigned int my_col_start =
+        cluster_block_rank * cols_per_block;
+    const unsigned int my_col_end =
+        (my_col_start + cols_per_block < n_cols)
+            ? (my_col_start + cols_per_block) : n_cols;        // clamp
+
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    // Cluster-distributed shared memory arrays.
+    // Each block holds at most COLS_PER_CLUSTER_BLOCK_MAX floats per array.
+    // Peer blocks address into each other's smem via map_shared_rank.
+    __shared__ float s_inhib_thr [COLS_PER_CLUSTER_BLOCK_MAX];
+    __shared__ float s_boost     [COLS_PER_CLUSTER_BLOCK_MAX];
+    __shared__ float s_active_duty[COLS_PER_CLUSTER_BLOCK_MAX];
+#endif
+
+    // TMA multicast input staging tile (T9) — HOPPER ONLY.
+    //
+    // On Hopper: cg::memcpy_async with cluster scope multicasts input to all
+    // 16 SMs, reducing DRAM traffic by ~16×.
+    // On Ampere: 32 KB smem allocation exceeds per-block budget when
+    // cooperatively launched (48 KB total, registers eat the rest). Skip the
+    // tile entirely — Stage A reads from GMEM directly (original path).
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    __shared__ __align__(16) unsigned char s_input_tile[INPUT_BITS_MAX];
+#endif
+
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    // Initial GMEM → smem load (reads state from previous forward call).
+    // Each block loads only its own slice; tid strides across the slice.
+    for (unsigned int c = my_col_start + tid; c < my_col_end; c += blockDim.x) {
+        const unsigned int off = c - my_col_start;
+        s_inhib_thr [off] = inhibition_threshold[c];
+        s_boost     [off] = boost[c];
+        s_active_duty[off] = active_duty[c];
+    }
+
+    // All blocks in the cluster must finish loading before any block
+    // starts reading peer smem inside the T-loop.
+    cluster.sync();
+#else
+    // Pre-Hopper: no smem caching needed — reads go directly to GMEM.
+    // Grid sync ensures all blocks have completed Phase 0 init before T-loop.
+    grid.sync();
+#endif
+
+    const unsigned int S   = cfg.synapses_per_col;
+    const unsigned int cpc = cfg.cells_per_column;
+    const unsigned int SPS = cfg.synapses_per_segment;
+    const unsigned int MSC = cfg.max_segments_per_cell;
+
+    // Main timestep loop.
+    for (unsigned int t = 0u; t < cfg.T; t++) {
+        const unsigned int inp_off      = t * cfg.input_bits;
+        const unsigned int col_base_out = t * n_cols;
+
+        unsigned int * curr_active = (t & 1u) ? cell_active_b : cell_active_a;
+        unsigned int * prev_active = (t & 1u) ? cell_active_a : cell_active_b;
+        unsigned int * curr_winner = (t & 1u) ? cell_winner_b : cell_winner_a;
+        unsigned int * prev_winner = (t & 1u) ? cell_winner_a : cell_winner_b;
+
+        // ---- Phase 0: clear curr bitsets for my cell range ----
+        const unsigned int my_cell_lo = col_lo * cpc;
+        const unsigned int my_cell_hi = col_hi * cpc;
+        if (cpc == 32u) {
+            // Fast path: one word per column.
+            for (unsigned int c = col_lo + lane; c < col_hi; c += 32u) {
+                curr_active[c] = 0u;
+                curr_winner[c] = 0u;
+            }
+        } else {
+            for (unsigned int cell = my_cell_lo + lane; cell < my_cell_hi; cell += 32u) {
+                unsigned int w = cell >> 5;
+                unsigned int m = 1u << (cell & 31u);
+                atomicAnd(&curr_active[w], ~m);
+                atomicAnd(&curr_winner[w], ~m);
+            }
+        }
+
+        // Block 0, lane 0, warp 0 resets step-scratch counters.
+        if (blockIdx.x == 0u && tid == 0u) {
+            step_scratch[0] = 0u;
+            step_scratch[1] = 0u;
+        }
+
+        // ---- BARRIER 1 ----
+        // Fence: make the above clear-bitsets + scratch writes globally
+        // visible before peer blocks observe "barrier arrived".
+        __threadfence();
+        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
+
+        // =========================================================
+        // T9: TMA MULTICAST INPUT STAGING
+        //
+        // Issue a single cluster-scope async DMA to broadcast this
+        // timestep's input slice into s_input_tile across all 16 SMs
+        // in the cluster simultaneously.  On Hopper sm_90a,
+        // cg::memcpy_async with cluster scope maps to the TMA
+        // hardware unit (cp.async.bulk.tensor multicast), reducing
+        // DRAM input traffic by ~16× vs each block fetching its own
+        // copy from GMEM.
+        //
+        // The staging is gated on cfg.input_bits <= INPUT_BITS_MAX.
+        // If the tile is too small (custom large input_bits), we fall
+        // back to per-thread GMEM reads in Stage A (identical to the
+        // original path; use_input_tile==false).
+        //
+        // Ordering: BARRIER 1 completes before we issue the DMA.
+        // The DMA completes before Stage A reads s_input_tile.
+        // =========================================================
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        const bool use_input_tile = (cfg.input_bits <= INPUT_BITS_MAX);
+        if (use_input_tile) {
+            auto tb = cg::this_thread_block();
+            cg::memcpy_async(tb, s_input_tile,
+                             inputs + inp_off,
+                             cfg.input_bits);
+            cg::wait(tb);
+            cluster.sync();
+        }
+#else
+        const bool use_input_tile = false;
+#endif
+
+        // =========================================================
+        // STAGE A: Spatial Pooler
+        //
+        // Hot per-column state (boost, inhibition_threshold,
+        // active_duty) is served from cluster DSMEM rather than
+        // GMEM for each of the T timesteps.  GMEM is written on
+        // update so state persists across forward calls.
+        // =========================================================
+        for (unsigned int c = col_lo; c < col_hi; c++) {
+            unsigned int base = c * S;
+            unsigned int local = 0u;
+            for (unsigned int s = lane; s < S; s += 32u) {
+                unsigned int b = syn_bit[base + s];
+                float p = syn_perm[base + s];
+                // T9: read from cluster-broadcast tile when available;
+                // fall back to direct GMEM when input_bits > INPUT_BITS_MAX.
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                unsigned int inp_byte = use_input_tile
+                    ? (unsigned int)s_input_tile[b]
+                    : (unsigned int)inputs[inp_off + b];
+#else
+                unsigned int inp_byte = (unsigned int)inputs[inp_off + b];
+#endif
+                unsigned int hit = ((inp_byte != 0u) && (p >= cfg.conn_thr)) ? 1u : 0u;
+                local += hit;
+            }
+            unsigned int overlap = warp_sum_u32(local);
+            overlap = __shfl_sync(0xffffffffu, overlap, 0);
+
+            // Read boost + threshold for column c.
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+            // Hopper: read from cluster-distributed shared memory.
+            const unsigned int owner_block  = c / cols_per_block;
+            const unsigned int owner_offset = c - owner_block * cols_per_block;
+            float boost_val = cluster.map_shared_rank(s_boost,      owner_block)[owner_offset];
+            float thr       = cluster.map_shared_rank(s_inhib_thr,  owner_block)[owner_offset];
+#else
+            // Pre-Hopper: read directly from global memory.
+            float boost_val = boost[c];
+            float thr       = inhibition_threshold[c];
+#endif
+
+            float boosted = (float)overlap * boost_val;
+            unsigned int is_active = (boosted > thr) ? 1u : 0u;
+
+            if (lane == 0) {
+                cols_out[col_base_out + c] = (unsigned char)is_active;
+                if (is_active) {
+                    atomicAdd(&step_scratch[0], 1u);
+                }
+            }
+
+            // SP learn (Hebbian) on active columns.
+            // T9: use tile for input reads here too.
+            if (cfg.learn && is_active) {
+                for (unsigned int s = lane; s < S; s += 32u) {
+                    unsigned int b = syn_bit[base + s];
+                    float p = syn_perm[base + s];
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                    unsigned int inp_byte = use_input_tile
+                        ? (unsigned int)s_input_tile[b]
+                        : (unsigned int)inputs[inp_off + b];
+#else
+                    unsigned int inp_byte = (unsigned int)inputs[inp_off + b];
+#endif
+                    if (inp_byte != 0u) {
+                        p += cfg.sp_inc;
+                        if (p > 1.0f) p = 1.0f;
+                    } else {
+                        p -= cfg.sp_dec;
+                        if (p < 0.0f) p = 0.0f;
+                    }
+                    syn_perm[base + s] = p;
+                }
+            }
+
+            // active_duty EMA + threshold adaptation.
+            // Writes go to both DSMEM (hot path, Hopper only) and GMEM (persistence).
+            if (lane == 0) {
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                float ad = cluster.map_shared_rank(s_active_duty, owner_block)[owner_offset];
+#else
+                float ad = active_duty[c];
+#endif
+                float sample = is_active ? 1.0f : 0.0f;
+                ad = (1.0f - cfg.duty_alpha) * ad + cfg.duty_alpha * sample;
+
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                // Writeback: peer smem (for next timestep read) + GMEM (persistence).
+                cluster.map_shared_rank(s_active_duty, owner_block)[owner_offset] = ad;
+#endif
+                active_duty[c] = ad;
+
+                // Threshold steers toward target sparsity.
+                float err = ad - cfg.sparsity_target;
+                float new_thr = thr + cfg.thr_adapt_rate * err * 100.0f;
+                if (new_thr < 0.1f) new_thr = 0.1f;
+                if (new_thr > 1000.0f) new_thr = 1000.0f;
+
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+                // Writeback: peer smem (for next timestep read) + GMEM (persistence).
+                cluster.map_shared_rank(s_inhib_thr, owner_block)[owner_offset] = new_thr;
+#endif
+                inhibition_threshold[c] = new_thr;
+            }
+        }
+
+        // ---- DSMEM WRITEBACK SYNC: peer-smem writes must be visible cluster-wide ----
+        //
+        // On Hopper: cluster.sync() ensures all peer smem writes from this
+        // timestep are visible to all blocks before Stage B / next t.
+        // On pre-Hopper: no smem peer writes occur (all state in GMEM),
+        // so no extra sync needed here — the grid barrier below suffices.
+#if !defined(HTM_DISABLE_CLUSTER) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+        cluster.sync();
+#endif
+
+        // ---- BARRIER 2: SP active_mask must be visible before TM reads ----
+        // Fence: flush cols_out + active_duty + inhibition_threshold + step_scratch
+        // writes to global memory before peers advance past this barrier.
+        __threadfence();
+        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
+
+        // =========================================================
+        // STAGE B: Temporal Memory
+        // =========================================================
+        for (unsigned int c = col_lo; c < col_hi; c++) {
+            unsigned int col_active = cols_out[col_base_out + c];
+            if (col_active == 0u) continue;
+
+            unsigned int base_cell = c * cpc;
+            unsigned int any_predicted = 0u;
+            unsigned int best_seg_id_for_grow = 0xFFFFFFFFu;
+            unsigned int best_pot_count = 0u;
+
+            for (unsigned int k = 0u; k < cpc; k++) {
+                unsigned int cell = base_cell + k;
+                unsigned int n_segs_here = cell_seg_count[cell];
+                if (n_segs_here > MSC) n_segs_here = MSC;
+                if (n_segs_here == 0u) continue;
+
+                unsigned int seg_base_id = cell * MSC;
+                unsigned int cell_is_predictive = 0u;
+
+                for (unsigned int ls = 0u; ls < n_segs_here; ls++) {
+                    unsigned int seg = seg_base_id + ls;
+                    unsigned int n_syn = seg_syn_count[seg];
+                    if (n_syn == 0u) continue;
+                    unsigned int syn_base = seg * SPS;
+
+                    unsigned int l_conn = 0u;
+                    unsigned int l_pot  = 0u;
+                    for (unsigned int s = lane; s < n_syn; s += 32u) {
+                        unsigned int presyn = syn_presyn[syn_base + s];
+                        unsigned int w = prev_active[presyn >> 5];
+                        unsigned int bit = (w >> (presyn & 31u)) & 1u;
+                        if (bit) {
+                            l_pot += 1u;
+                            int p = (int)tm_syn_perm[syn_base + s];
+                            if (p >= cfg.conn_thr_i16) l_conn += 1u;
+                        }
+                    }
+                    unsigned int tot_conn = warp_sum_u32(l_conn);
+                    unsigned int tot_pot  = warp_sum_u32(l_pot);
+                    tot_conn = __shfl_sync(0xffffffffu, tot_conn, 0);
+                    tot_pot  = __shfl_sync(0xffffffffu, tot_pot, 0);
+
+                    if (tot_conn >= cfg.activation_threshold) cell_is_predictive = 1u;
+                    if (tot_pot >= cfg.learning_threshold && tot_pot > best_pot_count) {
+                        best_pot_count = tot_pot;
+                        best_seg_id_for_grow = seg;
+                    }
+
+                    // Reinforce predicted-and-correct segment.
+                    if (cfg.learn && tot_conn >= cfg.activation_threshold) {
+                        for (unsigned int s = lane; s < n_syn; s += 32u) {
+                            unsigned int presyn = syn_presyn[syn_base + s];
+                            unsigned int w = prev_active[presyn >> 5];
+                            unsigned int bit = (w >> (presyn & 31u)) & 1u;
+                            int p = (int)tm_syn_perm[syn_base + s];
+                            if (bit) {
+                                int np = p + cfg.perm_inc_i16;
+                                if (np > 32767) np = 32767;
+                                tm_syn_perm[syn_base + s] = (short)np;
+                            } else {
+                                int np = p - cfg.perm_dec_i16;
+                                if (np < 0) np = 0;
+                                tm_syn_perm[syn_base + s] = (short)np;
+                            }
+                        }
+                    }
+                }
+
+                if (cell_is_predictive) {
+                    any_predicted = 1u;
+                    if (lane == 0) {
+                        unsigned int w = cell >> 5;
+                        unsigned int m = 1u << (cell & 31u);
+                        atomicOr(&curr_active[w], m);
+                        atomicOr(&curr_winner[w], m);
+                    }
+                }
+            }
+
+            // BURST if no predicted.
+            if (!any_predicted) {
+                if (lane == 0) {
+                    for (unsigned int k = 0u; k < cpc; k++) {
+                        unsigned int cell = base_cell + k;
+                        unsigned int w = cell >> 5;
+                        unsigned int m = 1u << (cell & 31u);
+                        atomicOr(&curr_active[w], m);
+                    }
+                    unsigned int win = base_cell;
+                    unsigned int ww = win >> 5;
+                    unsigned int wm = 1u << (win & 31u);
+                    atomicOr(&curr_winner[ww], wm);
+                    atomicAdd(&step_scratch[1], 1u);
+                }
+
+                if (cfg.learn) {
+                    unsigned int target_seg;
+                    unsigned int existing_syn;
+                    if (best_seg_id_for_grow != 0xFFFFFFFFu) {
+                        // Reuse best matching segment.
+                        target_seg = best_seg_id_for_grow;
+                        existing_syn = seg_syn_count[target_seg];
+                        target_seg = __shfl_sync(0xffffffffu, target_seg, 0);
+                        existing_syn = __shfl_sync(0xffffffffu, existing_syn, 0);
+
+                        // Reinforce its existing synapses.
+                        unsigned int syn_base = target_seg * SPS;
+                        for (unsigned int s = lane; s < existing_syn; s += 32u) {
+                            unsigned int presyn = syn_presyn[syn_base + s];
+                            unsigned int w = prev_active[presyn >> 5];
+                            unsigned int bit = (w >> (presyn & 31u)) & 1u;
+                            int p = (int)tm_syn_perm[syn_base + s];
+                            if (bit) {
+                                int np = p + cfg.perm_inc_i16;
+                                if (np > 32767) np = 32767;
+                                tm_syn_perm[syn_base + s] = (short)np;
+                            } else {
+                                int np = p - cfg.perm_dec_i16;
+                                if (np < 0) np = 0;
+                                tm_syn_perm[syn_base + s] = (short)np;
+                            }
+                        }
+                    } else {
+                        // Allocate new segment on winner cell (cell 0 of col).
+                        unsigned int new_seg = 0u;
+                        if (lane == 0) {
+                            unsigned int winner_cell = base_cell;
+                            unsigned int slot = atomicAdd(&cell_seg_count[winner_cell], 1u);
+                            if (slot >= MSC) slot = slot % MSC;
+                            new_seg = winner_cell * MSC + slot;
+                            seg_cell_id[new_seg] = winner_cell;
+                            seg_syn_count[new_seg] = 0u;
+                        }
+                        target_seg = __shfl_sync(0xffffffffu, new_seg, 0);
+                        existing_syn = 0u;
+                    }
+
+                    // Grow synapses to prev_winner cells — lane 0 serialized.
+                    unsigned int room = (SPS > existing_syn) ? (SPS - existing_syn) : 0u;
+                    unsigned int max_grow = (cfg.max_new_synapses < room) ? cfg.max_new_synapses : room;
+                    if (lane == 0 && max_grow > 0u) {
+                        unsigned int syn_base = target_seg * SPS;
+                        unsigned int grown = 0u;
+                        unsigned int start_off = (c * 2654435761u + cfg.iter_seed + t) % cfg.bits_words;
+                        for (unsigned int w_off = 0u;
+                             w_off < cfg.bits_words && grown < max_grow;
+                             w_off++) {
+                            unsigned int widx = (start_off + w_off) % cfg.bits_words;
+                            unsigned int word = prev_winner[widx];
+                            while (word != 0u && grown < max_grow) {
+                                unsigned int bit_pos = __ffs(word) - 1u;
+                                word &= ~(1u << bit_pos);
+                                unsigned int cell_id = widx * 32u + bit_pos;
+                                if (cell_id >= cfg.n_cells) continue;
+                                bool exists = false;
+                                for (unsigned int es = 0u; es < existing_syn + grown; es++) {
+                                    if (syn_presyn[syn_base + es] == cell_id) { exists = true; break; }
+                                }
+                                if (exists) continue;
+                                unsigned int write_idx = existing_syn + grown;
+                                if (write_idx >= SPS) break;
+                                syn_presyn[syn_base + write_idx] = cell_id;
+                                tm_syn_perm[syn_base + write_idx] = (short)cfg.initial_perm_i16;
+                                grown++;
+                            }
+                        }
+                        if (grown > 0u) {
+                            seg_syn_count[target_seg] = existing_syn + grown;
+                        }
+                    }
+                }
+            }
+        }
+
+        // ---- BARRIER 3: TM writes complete before anomaly + next-step read ----
+        // Fence: flush curr_active/curr_winner bitsets + tm_syn_perm +
+        // seg_syn_count + syn_presyn before peers advance and consume them as
+        // prev_active/prev_winner at t+1.
+        __threadfence();
+        fused_grid_barrier(grid, barrier_counters, 0u, phase++, cfg.cooperative_grid_sync);
+
+        // Write anomaly for step t.
+        if (blockIdx.x == 0u && tid == 0u) {
+            unsigned int total = step_scratch[0];
+            unsigned int bad   = step_scratch[1];
+            float anom = (total > 0u) ? ((float)bad / (float)total) : 0.0f;
+            anom_out[t] = anom;
+        }
+    }
+}
+
+// Single-region kernel (legacy call site).
+__global__ __launch_bounds__(256, 2)
+void htm_fused_step(FusedPtrs P, FusedConfig cfg) {
+    htm_fused_step_body(P, cfg);
+}
+
+// Batched kernel: one cooperative launch for B regions. grid.y = B,
+// grid.x = per-region block count. Each block reads its region's
+// FusedPtrs from the device array via blockIdx.y.
+__global__ __launch_bounds__(256, 2)
+void htm_fused_step_batched(const FusedPtrs* __restrict__ P_arr, FusedConfig cfg) {
+    const FusedPtrs P = P_arr[blockIdx.y];
+    htm_fused_step_body(P, cfg);
+}
+
+} // extern "C"

overlay/htm_rust/src/gpu/kernels/sp_boost_fused.cu

@@ -1,59 +1,59 @@
-// Fused mean-reduction + boost-update kernel.
-//
-// Inputs:
-//   active_duty[n] (f32)
-//   boost_strength (f32)
-//
-// Output:
-//   boost[n] (f32) = expf(-boost_strength * (active_duty[c] - mean))
-//
-// Launch: single block (1024 threads), shared mem for reduction. At n=2048
-// each thread handles 2 elements.
-
-extern "C" __global__
-void sp_boost_from_duty(
-    const float * __restrict__ active_duty,  // (n,)
-    float       * __restrict__ boost,        // (n,) in-place out
-    float         boost_strength,
-    unsigned int  n
-) {
-    extern __shared__ float smem_raw[];
-    float * smem = smem_raw;
-    const unsigned int tid = threadIdx.x;
-    const unsigned int bsz = blockDim.x;
-
-    // Phase 1: parallel sum of active_duty into smem[0..32] (warp-level).
-    float local_sum = 0.0f;
-    for (unsigned int i = tid; i < n; i += bsz) {
-        local_sum += active_duty[i];
-    }
-    // Warp reduction.
-    for (int off = 16; off > 0; off >>= 1) {
-        local_sum += __shfl_down_sync(0xffffffff, local_sum, off);
-    }
-    unsigned int lane = tid & 31;
-    unsigned int warp = tid >> 5;
-    if (lane == 0) smem[warp] = local_sum;
-    __syncthreads();
-
-    // Warp 0 reduces warp-sums.
-    __shared__ float mean_s;
-    if (warp == 0) {
-        unsigned int nwarps = (bsz + 31) / 32;
-        float v = (lane < nwarps) ? smem[lane] : 0.0f;
-        for (int off = 16; off > 0; off >>= 1) {
-            v += __shfl_down_sync(0xffffffff, v, off);
-        }
-        if (tid == 0) {
-            mean_s = v / (float)n;
-        }
-    }
-    __syncthreads();
-
-    // Phase 2: boost[c] = expf(-strength * (active_duty[c] - mean)).
-    float mean = mean_s;
-    for (unsigned int i = tid; i < n; i += bsz) {
-        float d = active_duty[i] - mean;
-        boost[i] = expf(-boost_strength * d);
-    }
-}
+// Fused mean-reduction + boost-update kernel.
+//
+// Inputs:
+//   active_duty[n] (f32)
+//   boost_strength (f32)
+//
+// Output:
+//   boost[n] (f32) = expf(-boost_strength * (active_duty[c] - mean))
+//
+// Launch: single block (1024 threads), shared mem for reduction. At n=2048
+// each thread handles 2 elements.
+
+extern "C" __global__
+void sp_boost_from_duty(
+    const float * __restrict__ active_duty,  // (n,)
+    float       * __restrict__ boost,        // (n,) in-place out
+    float         boost_strength,
+    unsigned int  n
+) {
+    extern __shared__ float smem_raw[];
+    float * smem = smem_raw;
+    const unsigned int tid = threadIdx.x;
+    const unsigned int bsz = blockDim.x;
+
+    // Phase 1: parallel sum of active_duty into smem[0..32] (warp-level).
+    float local_sum = 0.0f;
+    for (unsigned int i = tid; i < n; i += bsz) {
+        local_sum += active_duty[i];
+    }
+    // Warp reduction.
+    for (int off = 16; off > 0; off >>= 1) {
+        local_sum += __shfl_down_sync(0xffffffff, local_sum, off);
+    }
+    unsigned int lane = tid & 31;
+    unsigned int warp = tid >> 5;
+    if (lane == 0) smem[warp] = local_sum;
+    __syncthreads();
+
+    // Warp 0 reduces warp-sums.
+    __shared__ float mean_s;
+    if (warp == 0) {
+        unsigned int nwarps = (bsz + 31) / 32;
+        float v = (lane < nwarps) ? smem[lane] : 0.0f;
+        for (int off = 16; off > 0; off >>= 1) {
+            v += __shfl_down_sync(0xffffffff, v, off);
+        }
+        if (tid == 0) {
+            mean_s = v / (float)n;
+        }
+    }
+    __syncthreads();
+
+    // Phase 2: boost[c] = expf(-strength * (active_duty[c] - mean)).
+    float mean = mean_s;
+    for (unsigned int i = tid; i < n; i += bsz) {
+        float d = active_duty[i] - mean;
+        boost[i] = expf(-boost_strength * d);
+    }
+}

overlay/htm_rust/src/gpu/kernels/sp_duty.cu

@@ -1,45 +1,45 @@
-// Duty cycle + boost update kernel.
-//
-// For each column c (one thread each):
-//   active_sample    = active_mask[c] ? 1 : 0
-//   overlap_sample   = raw_overlap[c] >= stim_thr ? 1 : 0
-//   active_duty[c]   = (1-alpha) * active_duty[c]  + alpha * active_sample
-//   overlap_duty[c]  = (1-alpha) * overlap_duty[c] + alpha * overlap_sample
-//
-// Then, if learn:
-//   boost[c] = exp(-boost_strength * (active_duty[c] - mean_duty))
-// mean_duty is computed on the host (one reduction) and passed in.
-
-extern "C" __global__
-void sp_duty_update(
-    const unsigned char * __restrict__ active_mask,  // (n_columns,)
-    const unsigned int  * __restrict__ raw_overlap,  // (n_columns,)
-    float               * __restrict__ active_duty,  // (n_columns,) in-place
-    float               * __restrict__ overlap_duty, // (n_columns,) in-place
-    float               * __restrict__ boost,        // (n_columns,) in-place
-    float                 alpha,
-    float                 stim_thr,
-    float                 boost_strength,            // 0 to skip boost
-    float                 mean_duty,
-    unsigned int          learn_flag,                // 0 or 1
-    unsigned int          n_columns
-) {
-    unsigned int c = blockIdx.x * blockDim.x + threadIdx.x;
-    if (c >= n_columns) return;
-
-    float ad = active_duty[c];
-    float od = overlap_duty[c];
-
-    float a_sample = (active_mask[c] != 0) ? 1.0f : 0.0f;
-    float o_sample = ((float)raw_overlap[c] >= stim_thr) ? 1.0f : 0.0f;
-
-    ad = (1.0f - alpha) * ad + alpha * a_sample;
-    od = (1.0f - alpha) * od + alpha * o_sample;
-
-    active_duty[c]  = ad;
-    overlap_duty[c] = od;
-
-    if (learn_flag && boost_strength > 0.0f) {
-        boost[c] = expf(-boost_strength * (ad - mean_duty));
-    }
-}
+// Duty cycle + boost update kernel.
+//
+// For each column c (one thread each):
+//   active_sample    = active_mask[c] ? 1 : 0
+//   overlap_sample   = raw_overlap[c] >= stim_thr ? 1 : 0
+//   active_duty[c]   = (1-alpha) * active_duty[c]  + alpha * active_sample
+//   overlap_duty[c]  = (1-alpha) * overlap_duty[c] + alpha * overlap_sample
+//
+// Then, if learn:
+//   boost[c] = exp(-boost_strength * (active_duty[c] - mean_duty))
+// mean_duty is computed on the host (one reduction) and passed in.
+
+extern "C" __global__
+void sp_duty_update(
+    const unsigned char * __restrict__ active_mask,  // (n_columns,)
+    const unsigned int  * __restrict__ raw_overlap,  // (n_columns,)
+    float               * __restrict__ active_duty,  // (n_columns,) in-place
+    float               * __restrict__ overlap_duty, // (n_columns,) in-place
+    float               * __restrict__ boost,        // (n_columns,) in-place
+    float                 alpha,
+    float                 stim_thr,
+    float                 boost_strength,            // 0 to skip boost
+    float                 mean_duty,
+    unsigned int          learn_flag,                // 0 or 1
+    unsigned int          n_columns
+) {
+    unsigned int c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (c >= n_columns) return;
+
+    float ad = active_duty[c];
+    float od = overlap_duty[c];
+
+    float a_sample = (active_mask[c] != 0) ? 1.0f : 0.0f;
+    float o_sample = ((float)raw_overlap[c] >= stim_thr) ? 1.0f : 0.0f;
+
+    ad = (1.0f - alpha) * ad + alpha * a_sample;
+    od = (1.0f - alpha) * od + alpha * o_sample;
+
+    active_duty[c]  = ad;
+    overlap_duty[c] = od;
+
+    if (learn_flag && boost_strength > 0.0f) {
+        boost[c] = expf(-boost_strength * (ad - mean_duty));
+    }
+}

overlay/htm_rust/src/gpu/kernels/sp_learn.cu

@@ -1,45 +1,45 @@
-// SP Hebbian learning kernel.
-//
-// For each active (winner) column c, for each of its synapses s:
-//   if input[bit[c][s]] active: perm += inc
-//   else:                       perm -= dec
-// Clamp to [0, 1].
-//
-// Launch: one block per column (2048 blocks), but we predicate on
-// active_mask[c] to avoid launching k-specific blocks.
-//
-// This matches the CPU reference line-for-line:
-//   src/sp.rs lines 157-169.
-
-extern "C" __global__
-void sp_learn(
-    const unsigned char * __restrict__ active_mask,  // (n_columns,) 0/1
-    const unsigned char * __restrict__ inp,          // (input_bits,)
-    const unsigned int  * __restrict__ syn_bit,      // (n_columns * S,)
-    float               * __restrict__ syn_perm,     // (n_columns * S,) in-place
-    float                 inc,
-    float                 dec,
-    unsigned int          synapses_per_col,
-    unsigned int          n_columns
-) {
-    const unsigned int c = blockIdx.x;
-    if (c >= n_columns) return;
-    if (active_mask[c] == 0) return;
-
-    const unsigned int base = c * synapses_per_col;
-    const unsigned int tid = threadIdx.x;
-    const unsigned int bsz = blockDim.x;
-
-    for (unsigned int s = tid; s < synapses_per_col; s += bsz) {
-        unsigned int b = syn_bit[base + s];
-        float p = syn_perm[base + s];
-        if (inp[b] != 0) {
-            p += inc;
-            if (p > 1.0f) p = 1.0f;
-        } else {
-            p -= dec;
-            if (p < 0.0f) p = 0.0f;
-        }
-        syn_perm[base + s] = p;
-    }
-}
+// SP Hebbian learning kernel.
+//
+// For each active (winner) column c, for each of its synapses s:
+//   if input[bit[c][s]] active: perm += inc
+//   else:                       perm -= dec
+// Clamp to [0, 1].
+//
+// Launch: one block per column (2048 blocks), but we predicate on
+// active_mask[c] to avoid launching k-specific blocks.
+//
+// This matches the CPU reference line-for-line:
+//   src/sp.rs lines 157-169.
+
+extern "C" __global__
+void sp_learn(
+    const unsigned char * __restrict__ active_mask,  // (n_columns,) 0/1
+    const unsigned char * __restrict__ inp,          // (input_bits,)
+    const unsigned int  * __restrict__ syn_bit,      // (n_columns * S,)
+    float               * __restrict__ syn_perm,     // (n_columns * S,) in-place
+    float                 inc,
+    float                 dec,
+    unsigned int          synapses_per_col,
+    unsigned int          n_columns
+) {
+    const unsigned int c = blockIdx.x;
+    if (c >= n_columns) return;
+    if (active_mask[c] == 0) return;
+
+    const unsigned int base = c * synapses_per_col;
+    const unsigned int tid = threadIdx.x;
+    const unsigned int bsz = blockDim.x;
+
+    for (unsigned int s = tid; s < synapses_per_col; s += bsz) {
+        unsigned int b = syn_bit[base + s];
+        float p = syn_perm[base + s];
+        if (inp[b] != 0) {
+            p += inc;
+            if (p > 1.0f) p = 1.0f;
+        } else {
+            p -= dec;
+            if (p < 0.0f) p = 0.0f;
+        }
+        syn_perm[base + s] = p;
+    }
+}

overlay/htm_rust/src/gpu/kernels/sp_overlap.cu

@@ -1,78 +1,78 @@
-// SP overlap kernel.
-//
-// For each column c (one CUDA block), compute:
-//   overlap[c] = sum over its synapse list of {inp[bit[c][s]] && perm[c][s] >= conn_thr}
-//   boosted[c] = overlap[c] * boost[c]
-//   raw_overlap[c] = overlap[c]   (also returned so host can drive duty cycle)
-//
-// Memory layout (flat, column-major with per-column stride = synapses_per_col):
-//   syn_bit[c * S + s]  : u32 index into input SDR
-//   syn_perm[c * S + s] : f32 permanence in [0, 1]
-//   boost[c]           : f32
-//   inp[b]             : u8 0/1
-// Output:
-//   raw[c]     : u32
-//   boosted[c] : f32
-//
-// Launch:
-//   grid  = n_columns
-//   block = 128 (or 256) — one warp-sweep across synapses; many warps give
-//                          parallel reduction across S (typically S=40).
-//
-// At S=40 this is completely latency-bound; we coalesce reads and do a
-// warp-shuffle reduction. For clarity we use a simple block-wide shared-mem
-// reduction which is sufficient for S <= 1024 and has zero correctness risk.
-
-extern "C" __global__
-void sp_overlap(
-    const unsigned char * __restrict__ inp,     // (input_bits,)
-    const unsigned int  * __restrict__ syn_bit, // (n_columns * S,)
-    const float         * __restrict__ syn_perm,// (n_columns * S,)
-    const float         * __restrict__ boost,   // (n_columns,)
-    float                 conn_thr,
-    unsigned int          synapses_per_col,     // S
-    unsigned int          n_columns,
-    unsigned int        * __restrict__ raw_out,     // (n_columns,)
-    float               * __restrict__ boosted_out  // (n_columns,)
-) {
-    const unsigned int c = blockIdx.x;
-    if (c >= n_columns) return;
-
-    const unsigned int base = c * synapses_per_col;
-    const unsigned int tid = threadIdx.x;
-    const unsigned int bsz = blockDim.x;
-
-    // Per-thread partial count.
-    unsigned int local = 0;
-    for (unsigned int s = tid; s < synapses_per_col; s += bsz) {
-        unsigned int b = syn_bit[base + s];
-        float p = syn_perm[base + s];
-        // Branchless: only counts when input active AND perm connected.
-        // Using (inp != 0) to tolerate u8 layout.
-        unsigned int hit = ((inp[b] != 0) && (p >= conn_thr)) ? 1u : 0u;
-        local += hit;
-    }
-
-    // Block-wide reduction in shared memory.
-    __shared__ unsigned int smem[32];
-
-    // Warp-level reduction via shuffle.
-    unsigned int lane = tid & 31;
-    unsigned int warp = tid >> 5;
-    for (int off = 16; off > 0; off >>= 1) {
-        local += __shfl_down_sync(0xffffffff, local, off);
-    }
-    if (lane == 0) smem[warp] = local;
-    __syncthreads();
-
-    if (warp == 0) {
-        unsigned int v = (tid < (bsz + 31) / 32) ? smem[lane] : 0;
-        for (int off = 16; off > 0; off >>= 1) {
-            v += __shfl_down_sync(0xffffffff, v, off);
-        }
-        if (tid == 0) {
-            raw_out[c] = v;
-            boosted_out[c] = (float)v * boost[c];
-        }
-    }
-}
+// SP overlap kernel.
+//
+// For each column c (one CUDA block), compute:
+//   overlap[c] = sum over its synapse list of {inp[bit[c][s]] && perm[c][s] >= conn_thr}
+//   boosted[c] = overlap[c] * boost[c]
+//   raw_overlap[c] = overlap[c]   (also returned so host can drive duty cycle)
+//
+// Memory layout (flat, column-major with per-column stride = synapses_per_col):
+//   syn_bit[c * S + s]  : u32 index into input SDR
+//   syn_perm[c * S + s] : f32 permanence in [0, 1]
+//   boost[c]           : f32
+//   inp[b]             : u8 0/1
+// Output:
+//   raw[c]     : u32
+//   boosted[c] : f32
+//
+// Launch:
+//   grid  = n_columns
+//   block = 128 (or 256) — one warp-sweep across synapses; many warps give
+//                          parallel reduction across S (typically S=40).
+//
+// At S=40 this is completely latency-bound; we coalesce reads and do a
+// warp-shuffle reduction. For clarity we use a simple block-wide shared-mem
+// reduction which is sufficient for S <= 1024 and has zero correctness risk.
+
+extern "C" __global__
+void sp_overlap(
+    const unsigned char * __restrict__ inp,     // (input_bits,)
+    const unsigned int  * __restrict__ syn_bit, // (n_columns * S,)
+    const float         * __restrict__ syn_perm,// (n_columns * S,)
+    const float         * __restrict__ boost,   // (n_columns,)
+    float                 conn_thr,
+    unsigned int          synapses_per_col,     // S
+    unsigned int          n_columns,
+    unsigned int        * __restrict__ raw_out,     // (n_columns,)
+    float               * __restrict__ boosted_out  // (n_columns,)
+) {
+    const unsigned int c = blockIdx.x;
+    if (c >= n_columns) return;
+
+    const unsigned int base = c * synapses_per_col;
+    const unsigned int tid = threadIdx.x;
+    const unsigned int bsz = blockDim.x;
+
+    // Per-thread partial count.
+    unsigned int local = 0;
+    for (unsigned int s = tid; s < synapses_per_col; s += bsz) {
+        unsigned int b = syn_bit[base + s];
+        float p = syn_perm[base + s];
+        // Branchless: only counts when input active AND perm connected.
+        // Using (inp != 0) to tolerate u8 layout.
+        unsigned int hit = ((inp[b] != 0) && (p >= conn_thr)) ? 1u : 0u;
+        local += hit;
+    }
+
+    // Block-wide reduction in shared memory.
+    __shared__ unsigned int smem[32];
+
+    // Warp-level reduction via shuffle.
+    unsigned int lane = tid & 31;
+    unsigned int warp = tid >> 5;
+    for (int off = 16; off > 0; off >>= 1) {
+        local += __shfl_down_sync(0xffffffff, local, off);
+    }
+    if (lane == 0) smem[warp] = local;
+    __syncthreads();
+
+    if (warp == 0) {
+        unsigned int v = (tid < (bsz + 31) / 32) ? smem[lane] : 0;
+        for (int off = 16; off > 0; off >>= 1) {
+            v += __shfl_down_sync(0xffffffff, v, off);
+        }
+        if (tid == 0) {
+            raw_out[c] = v;
+            boosted_out[c] = (float)v * boost[c];
+        }
+    }
+}

overlay/htm_rust/src/gpu/kernels/sp_topk.cu

@@ -1,117 +1,117 @@
-// Top-K column selection.
-//
-// Inputs:
-//   boosted[n_columns] : f32 score
-// Output:
-//   active_mask[n_columns] : u8 0/1, exactly k ones
-//
-// Tie-breaking: when scores are equal, the LOWER column index wins (matches
-// CPU reference `select_nth_unstable_by` with secondary index comparator).
-//
-// Strategy: a single-block implementation. n_columns is typically 2048, which
-// fits comfortably in shared memory. We use a bitonic top-k via per-thread
-// radix-select of the (score, -index) key. At k≈41 of n=2048 the simplest
-// correct approach is a thresholding pass:
-//
-//   1. Radix-like bucket pass to find the k-th largest score.
-//   2. Mark winners = strictly-greater-than-threshold AND ties until count hits k.
-//
-// For strict index-ordered tie-break we materialise a 64-bit key:
-//   key = (float_to_sortable_u32(score) << 32) | (0xffffffff - index)
-// Larger key = (higher score) OR (same score, smaller index).
-//
-// Then we find the k-th largest 64-bit key via radix-select and mark all
-// columns whose key >= threshold. This is O(n_cols * log k) and well under
-// 100 μs for n=2048, k=41 on sm_86.
-//
-// For simplicity and correctness this kernel uses a single-block parallel
-// selection sort variant (find max → mark → zero → repeat, k iterations).
-// At k=41 this is 41 passes of 2048 threads = ~2048*41 = 84K ops, trivially
-// fast.
-
-extern "C" __global__
-void sp_topk_select(
-    const float * __restrict__ scores,    // (n_columns,)
-    unsigned int  n_columns,
-    unsigned int  k,
-    unsigned char * __restrict__ active_out  // (n_columns,)
-) {
-    extern __shared__ float smem[];
-    // Layout: smem[0..n] = working scores (we'll mark selected entries as -inf)
-    //         smem[n..n+32*2] = reduction scratch (score + index, per warp)
-    float * work = smem;
-    const unsigned int tid = threadIdx.x;
-    const unsigned int bsz = blockDim.x;
-
-    // Load scores into shared; also init active_out = 0.
-    for (unsigned int i = tid; i < n_columns; i += bsz) {
-        work[i] = scores[i];
-        active_out[i] = 0;
-    }
-    __syncthreads();
-
-    __shared__ int   winner_idx;
-    __shared__ float winner_score;
-
-    for (unsigned int iter = 0; iter < k; ++iter) {
-        // Find (argmax score, lowest index for ties).
-        float best_s = -INFINITY;
-        int   best_i = n_columns;   // sentinel larger than any index
-
-        for (unsigned int i = tid; i < n_columns; i += bsz) {
-            float s = work[i];
-            if (s > best_s || (s == best_s && (int)i < best_i)) {
-                best_s = s;
-                best_i = (int)i;
-            }
-        }
-
-        // Warp reduction. We reduce pairs (score, idx) keeping (max score, min idx on tie).
-        unsigned int mask = 0xffffffff;
-        for (int off = 16; off > 0; off >>= 1) {
-            float os = __shfl_down_sync(mask, best_s, off);
-            int   oi = __shfl_down_sync(mask, best_i, off);
-            if (os > best_s || (os == best_s && oi < best_i)) {
-                best_s = os;
-                best_i = oi;
-            }
-        }
-        // Warp 0 collects lane 0 values from other warps via shared mem.
-        __shared__ float warp_s[32];
-        __shared__ int   warp_i[32];
-        unsigned int lane = tid & 31;
-        unsigned int warp = tid >> 5;
-        if (lane == 0) {
-            warp_s[warp] = best_s;
-            warp_i[warp] = best_i;
-        }
-        __syncthreads();
-
-        if (warp == 0) {
-            unsigned int nwarps = (bsz + 31) / 32;
-            float s = (lane < nwarps) ? warp_s[lane] : -INFINITY;
-            int   i = (lane < nwarps) ? warp_i[lane] : (int)n_columns;
-            for (int off = 16; off > 0; off >>= 1) {
-                float os = __shfl_down_sync(mask, s, off);
-                int   oi = __shfl_down_sync(mask, i, off);
-                if (os > s || (os == s && oi < i)) {
-                    s = os;
-                    i = oi;
-                }
-            }
-            if (tid == 0) {
-                winner_score = s;
-                winner_idx = i;
-            }
-        }
-        __syncthreads();
-
-        if (tid == 0) {
-            if (winner_idx < (int)n_columns) {
-                active_out[winner_idx] = 1;
-                work[winner_idx] = -INFINITY;
-            }
-        }
-        __syncthreads();
-    }
-}
+// Top-K column selection.
+//
+// Inputs:
+//   boosted[n_columns] : f32 score
+// Output:
+//   active_mask[n_columns] : u8 0/1, exactly k ones
+//
+// Tie-breaking: when scores are equal, the LOWER column index wins (matches
+// CPU reference `select_nth_unstable_by` with secondary index comparator).
+//
+// Strategy: a single-block implementation. n_columns is typically 2048, which
+// fits comfortably in shared memory. We use a bitonic top-k via per-thread
+// radix-select of the (score, -index) key. At k≈41 of n=2048 the simplest
+// correct approach is a thresholding pass:
+//
+//   1. Radix-like bucket pass to find the k-th largest score.
+//   2. Mark winners = strictly-greater-than-threshold AND ties until count hits k.
+//
+// For strict index-ordered tie-break we materialise a 64-bit key:
+//   key = (float_to_sortable_u32(score) << 32) | (0xffffffff - index)
+// Larger key = (higher score) OR (same score, smaller index).
+//
+// Then we find the k-th largest 64-bit key via radix-select and mark all
+// columns whose key >= threshold. This is O(n_cols * log k) and well under
+// 100 μs for n=2048, k=41 on sm_86.
+//
+// For simplicity and correctness this kernel uses a single-block parallel
+// selection sort variant (find max → mark → zero → repeat, k iterations).
+// At k=41 this is 41 passes of 2048 threads = ~2048*41 = 84K ops, trivially
+// fast.
+
+extern "C" __global__
+void sp_topk_select(
+    const float * __restrict__ scores,    // (n_columns,)
+    unsigned int  n_columns,
+    unsigned int  k,
+    unsigned char * __restrict__ active_out  // (n_columns,)
+) {
+    extern __shared__ float smem[];
+    // Layout: smem[0..n] = working scores (we'll mark selected entries as -inf)
+    //         smem[n..n+32*2] = reduction scratch (score + index, per warp)
+    float * work = smem;
+    const unsigned int tid = threadIdx.x;
+    const unsigned int bsz = blockDim.x;
+
+    // Load scores into shared; also init active_out = 0.
+    for (unsigned int i = tid; i < n_columns; i += bsz) {
+        work[i] = scores[i];
+        active_out[i] = 0;
+    }
+    __syncthreads();
+
+    __shared__ int   winner_idx;
+    __shared__ float winner_score;
+
+    for (unsigned int iter = 0; iter < k; ++iter) {
+        // Find (argmax score, lowest index for ties).
+        float best_s = -INFINITY;
+        int   best_i = n_columns;   // sentinel larger than any index
+
+        for (unsigned int i = tid; i < n_columns; i += bsz) {
+            float s = work[i];
+            if (s > best_s || (s == best_s && (int)i < best_i)) {
+                best_s = s;
+                best_i = (int)i;
+            }
+        }
+
+        // Warp reduction. We reduce pairs (score, idx) keeping (max score, min idx on tie).
+        unsigned int mask = 0xffffffff;
+        for (int off = 16; off > 0; off >>= 1) {
+            float os = __shfl_down_sync(mask, best_s, off);
+            int   oi = __shfl_down_sync(mask, best_i, off);
+            if (os > best_s || (os == best_s && oi < best_i)) {
+                best_s = os;
+                best_i = oi;
+            }
+        }
+        // Warp 0 collects lane 0 values from other warps via shared mem.
+        __shared__ float warp_s[32];
+        __shared__ int   warp_i[32];
+        unsigned int lane = tid & 31;
+        unsigned int warp = tid >> 5;
+        if (lane == 0) {
+            warp_s[warp] = best_s;
+            warp_i[warp] = best_i;
+        }
+        __syncthreads();
+
+        if (warp == 0) {
+            unsigned int nwarps = (bsz + 31) / 32;
+            float s = (lane < nwarps) ? warp_s[lane] : -INFINITY;
+            int   i = (lane < nwarps) ? warp_i[lane] : (int)n_columns;
+            for (int off = 16; off > 0; off >>= 1) {
+                float os = __shfl_down_sync(mask, s, off);
+                int   oi = __shfl_down_sync(mask, i, off);
+                if (os > s || (os == s && oi < i)) {
+                    s = os;
+                    i = oi;
+                }
+            }
+            if (tid == 0) {
+                winner_score = s;
+                winner_idx = i;
+            }
+        }
+        __syncthreads();
+
+        if (tid == 0) {
+            if (winner_idx < (int)n_columns) {
+                active_out[winner_idx] = 1;
+                work[winner_idx] = -INFINITY;
+            }
+        }
+        __syncthreads();
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_activate.cu

@@ -1,66 +1,66 @@
-// TM activate kernel. See tm_predict.cu for TmConfig.
-
-struct TmConfig {
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int cells_per_column;
-    unsigned int synapses_per_segment;
-    unsigned int n_segments;
-    unsigned int n_cells;
-    unsigned int max_segments_per_cell;
-    unsigned int max_new_synapses;
-    int conn_thr_i16;
-    int perm_inc_i16;
-    int perm_dec_i16;
-    int predicted_seg_dec_i16;
-    int initial_perm_i16;
-    unsigned int iter_seed;
-    unsigned int n_cols;
-    unsigned int bits_words;
-};
-
-extern "C" __global__
-void tm_activate(
-    const unsigned char * __restrict__ sp_active_mask,
-    const unsigned char * __restrict__ col_predicted,
-    const unsigned int  * __restrict__ cell_predictive_bits,
-    unsigned int        * __restrict__ cell_active_bits,
-    unsigned int        * __restrict__ cell_winner_bits,
-    unsigned int        * __restrict__ unpredicted_count,
-    unsigned int        * __restrict__ burst_cols_flat,
-    unsigned int        * __restrict__ burst_cols_count,
-    TmConfig              cfg
-) {
-    unsigned int col = blockIdx.x * blockDim.x + threadIdx.x;
-    if (col >= cfg.n_cols) return;
-    if (sp_active_mask[col] == 0) return;
-
-    unsigned int base_cell = col * cfg.cells_per_column;
-
-    if (col_predicted[col]) {
-        for (unsigned int k = 0; k < cfg.cells_per_column; k++) {
-            unsigned int cell = base_cell + k;
-            unsigned int word_idx = cell >> 5;
-            unsigned int bit_mask = 1u << (cell & 31u);
-            unsigned int pred_word = cell_predictive_bits[word_idx];
-            if (pred_word & bit_mask) {
-                atomicOr(&cell_active_bits[word_idx], bit_mask);
-                atomicOr(&cell_winner_bits[word_idx], bit_mask);
-            }
-        }
-    } else {
-        atomicAdd(unpredicted_count, 1u);
-        for (unsigned int k = 0; k < cfg.cells_per_column; k++) {
-            unsigned int cell = base_cell + k;
-            unsigned int word_idx = cell >> 5;
-            unsigned int bit_mask = 1u << (cell & 31u);
-            atomicOr(&cell_active_bits[word_idx], bit_mask);
-        }
-        unsigned int winner = base_cell;
-        unsigned int word_idx = winner >> 5;
-        unsigned int bit_mask = 1u << (winner & 31u);
-        atomicOr(&cell_winner_bits[word_idx], bit_mask);
-        unsigned int slot = atomicAdd(burst_cols_count, 1u);
-        burst_cols_flat[slot] = col;
-    }
-}
+// TM activate kernel. See tm_predict.cu for TmConfig.
+
+struct TmConfig {
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int cells_per_column;
+    unsigned int synapses_per_segment;
+    unsigned int n_segments;
+    unsigned int n_cells;
+    unsigned int max_segments_per_cell;
+    unsigned int max_new_synapses;
+    int conn_thr_i16;
+    int perm_inc_i16;
+    int perm_dec_i16;
+    int predicted_seg_dec_i16;
+    int initial_perm_i16;
+    unsigned int iter_seed;
+    unsigned int n_cols;
+    unsigned int bits_words;
+};
+
+extern "C" __global__
+void tm_activate(
+    const unsigned char * __restrict__ sp_active_mask,
+    const unsigned char * __restrict__ col_predicted,
+    const unsigned int  * __restrict__ cell_predictive_bits,
+    unsigned int        * __restrict__ cell_active_bits,
+    unsigned int        * __restrict__ cell_winner_bits,
+    unsigned int        * __restrict__ unpredicted_count,
+    unsigned int        * __restrict__ burst_cols_flat,
+    unsigned int        * __restrict__ burst_cols_count,
+    TmConfig              cfg
+) {
+    unsigned int col = blockIdx.x * blockDim.x + threadIdx.x;
+    if (col >= cfg.n_cols) return;
+    if (sp_active_mask[col] == 0) return;
+
+    unsigned int base_cell = col * cfg.cells_per_column;
+
+    if (col_predicted[col]) {
+        for (unsigned int k = 0; k < cfg.cells_per_column; k++) {
+            unsigned int cell = base_cell + k;
+            unsigned int word_idx = cell >> 5;
+            unsigned int bit_mask = 1u << (cell & 31u);
+            unsigned int pred_word = cell_predictive_bits[word_idx];
+            if (pred_word & bit_mask) {
+                atomicOr(&cell_active_bits[word_idx], bit_mask);
+                atomicOr(&cell_winner_bits[word_idx], bit_mask);
+            }
+        }
+    } else {
+        atomicAdd(unpredicted_count, 1u);
+        for (unsigned int k = 0; k < cfg.cells_per_column; k++) {
+            unsigned int cell = base_cell + k;
+            unsigned int word_idx = cell >> 5;
+            unsigned int bit_mask = 1u << (cell & 31u);
+            atomicOr(&cell_active_bits[word_idx], bit_mask);
+        }
+        unsigned int winner = base_cell;
+        unsigned int word_idx = winner >> 5;
+        unsigned int bit_mask = 1u << (winner & 31u);
+        atomicOr(&cell_winner_bits[word_idx], bit_mask);
+        unsigned int slot = atomicAdd(burst_cols_count, 1u);
+        burst_cols_flat[slot] = col;
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_anomaly.cu

@@ -1,43 +1,43 @@
-// TM anomaly kernel.
-//
-// Computes:
-//   n_active = sum of sp_active_mask
-//   anomaly  = unpredicted_count / n_active   (if n_active > 0)
-//            = 0                              (else)
-//
-// Launch: single block, 256 threads.
-
-extern "C" __global__
-void tm_anomaly(
-    const unsigned char * __restrict__ sp_active_mask,
-    const unsigned int  * __restrict__ unpredicted_count,
-    float               * __restrict__ anomaly_out,       // (1,) or (t_slot,)
-    unsigned int          t_slot,
-    unsigned int          n_cols
-) {
-    const unsigned int tid = threadIdx.x;
-    __shared__ unsigned int n_active_s;
-
-    if (tid == 0) n_active_s = 0u;
-    __syncthreads();
-
-    unsigned int local = 0u;
-    for (unsigned int i = tid; i < n_cols; i += blockDim.x) {
-        if (sp_active_mask[i]) local += 1u;
-    }
-    // Warp reduce.
-    for (int off = 16; off > 0; off >>= 1) {
-        local += __shfl_down_sync(0xffffffffu, local, off);
-    }
-    if ((tid & 31u) == 0) {
-        atomicAdd(&n_active_s, local);
-    }
-    __syncthreads();
-
-    if (tid == 0) {
-        unsigned int total = n_active_s;
-        unsigned int bad = unpredicted_count[0];
-        float anom = (total > 0u) ? ((float)bad / (float)total) : 0.0f;
-        anomaly_out[t_slot] = anom;
-    }
-}
+// TM anomaly kernel.
+//
+// Computes:
+//   n_active = sum of sp_active_mask
+//   anomaly  = unpredicted_count / n_active   (if n_active > 0)
+//            = 0                              (else)
+//
+// Launch: single block, 256 threads.
+
+extern "C" __global__
+void tm_anomaly(
+    const unsigned char * __restrict__ sp_active_mask,
+    const unsigned int  * __restrict__ unpredicted_count,
+    float               * __restrict__ anomaly_out,       // (1,) or (t_slot,)
+    unsigned int          t_slot,
+    unsigned int          n_cols
+) {
+    const unsigned int tid = threadIdx.x;
+    __shared__ unsigned int n_active_s;
+
+    if (tid == 0) n_active_s = 0u;
+    __syncthreads();
+
+    unsigned int local = 0u;
+    for (unsigned int i = tid; i < n_cols; i += blockDim.x) {
+        if (sp_active_mask[i]) local += 1u;
+    }
+    // Warp reduce.
+    for (int off = 16; off > 0; off >>= 1) {
+        local += __shfl_down_sync(0xffffffffu, local, off);
+    }
+    if ((tid & 31u) == 0) {
+        atomicAdd(&n_active_s, local);
+    }
+    __syncthreads();
+
+    if (tid == 0) {
+        unsigned int total = n_active_s;
+        unsigned int bad = unpredicted_count[0];
+        float anom = (total > 0u) ? ((float)bad / (float)total) : 0.0f;
+        anomaly_out[t_slot] = anom;
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_grow.cu

@@ -1,155 +1,155 @@
-// TM grow+reinforce kernel.
-//
-// For each bursting column:
-//   If col_best_match[col] is non-zero (i.e. at least one matching segment
-//     with num_active_potential >= learning_threshold exists on cells in this col):
-//     Target = that matching segment.
-//     Reinforce its existing synapses: +inc if presyn in prev_active, -dec otherwise.
-//     Grow up to (max_new - current_syn_count) additional synapses to prev_winners.
-//   Else:
-//     Allocate a fresh segment slot on winner cell (cell 0 of col).
-//     Grow up to max_new synapses to prev_winners (no reinforce needed — new seg).
-//
-// This mirrors the CPU TM burst logic.
-
-struct TmConfig {
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int cells_per_column;
-    unsigned int synapses_per_segment;
-    unsigned int n_segments;
-    unsigned int n_cells;
-    unsigned int max_segments_per_cell;
-    unsigned int max_new_synapses;
-    int conn_thr_i16;
-    int perm_inc_i16;
-    int perm_dec_i16;
-    int predicted_seg_dec_i16;
-    int initial_perm_i16;
-    unsigned int iter_seed;
-    unsigned int n_cols;
-    unsigned int bits_words;
-};
-
-extern "C" __global__
-void tm_grow(
-    unsigned int       * __restrict__ seg_cell_id,
-    unsigned int       * __restrict__ seg_syn_count,
-    unsigned int       * __restrict__ syn_presyn,
-    short              * __restrict__ syn_perm,
-    unsigned int       * __restrict__ cell_seg_count,
-    const unsigned int * __restrict__ burst_cols_flat,
-    const unsigned int * __restrict__ burst_cols_count,
-    const unsigned int * __restrict__ prev_winner_bits,
-    const unsigned int * __restrict__ prev_active_bits,
-    const unsigned int * __restrict__ col_best_match,
-    TmConfig             cfg
-) {
-    const unsigned int b = blockIdx.x;
-    const unsigned int n_burst_cols = burst_cols_count[0];
-    if (b >= n_burst_cols) return;
-    const unsigned int tid = threadIdx.x;
-
-    const unsigned int col = burst_cols_flat[b];
-
-    __shared__ unsigned int shared_seg_id;
-    __shared__ unsigned int shared_existing_syn_count;
-    __shared__ unsigned int shared_grown;
-    __shared__ unsigned int shared_is_new;
-    __shared__ unsigned int shared_start_offset;
-
-    if (tid == 0) {
-        unsigned int match_key = col_best_match[col];
-        if (match_key != 0u) {
-            // Reuse matching segment.
-            unsigned int seg_id = match_key & 0x1FFFFFu;
-            shared_seg_id = seg_id;
-            shared_existing_syn_count = seg_syn_count[seg_id];
-            shared_is_new = 0u;
-        } else {
-            // Allocate new segment on winner cell (cell 0 of col).
-            unsigned int winner_cell = col * cfg.cells_per_column;
-            unsigned int slot = atomicAdd(&cell_seg_count[winner_cell], 1u);
-            if (slot >= cfg.max_segments_per_cell) {
-                slot = slot % cfg.max_segments_per_cell;
-            }
-            unsigned int seg_id = winner_cell * cfg.max_segments_per_cell + slot;
-            seg_cell_id[seg_id] = winner_cell;
-            seg_syn_count[seg_id] = 0;
-            shared_seg_id = seg_id;
-            shared_existing_syn_count = 0u;
-            shared_is_new = 1u;
-        }
-        shared_grown = 0u;
-        shared_start_offset = (b * 2654435761u + cfg.iter_seed) % cfg.bits_words;
-    }
-    __syncthreads();
-
-    const unsigned int seg_id = shared_seg_id;
-    const unsigned int seg_base = seg_id * cfg.synapses_per_segment;
-    const unsigned int existing_syn = shared_existing_syn_count;
-    const unsigned int is_new = shared_is_new;
-    const unsigned int start = shared_start_offset;
-
-    // PHASE 1: If reusing, reinforce existing synapses.
-    if (!is_new) {
-        for (unsigned int s = tid; s < existing_syn; s += 32u) {
-            unsigned int presyn = syn_presyn[seg_base + s];
-            unsigned int word = prev_active_bits[presyn >> 5];
-            unsigned int bit = (word >> (presyn & 31u)) & 1u;
-            int p = (int)syn_perm[seg_base + s];
-            if (bit) {
-                int np = p + cfg.perm_inc_i16;
-                if (np > 32767) np = 32767;
-                syn_perm[seg_base + s] = (short)np;
-            } else {
-                int np = p - cfg.perm_dec_i16;
-                if (np < 0) np = 0;
-                syn_perm[seg_base + s] = (short)np;
-            }
-        }
-        __syncthreads();
-    }
-
-    // PHASE 2: Grow up to `max_new_synapses` (or room) synapses to prev_winners
-    // that aren't already presynaptic to this segment.
-    const unsigned int room = (cfg.synapses_per_segment > existing_syn)
-        ? (cfg.synapses_per_segment - existing_syn) : 0u;
-    const unsigned int max_grow = (cfg.max_new_synapses < room) ? cfg.max_new_synapses : room;
-
-    for (unsigned int w_off = 0; w_off < cfg.bits_words; w_off += 32u) {
-        if (shared_grown >= max_grow) break;
-        unsigned int widx = (start + w_off + tid) % cfg.bits_words;
-        unsigned int word = prev_winner_bits[widx];
-        while (word != 0u) {
-            if (shared_grown >= max_grow) break;
-            unsigned int bit_pos = __ffs(word) - 1u;
-            word &= ~(1u << bit_pos);
-            unsigned int cell = widx * 32u + bit_pos;
-            if (cell >= cfg.n_cells) continue;
-
-            // Skip if already presynaptic (O(existing_syn) scan; usually small).
-            bool exists = false;
-            for (unsigned int s = 0; s < existing_syn; s++) {
-                if (syn_presyn[seg_base + s] == cell) { exists = true; break; }
-            }
-            if (exists) continue;
-
-            unsigned int slot = atomicAdd(&shared_grown, 1u);
-            if (slot >= max_grow) break;
-            unsigned int write_idx = existing_syn + slot;
-            if (write_idx >= cfg.synapses_per_segment) break;
-            syn_presyn[seg_base + write_idx] = cell;
-            syn_perm[seg_base + write_idx] = (short)cfg.initial_perm_i16;
-        }
-    }
-    __syncthreads();
-
-    if (tid == 0) {
-        unsigned int grown = shared_grown;
-        if (grown > max_grow) grown = max_grow;
-        unsigned int new_count = existing_syn + grown;
-        if (new_count > cfg.synapses_per_segment) new_count = cfg.synapses_per_segment;
-        seg_syn_count[seg_id] = new_count;
-    }
-}
+// TM grow+reinforce kernel.
+//
+// For each bursting column:
+//   If col_best_match[col] is non-zero (i.e. at least one matching segment
+//     with num_active_potential >= learning_threshold exists on cells in this col):
+//     Target = that matching segment.
+//     Reinforce its existing synapses: +inc if presyn in prev_active, -dec otherwise.
+//     Grow up to (max_new - current_syn_count) additional synapses to prev_winners.
+//   Else:
+//     Allocate a fresh segment slot on winner cell (cell 0 of col).
+//     Grow up to max_new synapses to prev_winners (no reinforce needed — new seg).
+//
+// This mirrors the CPU TM burst logic.
+
+struct TmConfig {
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int cells_per_column;
+    unsigned int synapses_per_segment;
+    unsigned int n_segments;
+    unsigned int n_cells;
+    unsigned int max_segments_per_cell;
+    unsigned int max_new_synapses;
+    int conn_thr_i16;
+    int perm_inc_i16;
+    int perm_dec_i16;
+    int predicted_seg_dec_i16;
+    int initial_perm_i16;
+    unsigned int iter_seed;
+    unsigned int n_cols;
+    unsigned int bits_words;
+};
+
+extern "C" __global__
+void tm_grow(
+    unsigned int       * __restrict__ seg_cell_id,
+    unsigned int       * __restrict__ seg_syn_count,
+    unsigned int       * __restrict__ syn_presyn,
+    short              * __restrict__ syn_perm,
+    unsigned int       * __restrict__ cell_seg_count,
+    const unsigned int * __restrict__ burst_cols_flat,
+    const unsigned int * __restrict__ burst_cols_count,
+    const unsigned int * __restrict__ prev_winner_bits,
+    const unsigned int * __restrict__ prev_active_bits,
+    const unsigned int * __restrict__ col_best_match,
+    TmConfig             cfg
+) {
+    const unsigned int b = blockIdx.x;
+    const unsigned int n_burst_cols = burst_cols_count[0];
+    if (b >= n_burst_cols) return;
+    const unsigned int tid = threadIdx.x;
+
+    const unsigned int col = burst_cols_flat[b];
+
+    __shared__ unsigned int shared_seg_id;
+    __shared__ unsigned int shared_existing_syn_count;
+    __shared__ unsigned int shared_grown;
+    __shared__ unsigned int shared_is_new;
+    __shared__ unsigned int shared_start_offset;
+
+    if (tid == 0) {
+        unsigned int match_key = col_best_match[col];
+        if (match_key != 0u) {
+            // Reuse matching segment.
+            unsigned int seg_id = match_key & 0x1FFFFFu;
+            shared_seg_id = seg_id;
+            shared_existing_syn_count = seg_syn_count[seg_id];
+            shared_is_new = 0u;
+        } else {
+            // Allocate new segment on winner cell (cell 0 of col).
+            unsigned int winner_cell = col * cfg.cells_per_column;
+            unsigned int slot = atomicAdd(&cell_seg_count[winner_cell], 1u);
+            if (slot >= cfg.max_segments_per_cell) {
+                slot = slot % cfg.max_segments_per_cell;
+            }
+            unsigned int seg_id = winner_cell * cfg.max_segments_per_cell + slot;
+            seg_cell_id[seg_id] = winner_cell;
+            seg_syn_count[seg_id] = 0;
+            shared_seg_id = seg_id;
+            shared_existing_syn_count = 0u;
+            shared_is_new = 1u;
+        }
+        shared_grown = 0u;
+        shared_start_offset = (b * 2654435761u + cfg.iter_seed) % cfg.bits_words;
+    }
+    __syncthreads();
+
+    const unsigned int seg_id = shared_seg_id;
+    const unsigned int seg_base = seg_id * cfg.synapses_per_segment;
+    const unsigned int existing_syn = shared_existing_syn_count;
+    const unsigned int is_new = shared_is_new;
+    const unsigned int start = shared_start_offset;
+
+    // PHASE 1: If reusing, reinforce existing synapses.
+    if (!is_new) {
+        for (unsigned int s = tid; s < existing_syn; s += 32u) {
+            unsigned int presyn = syn_presyn[seg_base + s];
+            unsigned int word = prev_active_bits[presyn >> 5];
+            unsigned int bit = (word >> (presyn & 31u)) & 1u;
+            int p = (int)syn_perm[seg_base + s];
+            if (bit) {
+                int np = p + cfg.perm_inc_i16;
+                if (np > 32767) np = 32767;
+                syn_perm[seg_base + s] = (short)np;
+            } else {
+                int np = p - cfg.perm_dec_i16;
+                if (np < 0) np = 0;
+                syn_perm[seg_base + s] = (short)np;
+            }
+        }
+        __syncthreads();
+    }
+
+    // PHASE 2: Grow up to `max_new_synapses` (or room) synapses to prev_winners
+    // that aren't already presynaptic to this segment.
+    const unsigned int room = (cfg.synapses_per_segment > existing_syn)
+        ? (cfg.synapses_per_segment - existing_syn) : 0u;
+    const unsigned int max_grow = (cfg.max_new_synapses < room) ? cfg.max_new_synapses : room;
+
+    for (unsigned int w_off = 0; w_off < cfg.bits_words; w_off += 32u) {
+        if (shared_grown >= max_grow) break;
+        unsigned int widx = (start + w_off + tid) % cfg.bits_words;
+        unsigned int word = prev_winner_bits[widx];
+        while (word != 0u) {
+            if (shared_grown >= max_grow) break;
+            unsigned int bit_pos = __ffs(word) - 1u;
+            word &= ~(1u << bit_pos);
+            unsigned int cell = widx * 32u + bit_pos;
+            if (cell >= cfg.n_cells) continue;
+
+            // Skip if already presynaptic (O(existing_syn) scan; usually small).
+            bool exists = false;
+            for (unsigned int s = 0; s < existing_syn; s++) {
+                if (syn_presyn[seg_base + s] == cell) { exists = true; break; }
+            }
+            if (exists) continue;
+
+            unsigned int slot = atomicAdd(&shared_grown, 1u);
+            if (slot >= max_grow) break;
+            unsigned int write_idx = existing_syn + slot;
+            if (write_idx >= cfg.synapses_per_segment) break;
+            syn_presyn[seg_base + write_idx] = cell;
+            syn_perm[seg_base + write_idx] = (short)cfg.initial_perm_i16;
+        }
+    }
+    __syncthreads();
+
+    if (tid == 0) {
+        unsigned int grown = shared_grown;
+        if (grown > max_grow) grown = max_grow;
+        unsigned int new_count = existing_syn + grown;
+        if (new_count > cfg.synapses_per_segment) new_count = cfg.synapses_per_segment;
+        seg_syn_count[seg_id] = new_count;
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_learn.cu

@@ -1,75 +1,75 @@
-// TM learn (reinforce correctly predicted segments) — cell-grouped launch.
-//
-// Grid: n_cells.
-// For each cell in a predicted, SP-active column: iterate its segments.
-// For each segment with num_active_connected >= activation_threshold,
-// reinforce its synapses against prev_active_bits.
-
-struct TmConfig {
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int cells_per_column;
-    unsigned int synapses_per_segment;
-    unsigned int n_segments;
-    unsigned int n_cells;
-    unsigned int max_segments_per_cell;
-    unsigned int max_new_synapses;
-    int conn_thr_i16;
-    int perm_inc_i16;
-    int perm_dec_i16;
-    int predicted_seg_dec_i16;
-    int initial_perm_i16;
-    unsigned int iter_seed;
-    unsigned int n_cols;
-    unsigned int bits_words;
-};
-
-extern "C" __global__
-void tm_learn_reinforce(
-    const unsigned int * __restrict__ seg_cell_id,
-    const unsigned int * __restrict__ seg_syn_count,
-    const unsigned int * __restrict__ syn_presyn,
-    short              * __restrict__ syn_perm,
-    const unsigned int * __restrict__ seg_num_active_connected,
-    const unsigned int * __restrict__ prev_active_bits,
-    const unsigned char * __restrict__ sp_active_mask,
-    const unsigned char * __restrict__ col_predicted,
-    const unsigned int * __restrict__ cell_seg_count,
-    TmConfig             cfg
-) {
-    const unsigned int cell = blockIdx.x;
-    if (cell >= cfg.n_cells) return;
-    const unsigned int col = cell / cfg.cells_per_column;
-    if (sp_active_mask[col] == 0) return;
-    if (col_predicted[col] == 0) return;
-
-    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
-    if (n_segs_here == 0) return;
-
-    const unsigned int tid = threadIdx.x;
-    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
-
-    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
-        const unsigned int seg = seg_base_id + local_seg;
-        if (seg_num_active_connected[seg] < cfg.activation_threshold) continue;
-        const unsigned int n_syn = seg_syn_count[seg];
-        if (n_syn == 0) continue;
-        const unsigned int syn_base = seg * cfg.synapses_per_segment;
-
-        for (unsigned int s = tid; s < n_syn; s += 32u) {
-            unsigned int presyn = syn_presyn[syn_base + s];
-            unsigned int word = prev_active_bits[presyn >> 5];
-            unsigned int bit = (word >> (presyn & 31u)) & 1u;
-            int p = (int)syn_perm[syn_base + s];
-            if (bit) {
-                int np = p + cfg.perm_inc_i16;
-                if (np > 32767) np = 32767;
-                syn_perm[syn_base + s] = (short)np;
-            } else {
-                int np = p - cfg.perm_dec_i16;
-                if (np < 0) np = 0;
-                syn_perm[syn_base + s] = (short)np;
-            }
-        }
-    }
-}
+// TM learn (reinforce correctly predicted segments) — cell-grouped launch.
+//
+// Grid: n_cells.
+// For each cell in a predicted, SP-active column: iterate its segments.
+// For each segment with num_active_connected >= activation_threshold,
+// reinforce its synapses against prev_active_bits.
+
+struct TmConfig {
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int cells_per_column;
+    unsigned int synapses_per_segment;
+    unsigned int n_segments;
+    unsigned int n_cells;
+    unsigned int max_segments_per_cell;
+    unsigned int max_new_synapses;
+    int conn_thr_i16;
+    int perm_inc_i16;
+    int perm_dec_i16;
+    int predicted_seg_dec_i16;
+    int initial_perm_i16;
+    unsigned int iter_seed;
+    unsigned int n_cols;
+    unsigned int bits_words;
+};
+
+extern "C" __global__
+void tm_learn_reinforce(
+    const unsigned int * __restrict__ seg_cell_id,
+    const unsigned int * __restrict__ seg_syn_count,
+    const unsigned int * __restrict__ syn_presyn,
+    short              * __restrict__ syn_perm,
+    const unsigned int * __restrict__ seg_num_active_connected,
+    const unsigned int * __restrict__ prev_active_bits,
+    const unsigned char * __restrict__ sp_active_mask,
+    const unsigned char * __restrict__ col_predicted,
+    const unsigned int * __restrict__ cell_seg_count,
+    TmConfig             cfg
+) {
+    const unsigned int cell = blockIdx.x;
+    if (cell >= cfg.n_cells) return;
+    const unsigned int col = cell / cfg.cells_per_column;
+    if (sp_active_mask[col] == 0) return;
+    if (col_predicted[col] == 0) return;
+
+    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
+    if (n_segs_here == 0) return;
+
+    const unsigned int tid = threadIdx.x;
+    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
+
+    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
+        const unsigned int seg = seg_base_id + local_seg;
+        if (seg_num_active_connected[seg] < cfg.activation_threshold) continue;
+        const unsigned int n_syn = seg_syn_count[seg];
+        if (n_syn == 0) continue;
+        const unsigned int syn_base = seg * cfg.synapses_per_segment;
+
+        for (unsigned int s = tid; s < n_syn; s += 32u) {
+            unsigned int presyn = syn_presyn[syn_base + s];
+            unsigned int word = prev_active_bits[presyn >> 5];
+            unsigned int bit = (word >> (presyn & 31u)) & 1u;
+            int p = (int)syn_perm[syn_base + s];
+            if (bit) {
+                int np = p + cfg.perm_inc_i16;
+                if (np > 32767) np = 32767;
+                syn_perm[syn_base + s] = (short)np;
+            } else {
+                int np = p - cfg.perm_dec_i16;
+                if (np < 0) np = 0;
+                syn_perm[syn_base + s] = (short)np;
+            }
+        }
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_predict.cu

@@ -1,102 +1,102 @@
-// TM predict kernel — cell-grouped launch.
-//
-// Grid: n_cells blocks (one per cell).
-// Block: 32 threads (one warp).
-//
-// Each block iterates the segments owned by its cell (count in cell_seg_count[cell]).
-// For each live segment, counts active connected/potential synapses against
-// prev_active_bits. Updates per-segment counters, cell_predictive bit, and
-// col_predicted flag.
-
-struct TmConfig {
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int cells_per_column;
-    unsigned int synapses_per_segment;
-    unsigned int n_segments;
-    unsigned int n_cells;
-    unsigned int max_segments_per_cell;
-    unsigned int max_new_synapses;
-    int conn_thr_i16;
-    int perm_inc_i16;
-    int perm_dec_i16;
-    int predicted_seg_dec_i16;
-    int initial_perm_i16;
-    unsigned int iter_seed;
-    unsigned int n_cols;
-    unsigned int bits_words;
-};
-
-extern "C" __global__
-void tm_predict(
-    const unsigned int * __restrict__ seg_cell_id,
-    const unsigned int * __restrict__ seg_syn_count,
-    const unsigned int * __restrict__ syn_presyn,
-    const short        * __restrict__ syn_perm,
-    const unsigned int * __restrict__ cell_active_bits,
-    unsigned int       * __restrict__ cell_predictive_bits,
-    unsigned char      * __restrict__ col_predicted,
-    unsigned int       * __restrict__ seg_num_active_connected,
-    unsigned int       * __restrict__ seg_num_active_potential,
-    unsigned int       * __restrict__ col_best_match,
-    const unsigned int * __restrict__ cell_seg_count,
-    TmConfig             cfg
-) {
-    const unsigned int cell = blockIdx.x;
-    if (cell >= cfg.n_cells) return;
-
-    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
-    if (n_segs_here == 0) return;
-
-    const unsigned int tid = threadIdx.x;
-    const unsigned int col = cell / cfg.cells_per_column;
-    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
-
-    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
-        const unsigned int seg = seg_base_id + local_seg;
-        const unsigned int n_syn = seg_syn_count[seg];
-        if (n_syn == 0) {
-            if (tid == 0) {
-                seg_num_active_connected[seg] = 0;
-                seg_num_active_potential[seg] = 0;
-            }
-            continue;
-        }
-        const unsigned int syn_base = seg * cfg.synapses_per_segment;
-
-        unsigned int local_conn = 0;
-        unsigned int local_pot = 0;
-        for (unsigned int s = tid; s < n_syn; s += 32u) {
-            unsigned int presyn = syn_presyn[syn_base + s];
-            unsigned int word = cell_active_bits[presyn >> 5];
-            unsigned int bit  = (word >> (presyn & 31u)) & 1u;
-            if (bit) {
-                local_pot += 1u;
-                int p = (int)syn_perm[syn_base + s];
-                if (p >= cfg.conn_thr_i16) {
-                    local_conn += 1u;
-                }
-            }
-        }
-        for (int off = 16; off > 0; off >>= 1) {
-            local_conn += __shfl_down_sync(0xffffffffu, local_conn, off);
-            local_pot  += __shfl_down_sync(0xffffffffu, local_pot,  off);
-        }
-
-        if (tid == 0) {
-            seg_num_active_connected[seg] = local_conn;
-            seg_num_active_potential[seg] = local_pot;
-            if (local_conn >= cfg.activation_threshold) {
-                unsigned int word_idx = cell >> 5;
-                unsigned int bit_mask = 1u << (cell & 31u);
-                atomicOr(&cell_predictive_bits[word_idx], bit_mask);
-                col_predicted[col] = 1;
-            }
-            if (local_pot >= cfg.learning_threshold) {
-                unsigned int pot_c = local_pot > 2047u ? 2047u : local_pot;
-                unsigned int key = (pot_c << 21) | (seg & 0x1FFFFFu);
-                atomicMax(&col_best_match[col], key);
-            }
-        }
-    }
-}
+// TM predict kernel — cell-grouped launch.
+//
+// Grid: n_cells blocks (one per cell).
+// Block: 32 threads (one warp).
+//
+// Each block iterates the segments owned by its cell (count in cell_seg_count[cell]).
+// For each live segment, counts active connected/potential synapses against
+// prev_active_bits. Updates per-segment counters, cell_predictive bit, and
+// col_predicted flag.
+
+struct TmConfig {
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int cells_per_column;
+    unsigned int synapses_per_segment;
+    unsigned int n_segments;
+    unsigned int n_cells;
+    unsigned int max_segments_per_cell;
+    unsigned int max_new_synapses;
+    int conn_thr_i16;
+    int perm_inc_i16;
+    int perm_dec_i16;
+    int predicted_seg_dec_i16;
+    int initial_perm_i16;
+    unsigned int iter_seed;
+    unsigned int n_cols;
+    unsigned int bits_words;
+};
+
+extern "C" __global__
+void tm_predict(
+    const unsigned int * __restrict__ seg_cell_id,
+    const unsigned int * __restrict__ seg_syn_count,
+    const unsigned int * __restrict__ syn_presyn,
+    const short        * __restrict__ syn_perm,
+    const unsigned int * __restrict__ cell_active_bits,
+    unsigned int       * __restrict__ cell_predictive_bits,
+    unsigned char      * __restrict__ col_predicted,
+    unsigned int       * __restrict__ seg_num_active_connected,
+    unsigned int       * __restrict__ seg_num_active_potential,
+    unsigned int       * __restrict__ col_best_match,
+    const unsigned int * __restrict__ cell_seg_count,
+    TmConfig             cfg
+) {
+    const unsigned int cell = blockIdx.x;
+    if (cell >= cfg.n_cells) return;
+
+    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
+    if (n_segs_here == 0) return;
+
+    const unsigned int tid = threadIdx.x;
+    const unsigned int col = cell / cfg.cells_per_column;
+    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
+
+    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
+        const unsigned int seg = seg_base_id + local_seg;
+        const unsigned int n_syn = seg_syn_count[seg];
+        if (n_syn == 0) {
+            if (tid == 0) {
+                seg_num_active_connected[seg] = 0;
+                seg_num_active_potential[seg] = 0;
+            }
+            continue;
+        }
+        const unsigned int syn_base = seg * cfg.synapses_per_segment;
+
+        unsigned int local_conn = 0;
+        unsigned int local_pot = 0;
+        for (unsigned int s = tid; s < n_syn; s += 32u) {
+            unsigned int presyn = syn_presyn[syn_base + s];
+            unsigned int word = cell_active_bits[presyn >> 5];
+            unsigned int bit  = (word >> (presyn & 31u)) & 1u;
+            if (bit) {
+                local_pot += 1u;
+                int p = (int)syn_perm[syn_base + s];
+                if (p >= cfg.conn_thr_i16) {
+                    local_conn += 1u;
+                }
+            }
+        }
+        for (int off = 16; off > 0; off >>= 1) {
+            local_conn += __shfl_down_sync(0xffffffffu, local_conn, off);
+            local_pot  += __shfl_down_sync(0xffffffffu, local_pot,  off);
+        }
+
+        if (tid == 0) {
+            seg_num_active_connected[seg] = local_conn;
+            seg_num_active_potential[seg] = local_pot;
+            if (local_conn >= cfg.activation_threshold) {
+                unsigned int word_idx = cell >> 5;
+                unsigned int bit_mask = 1u << (cell & 31u);
+                atomicOr(&cell_predictive_bits[word_idx], bit_mask);
+                col_predicted[col] = 1;
+            }
+            if (local_pot >= cfg.learning_threshold) {
+                unsigned int pot_c = local_pot > 2047u ? 2047u : local_pot;
+                unsigned int key = (pot_c << 21) | (seg & 0x1FFFFFu);
+                atomicMax(&col_best_match[col], key);
+            }
+        }
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_punish.cu

@@ -1,64 +1,64 @@
-// TM punish — cell-grouped launch.
-
-struct TmConfig {
-    unsigned int activation_threshold;
-    unsigned int learning_threshold;
-    unsigned int cells_per_column;
-    unsigned int synapses_per_segment;
-    unsigned int n_segments;
-    unsigned int n_cells;
-    unsigned int max_segments_per_cell;
-    unsigned int max_new_synapses;
-    int conn_thr_i16;
-    int perm_inc_i16;
-    int perm_dec_i16;
-    int predicted_seg_dec_i16;
-    int initial_perm_i16;
-    unsigned int iter_seed;
-    unsigned int n_cols;
-    unsigned int bits_words;
-};
-
-extern "C" __global__
-void tm_punish(
-    const unsigned int  * __restrict__ seg_cell_id,
-    const unsigned int  * __restrict__ seg_syn_count,
-    const unsigned int  * __restrict__ syn_presyn,
-    short               * __restrict__ syn_perm,
-    const unsigned int  * __restrict__ seg_num_active_potential,
-    const unsigned int  * __restrict__ prev_active_bits,
-    const unsigned char * __restrict__ sp_active_mask,
-    const unsigned int  * __restrict__ cell_seg_count,
-    TmConfig              cfg
-) {
-    const unsigned int cell = blockIdx.x;
-    if (cell >= cfg.n_cells) return;
-    const unsigned int col = cell / cfg.cells_per_column;
-    if (sp_active_mask[col] != 0) return;   // skip: col became active
-
-    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
-    if (n_segs_here == 0) return;
-
-    const unsigned int tid = threadIdx.x;
-    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
-
-    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
-        const unsigned int seg = seg_base_id + local_seg;
-        if (seg_num_active_potential[seg] < cfg.learning_threshold) continue;
-        const unsigned int n_syn = seg_syn_count[seg];
-        if (n_syn == 0) continue;
-        const unsigned int syn_base = seg * cfg.synapses_per_segment;
-
-        for (unsigned int s = tid; s < n_syn; s += 32u) {
-            unsigned int presyn = syn_presyn[syn_base + s];
-            unsigned int word = prev_active_bits[presyn >> 5];
-            unsigned int bit = (word >> (presyn & 31u)) & 1u;
-            if (bit) {
-                int p = (int)syn_perm[syn_base + s];
-                int np = p - cfg.predicted_seg_dec_i16;
-                if (np < 0) np = 0;
-                syn_perm[syn_base + s] = (short)np;
-            }
-        }
-    }
-}
+// TM punish — cell-grouped launch.
+
+struct TmConfig {
+    unsigned int activation_threshold;
+    unsigned int learning_threshold;
+    unsigned int cells_per_column;
+    unsigned int synapses_per_segment;
+    unsigned int n_segments;
+    unsigned int n_cells;
+    unsigned int max_segments_per_cell;
+    unsigned int max_new_synapses;
+    int conn_thr_i16;
+    int perm_inc_i16;
+    int perm_dec_i16;
+    int predicted_seg_dec_i16;
+    int initial_perm_i16;
+    unsigned int iter_seed;
+    unsigned int n_cols;
+    unsigned int bits_words;
+};
+
+extern "C" __global__
+void tm_punish(
+    const unsigned int  * __restrict__ seg_cell_id,
+    const unsigned int  * __restrict__ seg_syn_count,
+    const unsigned int  * __restrict__ syn_presyn,
+    short               * __restrict__ syn_perm,
+    const unsigned int  * __restrict__ seg_num_active_potential,
+    const unsigned int  * __restrict__ prev_active_bits,
+    const unsigned char * __restrict__ sp_active_mask,
+    const unsigned int  * __restrict__ cell_seg_count,
+    TmConfig              cfg
+) {
+    const unsigned int cell = blockIdx.x;
+    if (cell >= cfg.n_cells) return;
+    const unsigned int col = cell / cfg.cells_per_column;
+    if (sp_active_mask[col] != 0) return;   // skip: col became active
+
+    const unsigned int n_segs_here = min(cell_seg_count[cell], cfg.max_segments_per_cell);
+    if (n_segs_here == 0) return;
+
+    const unsigned int tid = threadIdx.x;
+    const unsigned int seg_base_id = cell * cfg.max_segments_per_cell;
+
+    for (unsigned int local_seg = 0; local_seg < n_segs_here; local_seg++) {
+        const unsigned int seg = seg_base_id + local_seg;
+        if (seg_num_active_potential[seg] < cfg.learning_threshold) continue;
+        const unsigned int n_syn = seg_syn_count[seg];
+        if (n_syn == 0) continue;
+        const unsigned int syn_base = seg * cfg.synapses_per_segment;
+
+        for (unsigned int s = tid; s < n_syn; s += 32u) {
+            unsigned int presyn = syn_presyn[syn_base + s];
+            unsigned int word = prev_active_bits[presyn >> 5];
+            unsigned int bit = (word >> (presyn & 31u)) & 1u;
+            if (bit) {
+                int p = (int)syn_perm[syn_base + s];
+                int np = p - cfg.predicted_seg_dec_i16;
+                if (np < 0) np = 0;
+                syn_perm[syn_base + s] = (short)np;
+            }
+        }
+    }
+}

overlay/htm_rust/src/gpu/kernels/tm_reset.cu

@@ -1,36 +1,36 @@
-// TM reset-per-step kernel.
-
-extern "C" __global__
-void tm_reset_step(
-    unsigned int * __restrict__ cell_active_bits,
-    unsigned int * __restrict__ cell_winner_bits,
-    unsigned int * __restrict__ cell_predictive_bits,
-    unsigned int * __restrict__ prev_active_bits,
-    unsigned int * __restrict__ prev_winner_bits,
-    unsigned char * __restrict__ col_predicted,
-    unsigned int * __restrict__ unpredicted_count,
-    unsigned int * __restrict__ burst_cols_count,
-    unsigned int * __restrict__ col_best_match,
-    unsigned int   bits_words,
-    unsigned int   n_cols
-) {
-    unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x;
-
-    if (tid_global < bits_words) {
-        prev_active_bits[tid_global] = cell_active_bits[tid_global];
-        prev_winner_bits[tid_global] = cell_winner_bits[tid_global];
-        cell_active_bits[tid_global] = 0u;
-        cell_winner_bits[tid_global] = 0u;
-        cell_predictive_bits[tid_global] = 0u;
-    }
-
-    if (tid_global < n_cols) {
-        col_predicted[tid_global] = 0;
-        col_best_match[tid_global] = 0u;
-    }
-
-    if (tid_global == 0) {
-        unpredicted_count[0] = 0u;
-        burst_cols_count[0] = 0u;
-    }
-}
+// TM reset-per-step kernel.
+
+extern "C" __global__
+void tm_reset_step(
+    unsigned int * __restrict__ cell_active_bits,
+    unsigned int * __restrict__ cell_winner_bits,
+    unsigned int * __restrict__ cell_predictive_bits,
+    unsigned int * __restrict__ prev_active_bits,
+    unsigned int * __restrict__ prev_winner_bits,
+    unsigned char * __restrict__ col_predicted,
+    unsigned int * __restrict__ unpredicted_count,
+    unsigned int * __restrict__ burst_cols_count,
+    unsigned int * __restrict__ col_best_match,
+    unsigned int   bits_words,
+    unsigned int   n_cols
+) {
+    unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid_global < bits_words) {
+        prev_active_bits[tid_global] = cell_active_bits[tid_global];
+        prev_winner_bits[tid_global] = cell_winner_bits[tid_global];
+        cell_active_bits[tid_global] = 0u;
+        cell_winner_bits[tid_global] = 0u;
+        cell_predictive_bits[tid_global] = 0u;
+    }
+
+    if (tid_global < n_cols) {
+        col_predicted[tid_global] = 0;
+        col_best_match[tid_global] = 0u;
+    }
+
+    if (tid_global == 0) {
+        unpredicted_count[0] = 0u;
+        burst_cols_count[0] = 0u;
+    }
+}

overlay/htm_rust/src/gpu/mod.rs

@@ -1,549 +1,549 @@
-//! GPU backend for HTM.
-//!
-//! Full-GPU pipeline (SP + TM). Per-step state lives entirely on device; the
-//! batch API (`step_many_gpu`) uploads T steps of input once, runs T iterations
-//! of the full HTM pipeline on GPU, and copies (T, n_cols) u8 + (T,) f32 back
-//! to the host in one shot.
-//!
-//! TM parity with the CPU reference is approximate:
-//!   - Segment growth: winner = cell 0 of bursting column (CPU picks
-//!     least-used-cell with RNG tiebreak). This is a pragmatic simplification
-//!     for GPU atomicity; learning dynamics are preserved.
-//!   - Permanences stored as i16 (scaled 0..32767). Rounding differs from
-//!     f32 by <= 1 ULP of the scale factor (≈ 3e-5) — inside any meaningful
-//!     HTM learning quantum.
-
-#![cfg(feature = "gpu")]
-
-pub mod sp_gpu;
-pub mod tm_gpu;
-pub mod fused;
-
-#[cfg(test)]
-mod tests;
-
-use std::mem::ManuallyDrop;
-
-use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyTuple};
-use numpy::{PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray2, PyUntypedArrayMethods};
-
-use crate::region::HTMRegionCore;
-use crate::sp::SpatialPoolerConfig;
-use sp_gpu::SpatialPoolerGpu;
-use tm_gpu::TemporalMemoryGpu;
-use fused::FusedState;
-
-/// Extract (device_ptr, shape, typestr) from a `__cuda_array_interface__` dict.
-/// Returns Err if the dict is malformed. Used by `step_many_cuda` to wrap
-/// torch-owned CUDA allocations zero-copy.
-fn cai_parse(cai: &Bound<'_, PyDict>) -> PyResult<(u64, Vec<usize>, String)> {
-    // `data` is a (ptr: int, readonly: bool) tuple.
-    let data_obj = cai.get_item("data")?
-        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'data'"))?;
-    let data_tup: Bound<'_, PyTuple> = data_obj.downcast_into()
-        .map_err(|_| pyo3::exceptions::PyValueError::new_err("CAI 'data' must be a tuple"))?;
-    let ptr: u64 = data_tup.get_item(0)?.extract()?;
-
-    // `shape` is a tuple of ints.
-    let shape_obj = cai.get_item("shape")?
-        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'shape'"))?;
-    let shape_tup: Bound<'_, PyTuple> = shape_obj.downcast_into()
-        .map_err(|_| pyo3::exceptions::PyValueError::new_err("CAI 'shape' must be a tuple"))?;
-    let shape: Vec<usize> = (0..shape_tup.len())
-        .map(|i| shape_tup.get_item(i).and_then(|v| v.extract::<usize>()))
-        .collect::<PyResult<Vec<_>>>()?;
-
-    // `typestr` (e.g. "|u1", "<f4").
-    let typestr_obj = cai.get_item("typestr")?
-        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'typestr'"))?;
-    let typestr: String = typestr_obj.extract()?;
-
-    // Reject non-contiguous tensors — we don't handle strides.
-    if let Some(strides) = cai.get_item("strides")? {
-        if !strides.is_none() {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "CAI 'strides' must be None (tensor must be contiguous)",
-            ));
-        }
-    }
-
-    Ok((ptr, shape, typestr))
-}
-
-/// Python-exposed GPU HTM region. Drop-in replacement for `HTMRegion`.
-#[pyclass(module = "htm_rust")]
-pub struct HTMRegionGpu {
-    pub(super) sp_gpu: SpatialPoolerGpu,
-    pub(super) tm_gpu: TemporalMemoryGpu,
-    pub(super) fused_state: FusedState,
-    pub(super) n_columns: usize,
-    pub(super) input_bits: usize,
-    pub(super) cells_per_column: usize,
-}
-
-#[pymethods]
-impl HTMRegionGpu {
-    #[new]
-    #[pyo3(signature = (input_bits, n_columns, cells_per_column, seed=42))]
-    fn new(
-        input_bits: usize,
-        n_columns: usize,
-        cells_per_column: usize,
-        seed: u64,
-    ) -> PyResult<Self> {
-        if input_bits == 0 || n_columns == 0 || cells_per_column == 0 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "input_bits, n_columns, cells_per_column must all be > 0",
-            ));
-        }
-        // CPU reference for deterministic SP init.
-        let cpu_ref = HTMRegionCore::new(input_bits, n_columns, cells_per_column, seed);
-        let sp_cfg: &SpatialPoolerConfig = &cpu_ref.sp.cfg;
-        let sp_gpu = SpatialPoolerGpu::from_cpu(&cpu_ref.sp).map_err(|e| {
-            pyo3::exceptions::PyRuntimeError::new_err(format!(
-                "GPU SP init failed: {e:?}. Config: input_bits={}, n_columns={}",
-                sp_cfg.input_bits, sp_cfg.n_columns,
-            ))
-        })?;
-        let dev = sp_gpu.dev_ref().clone();
-        let tm_gpu = TemporalMemoryGpu::new(dev.clone(), n_columns, cells_per_column).map_err(|e| {
-            pyo3::exceptions::PyRuntimeError::new_err(format!(
-                "GPU TM init failed: {e:?}",
-            ))
-        })?;
-        let initial_threshold = sp_gpu.initial_threshold_estimate();
-        let fused_state = FusedState::new(dev, n_columns, cells_per_column, initial_threshold)
-            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!(
-                "GPU fused state init failed: {e:?}",
-            )))?;
-        Ok(Self {
-            sp_gpu,
-            tm_gpu,
-            fused_state,
-            n_columns,
-            input_bits,
-            cells_per_column,
-        })
-    }
-
-    #[getter] fn input_bits(&self) -> usize { self.input_bits }
-    #[getter] fn n_columns(&self) -> usize { self.n_columns }
-    #[getter] fn cells_per_column(&self) -> usize { self.cells_per_column }
-
-    /// Process T timesteps in one call on GPU. Per-step state (SP + TM) stays
-    /// on device; only the final (T, n_cols) mask and (T,) anomaly are copied
-    /// to the host at the end.
-    #[pyo3(signature = (inputs, learn=true))]
-    fn step_many_gpu<'py>(
-        &mut self,
-        py: Python<'py>,
-        inputs: PyReadonlyArray2<'py, bool>,
-        learn: bool,
-    ) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray1<f32>>)> {
-        let shape = inputs.shape();
-        if shape.len() != 2 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "inputs must be 2-D (T, input_bits)",
-            ));
-        }
-        let t = shape[0];
-        let bits = shape[1];
-        if bits != self.input_bits {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "inputs last dim {bits} != expected input_bits {}",
-                self.input_bits,
-            )));
-        }
-        let slice = inputs.as_slice()?;
-        let n_cols = self.n_columns;
-        let input_vec: Vec<bool> = slice.to_vec();
-
-        let result = py.allow_threads(|| -> Result<(Vec<u8>, Vec<f32>), String> {
-            // 1. Upload T*input_bits bytes (32 MB at T=2048, bits=16384).
-            let sdr_u8_all: Vec<u8> = input_vec.iter().map(|&b| b as u8).collect();
-            let inputs_dev = self
-                .sp_gpu
-                .dev_ref()
-                .htod_sync_copy(&sdr_u8_all)
-                .map_err(|e| format!("H2D inputs: {e:?}"))?;
-
-            // 2. Allocate output buffers on device.
-            let mut cols_dev = self.sp_gpu.dev_ref()
-                .alloc_zeros::<u8>(t * n_cols)
-                .map_err(|e| format!("alloc cols: {e:?}"))?;
-            let mut anom_dev = self.sp_gpu.dev_ref()
-                .alloc_zeros::<f32>(t)
-                .map_err(|e| format!("alloc anom: {e:?}"))?;
-
-            // 3. Run T steps of SP + TM on GPU with NO per-step host sync.
-            self.sp_gpu.step_batch_with_tm(
-                &inputs_dev,
-                t,
-                self.input_bits,
-                learn,
-                &mut cols_dev,
-                &mut anom_dev,
-                &mut self.tm_gpu,
-            ).map_err(|e| format!("step_batch_with_tm: {e:?}"))?;
-
-            // 4. ONE D2H for the whole run (T * n_cols bytes + T floats).
-            let cols_host: Vec<u8> = self.sp_gpu.dev_ref()
-                .dtoh_sync_copy(&cols_dev)
-                .map_err(|e| format!("D2H cols: {e:?}"))?;
-            let anom_host: Vec<f32> = self.sp_gpu.dev_ref()
-                .dtoh_sync_copy(&anom_dev)
-                .map_err(|e| format!("D2H anom: {e:?}"))?;
-
-            Ok((cols_host, anom_host))
-        });
-
-        let (cols_u8, anom) = result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
-
-        let cols_f32: Vec<f32> = cols_u8.iter().map(|&b| b as f32).collect();
-        let cols_arr = numpy::PyArray1::from_vec_bound(py, cols_f32)
-            .reshape([t, n_cols])
-            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
-        let anom_arr = numpy::PyArray1::from_vec_bound(py, anom);
-        Ok((cols_arr, anom_arr))
-    }
-
-    /// Zero-copy CUDA path: accept torch tensors via __cuda_array_interface__,
-    /// write outputs directly into caller-allocated torch tensors. Skips the
-    /// host round-trip that `step_many_gpu` pays on every call (sdr.cpu() +
-    /// two D2H copies at the end). This is the hot path for `train.py`.
-    ///
-    /// Contract:
-    ///   sdr_cai.shape  == (T, input_bits), dtype u8   (0/1 mask)
-    ///   cols_cai.shape == (T, n_columns),  dtype u8   (written)
-    ///   anom_cai.shape == (T,),            dtype f32  (written)
-    /// All three tensors must live on the SAME CUDA device as this region.
-    ///
-    /// The torch tensors still own their memory — this method only wraps
-    /// them as borrowed CudaSlice views (via ManuallyDrop) so cudarc's Drop
-    /// impl can't free pytorch's allocator.
-    #[pyo3(signature = (sdr_cai, cols_cai, anom_cai, learn=true))]
-    fn step_many_cuda(
-        &mut self,
-        py: Python<'_>,
-        sdr_cai: &Bound<'_, PyDict>,
-        cols_cai: &Bound<'_, PyDict>,
-        anom_cai: &Bound<'_, PyDict>,
-        learn: bool,
-    ) -> PyResult<()> {
-        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(sdr_cai)?;
-        let (cols_ptr, cols_shape, cols_type) = cai_parse(cols_cai)?;
-        let (anom_ptr, anom_shape, anom_type) = cai_parse(anom_cai)?;
-
-        // typestr sanity. numpy u1 is what torch.uint8 exports.
-        if sdr_type != "|u1" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "sdr_cai typestr must be '|u1' (uint8), got {sdr_type}",
-            )));
-        }
-        if cols_type != "|u1" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "cols_cai typestr must be '|u1' (uint8), got {cols_type}",
-            )));
-        }
-        if anom_type != "<f4" && anom_type != "=f4" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "anom_cai typestr must be '<f4' (float32), got {anom_type}",
-            )));
-        }
-
-        // Shape validation.
-        if sdr_shape.len() != 2 || sdr_shape[1] != self.input_bits {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "sdr_cai shape {sdr_shape:?} != (T, {})",
-                self.input_bits,
-            )));
-        }
-        let t = sdr_shape[0];
-        if cols_shape != [t, self.n_columns] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "cols_cai shape {cols_shape:?} != ({t}, {})",
-                self.n_columns,
-            )));
-        }
-        if anom_shape != [t] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "anom_cai shape {anom_shape:?} != ({t},)",
-            )));
-        }
-
-        let dev = self.sp_gpu.dev_ref().clone();
-        let n_cols = self.n_columns;
-        let input_bits = self.input_bits;
-
-        let result = py.allow_threads(|| -> Result<(), String> {
-            // SAFETY:
-            // - ptrs came from torch CUDA tensors validated non-null by the
-            //   __cuda_array_interface__ contract.
-            // - lens computed from validated shapes.
-            // - We wrap the returned CudaSlice in ManuallyDrop so cudarc's
-            //   Drop (which calls cuMemFree) never runs against torch memory.
-            //   The underlying allocation is owned+freed by torch.
-            // - The slices are used only for the duration of this call;
-            //   torch guarantees the backing tensors are live across it
-            //   (Python holds refs on the wrapping tensors).
-            let inputs_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<u8>(sdr_ptr, t * input_bits)
-            });
-            let mut cols_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols)
-            });
-            let mut anom_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<f32>(anom_ptr, t)
-            });
-
-            self.sp_gpu.step_batch_with_tm(
-                &inputs_dev,
-                t,
-                input_bits,
-                learn,
-                &mut cols_dev,
-                &mut anom_dev,
-                &mut self.tm_gpu,
-            ).map_err(|e| format!("step_batch_with_tm: {e:?}"))?;
-
-            // Synchronize: kernel writes must be visible to the next torch
-            // op that reads cols/anom. Pytorch's default stream is stream 0,
-            // and cudarc launches on its own stream — a full device sync
-            // is the simplest correct barrier. (Could narrow to a stream
-            // wait event in PR 2.)
-            // No dev.synchronize() here: caller must explicitly sync via the
-            // `device_sync()` method (or PyTorch auto-syncs when the output
-            // tensor is next consumed). Removing the per-launch barrier lets
-            // subsequent GPU work (mamba3 fwd, etc.) overlap in time.
-            Ok(())
-        });
-
-        result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
-        Ok(())
-    }
-
-    /// Clear TM state on the GPU.
-    fn reset(&mut self) -> PyResult<()> {
-        self.tm_gpu.reset().map_err(|e| {
-            pyo3::exceptions::PyRuntimeError::new_err(format!("GPU TM reset: {e:?}"))
-        })?;
-        self.fused_state.reset().map_err(|e| {
-            pyo3::exceptions::PyRuntimeError::new_err(format!("GPU fused reset: {e:?}"))
-        })
-    }
-
-    /// FUSED MEGAKERNEL PATH: single CUDA launch for the entire T-step
-    /// forward (SP + TM all in one). Accepts torch CUDA tensors via
-    /// `__cuda_array_interface__` (zero-copy). Writes active-column mask +
-    /// anomaly directly into caller-allocated torch tensors.
-    ///
-    /// Semantics diverge from `step_many_cuda` in one important way: column
-    /// activation uses per-column threshold inhibition instead of global
-    /// top-K. The threshold is EMA-adapted per column toward the sparsity
-    /// target. See `docs/GPU_HTM.md` §Fused Kernel.
-    #[pyo3(signature = (sdr_cai, cols_cai, anom_cai, learn=true))]
-    fn step_many_fused_cuda(
-        &mut self,
-        py: Python<'_>,
-        sdr_cai: &Bound<'_, PyDict>,
-        cols_cai: &Bound<'_, PyDict>,
-        anom_cai: &Bound<'_, PyDict>,
-        learn: bool,
-    ) -> PyResult<()> {
-        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(sdr_cai)?;
-        let (cols_ptr, cols_shape, cols_type) = cai_parse(cols_cai)?;
-        let (anom_ptr, anom_shape, anom_type) = cai_parse(anom_cai)?;
-
-        if sdr_type != "|u1" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "sdr_cai typestr must be '|u1' (uint8), got {sdr_type}",
-            )));
-        }
-        if cols_type != "|u1" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "cols_cai typestr must be '|u1' (uint8), got {cols_type}",
-            )));
-        }
-        if anom_type != "<f4" && anom_type != "=f4" {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "anom_cai typestr must be '<f4' (float32), got {anom_type}",
-            )));
-        }
-
-        if sdr_shape.len() != 2 || sdr_shape[1] != self.input_bits {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "sdr_cai shape {sdr_shape:?} != (T, {})",
-                self.input_bits,
-            )));
-        }
-        let t = sdr_shape[0];
-        if cols_shape != [t, self.n_columns] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "cols_cai shape {cols_shape:?} != ({t}, {})",
-                self.n_columns,
-            )));
-        }
-        if anom_shape != [t] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "anom_cai shape {anom_shape:?} != ({t},)",
-            )));
-        }
-
-        let dev = self.sp_gpu.dev_ref().clone();
-        let n_cols = self.n_columns;
-        let input_bits = self.input_bits;
-
-        let result = py.allow_threads(|| -> Result<(), String> {
-            let inputs_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<u8>(sdr_ptr, t * input_bits)
-            });
-            let mut cols_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols)
-            });
-            let mut anom_dev = ManuallyDrop::new(unsafe {
-                dev.upgrade_device_ptr::<f32>(anom_ptr, t)
-            });
-
-            fused::launch_fused(
-                &mut self.sp_gpu,
-                &mut self.tm_gpu,
-                &mut self.fused_state,
-                &inputs_dev,
-                &mut cols_dev,
-                &mut anom_dev,
-                t,
-                input_bits,
-                learn,
-            ).map_err(|e| format!("launch_fused: {e:?}"))?;
-
-            // No dev.synchronize() here: caller must explicitly sync via the
-            // `device_sync()` method (or PyTorch auto-syncs when the output
-            // tensor is next consumed). Removing the per-launch barrier lets
-            // subsequent GPU work (mamba3 fwd, etc.) overlap in time.
-            Ok(())
-        });
-
-        result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
-        Ok(())
-    }
-
-    /// Explicit device synchronization — the caller must invoke this after
-    /// all batched `step_many_*_cuda` calls complete, before reading the
-    /// output tensors from a different CUDA stream. Equivalent to the old
-    /// per-call `dev.synchronize()` that was removed for overlap.
-    fn device_sync(&self) -> PyResult<()> {
-        let dev = self.sp_gpu.dev_ref();
-        dev.synchronize()
-            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("sync: {e:?}")))?;
-        Ok(())
-    }
-}
-
-/// Batch B regions into ONE cooperative kernel launch. Breaks through the
-/// CUDA cooperative-kernel device-level serialization: a single cooperative
-/// launch with grid.y=B processes all regions concurrently — ~B× speedup
-/// over B sequential launches.
-///
-/// All regions must have the same config (input_bits, n_columns,
-/// cells_per_column). Each region keeps its independent GPU state.
-/// Does NOT sync; caller must invoke `device_sync()` on any region
-/// afterwards (or rely on a downstream torch op to auto-sync).
-#[pyfunction]
-#[pyo3(signature = (regions, sdr_cais, cols_cais, anom_cais, learn=true))]
-fn step_batch_fused_cuda(
-    py: Python<'_>,
-    regions: Vec<Py<HTMRegionGpu>>,
-    sdr_cais: Vec<Bound<'_, PyDict>>,
-    cols_cais: Vec<Bound<'_, PyDict>>,
-    anom_cais: Vec<Bound<'_, PyDict>>,
-    learn: bool,
-) -> PyResult<()> {
-    let b = regions.len();
-    if b == 0 {
-        return Err(pyo3::exceptions::PyValueError::new_err("regions is empty"));
-    }
-    if sdr_cais.len() != b || cols_cais.len() != b || anom_cais.len() != b {
-        return Err(pyo3::exceptions::PyValueError::new_err(
-            "sdr_cais / cols_cais / anom_cais length must match regions",
-        ));
-    }
-
-    // Parse all CAI dicts; collect device pointers. Validate shapes/dtypes.
-    let mut sdr_ptrs = Vec::with_capacity(b);
-    let mut cols_ptrs = Vec::with_capacity(b);
-    let mut anom_ptrs = Vec::with_capacity(b);
-    let (input_bits, n_columns, t) = {
-        let r0 = regions[0].bind(py).borrow();
-        (r0.input_bits, r0.n_columns, {
-            let (_p, sh, _ty) = cai_parse(&sdr_cais[0])?;
-            if sh.len() != 2 {
-                return Err(pyo3::exceptions::PyValueError::new_err(
-                    format!("sdr_cai must be 2-D (T, input_bits), got {sh:?}"),
-                ));
-            }
-            sh[0]
-        })
-    };
-
-    for i in 0..b {
-        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(&sdr_cais[i])?;
-        let (cols_ptr, cols_shape, cols_type) = cai_parse(&cols_cais[i])?;
-        let (anom_ptr, anom_shape, anom_type) = cai_parse(&anom_cais[i])?;
-        if sdr_type != "|u1" || cols_type != "|u1" {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "sdr/cols typestr must be '|u1' (uint8)",
-            ));
-        }
-        if anom_type != "<f4" && anom_type != "=f4" {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "anom typestr must be '<f4' (float32)",
-            ));
-        }
-        if sdr_shape != [t, input_bits] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "sdr[{i}] shape {sdr_shape:?} != ({t}, {input_bits})"
-            )));
-        }
-        if cols_shape != [t, n_columns] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "cols[{i}] shape {cols_shape:?} != ({t}, {n_columns})"
-            )));
-        }
-        if anom_shape != [t] {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "anom[{i}] shape {anom_shape:?} != ({t},)"
-            )));
-        }
-        sdr_ptrs.push(sdr_ptr);
-        cols_ptrs.push(cols_ptr);
-        anom_ptrs.push(anom_ptr);
-    }
-
-    // Exclusively borrow each region. PyRefMut guarantees uniqueness.
-    let mut region_refs: Vec<pyo3::PyRefMut<HTMRegionGpu>> =
-        regions.iter().map(|p| p.bind(py).borrow_mut()).collect();
-    // Collect raw mutable pointers — each PyRefMut exclusively borrows its
-    // region for the lifetime of this call, so pointers stay valid and
-    // unique. launch_fused_batched_raw only dereferences one region at a
-    // time, not constructing an aliased slice.
-    let raw_ptrs: Vec<*mut HTMRegionGpu> = region_refs
-        .iter_mut()
-        .map(|r| &mut **r as *mut HTMRegionGpu)
-        .collect();
-
-    // No allow_threads: raw pointers aren't Send. The launch is GPU-queued
-    // and sync'd downstream; holding the GIL for the duration is cheap.
-    fused::launch_fused_batched_raw(
-        &raw_ptrs, &sdr_ptrs, &cols_ptrs, &anom_ptrs,
-        t, input_bits, learn,
-    )
-    .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("launch_fused_batched: {e:?}")))?;
-    Ok(())
-}
-
-pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<HTMRegionGpu>()?;
-    m.add_function(pyo3::wrap_pyfunction!(step_batch_fused_cuda, m)?)?;
-    Ok(())
-}
+//! GPU backend for HTM.
+//!
+//! Full-GPU pipeline (SP + TM). Per-step state lives entirely on device; the
+//! batch API (`step_many_gpu`) uploads T steps of input once, runs T iterations
+//! of the full HTM pipeline on GPU, and copies (T, n_cols) u8 + (T,) f32 back
+//! to the host in one shot.
+//!
+//! TM parity with the CPU reference is approximate:
+//!   - Segment growth: winner = cell 0 of bursting column (CPU picks
+//!     least-used-cell with RNG tiebreak). This is a pragmatic simplification
+//!     for GPU atomicity; learning dynamics are preserved.
+//!   - Permanences stored as i16 (scaled 0..32767). Rounding differs from
+//!     f32 by <= 1 ULP of the scale factor (≈ 3e-5) — inside any meaningful
+//!     HTM learning quantum.
+
+#![cfg(feature = "gpu")]
+
+pub mod sp_gpu;
+pub mod tm_gpu;
+pub mod fused;
+
+#[cfg(test)]
+mod tests;
+
+use std::mem::ManuallyDrop;
+
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+use numpy::{PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray2, PyUntypedArrayMethods};
+
+use crate::region::HTMRegionCore;
+use crate::sp::SpatialPoolerConfig;
+use sp_gpu::SpatialPoolerGpu;
+use tm_gpu::TemporalMemoryGpu;
+use fused::FusedState;
+
+/// Extract (device_ptr, shape, typestr) from a `__cuda_array_interface__` dict.
+/// Returns Err if the dict is malformed. Used by `step_many_cuda` to wrap
+/// torch-owned CUDA allocations zero-copy.
+fn cai_parse(cai: &Bound<'_, PyDict>) -> PyResult<(u64, Vec<usize>, String)> {
+    // `data` is a (ptr: int, readonly: bool) tuple.
+    let data_obj = cai.get_item("data")?
+        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'data'"))?;
+    let data_tup: Bound<'_, PyTuple> = data_obj.downcast_into()
+        .map_err(|_| pyo3::exceptions::PyValueError::new_err("CAI 'data' must be a tuple"))?;
+    let ptr: u64 = data_tup.get_item(0)?.extract()?;
+
+    // `shape` is a tuple of ints.
+    let shape_obj = cai.get_item("shape")?
+        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'shape'"))?;
+    let shape_tup: Bound<'_, PyTuple> = shape_obj.downcast_into()
+        .map_err(|_| pyo3::exceptions::PyValueError::new_err("CAI 'shape' must be a tuple"))?;
+    let shape: Vec<usize> = (0..shape_tup.len())
+        .map(|i| shape_tup.get_item(i).and_then(|v| v.extract::<usize>()))
+        .collect::<PyResult<Vec<_>>>()?;
+
+    // `typestr` (e.g. "|u1", "<f4").
+    let typestr_obj = cai.get_item("typestr")?
+        .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("CAI missing 'typestr'"))?;
+    let typestr: String = typestr_obj.extract()?;
+
+    // Reject non-contiguous tensors — we don't handle strides.
+    if let Some(strides) = cai.get_item("strides")? {
+        if !strides.is_none() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "CAI 'strides' must be None (tensor must be contiguous)",
+            ));
+        }
+    }
+
+    Ok((ptr, shape, typestr))
+}
+
+/// Python-exposed GPU HTM region. Drop-in replacement for `HTMRegion`.
+#[pyclass(module = "htm_rust")]
+pub struct HTMRegionGpu {
+    pub(super) sp_gpu: SpatialPoolerGpu,
+    pub(super) tm_gpu: TemporalMemoryGpu,
+    pub(super) fused_state: FusedState,
+    pub(super) n_columns: usize,
+    pub(super) input_bits: usize,
+    pub(super) cells_per_column: usize,
+}
+
+#[pymethods]
+impl HTMRegionGpu {
+    #[new]
+    #[pyo3(signature = (input_bits, n_columns, cells_per_column, seed=42))]
+    fn new(
+        input_bits: usize,
+        n_columns: usize,
+        cells_per_column: usize,
+        seed: u64,
+    ) -> PyResult<Self> {
+        if input_bits == 0 || n_columns == 0 || cells_per_column == 0 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "input_bits, n_columns, cells_per_column must all be > 0",
+            ));
+        }
+        // CPU reference for deterministic SP init.
+        let cpu_ref = HTMRegionCore::new(input_bits, n_columns, cells_per_column, seed);
+        let sp_cfg: &SpatialPoolerConfig = &cpu_ref.sp.cfg;
+        let sp_gpu = SpatialPoolerGpu::from_cpu(&cpu_ref.sp).map_err(|e| {
+            pyo3::exceptions::PyRuntimeError::new_err(format!(
+                "GPU SP init failed: {e:?}. Config: input_bits={}, n_columns={}",
+                sp_cfg.input_bits, sp_cfg.n_columns,
+            ))
+        })?;
+        let dev = sp_gpu.dev_ref().clone();
+        let tm_gpu = TemporalMemoryGpu::new(dev.clone(), n_columns, cells_per_column).map_err(|e| {
+            pyo3::exceptions::PyRuntimeError::new_err(format!(
+                "GPU TM init failed: {e:?}",
+            ))
+        })?;
+        let initial_threshold = sp_gpu.initial_threshold_estimate();
+        let fused_state = FusedState::new(dev, n_columns, cells_per_column, initial_threshold)
+            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!(
+                "GPU fused state init failed: {e:?}",
+            )))?;
+        Ok(Self {
+            sp_gpu,
+            tm_gpu,
+            fused_state,
+            n_columns,
+            input_bits,
+            cells_per_column,
+        })
+    }
+
+    #[getter] fn input_bits(&self) -> usize { self.input_bits }
+    #[getter] fn n_columns(&self) -> usize { self.n_columns }
+    #[getter] fn cells_per_column(&self) -> usize { self.cells_per_column }
+
+    /// Process T timesteps in one call on GPU. Per-step state (SP + TM) stays
+    /// on device; only the final (T, n_cols) mask and (T,) anomaly are copied
+    /// to the host at the end.
+    #[pyo3(signature = (inputs, learn=true))]
+    fn step_many_gpu<'py>(
+        &mut self,
+        py: Python<'py>,
+        inputs: PyReadonlyArray2<'py, bool>,
+        learn: bool,
+    ) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray1<f32>>)> {
+        let shape = inputs.shape();
+        if shape.len() != 2 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "inputs must be 2-D (T, input_bits)",
+            ));
+        }
+        let t = shape[0];
+        let bits = shape[1];
+        if bits != self.input_bits {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "inputs last dim {bits} != expected input_bits {}",
+                self.input_bits,
+            )));
+        }
+        let slice = inputs.as_slice()?;
+        let n_cols = self.n_columns;
+        let input_vec: Vec<bool> = slice.to_vec();
+
+        let result = py.allow_threads(|| -> Result<(Vec<u8>, Vec<f32>), String> {
+            // 1. Upload T*input_bits bytes (32 MB at T=2048, bits=16384).
+            let sdr_u8_all: Vec<u8> = input_vec.iter().map(|&b| b as u8).collect();
+            let inputs_dev = self
+                .sp_gpu
+                .dev_ref()
+                .htod_sync_copy(&sdr_u8_all)
+                .map_err(|e| format!("H2D inputs: {e:?}"))?;
+
+            // 2. Allocate output buffers on device.
+            let mut cols_dev = self.sp_gpu.dev_ref()
+                .alloc_zeros::<u8>(t * n_cols)
+                .map_err(|e| format!("alloc cols: {e:?}"))?;
+            let mut anom_dev = self.sp_gpu.dev_ref()
+                .alloc_zeros::<f32>(t)
+                .map_err(|e| format!("alloc anom: {e:?}"))?;
+
+            // 3. Run T steps of SP + TM on GPU with NO per-step host sync.
+            self.sp_gpu.step_batch_with_tm(
+                &inputs_dev,
+                t,
+                self.input_bits,
+                learn,
+                &mut cols_dev,
+                &mut anom_dev,
+                &mut self.tm_gpu,
+            ).map_err(|e| format!("step_batch_with_tm: {e:?}"))?;
+
+            // 4. ONE D2H for the whole run (T * n_cols bytes + T floats).
+            let cols_host: Vec<u8> = self.sp_gpu.dev_ref()
+                .dtoh_sync_copy(&cols_dev)
+                .map_err(|e| format!("D2H cols: {e:?}"))?;
+            let anom_host: Vec<f32> = self.sp_gpu.dev_ref()
+                .dtoh_sync_copy(&anom_dev)
+                .map_err(|e| format!("D2H anom: {e:?}"))?;
+
+            Ok((cols_host, anom_host))
+        });
+
+        let (cols_u8, anom) = result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
+
+        let cols_f32: Vec<f32> = cols_u8.iter().map(|&b| b as f32).collect();
+        let cols_arr = numpy::PyArray1::from_vec_bound(py, cols_f32)
+            .reshape([t, n_cols])
+            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
+        let anom_arr = numpy::PyArray1::from_vec_bound(py, anom);
+        Ok((cols_arr, anom_arr))
+    }
+
+    /// Zero-copy CUDA path: accept torch tensors via __cuda_array_interface__,
+    /// write outputs directly into caller-allocated torch tensors. Skips the
+    /// host round-trip that `step_many_gpu` pays on every call (sdr.cpu() +
+    /// two D2H copies at the end). This is the hot path for `train.py`.
+    ///
+    /// Contract:
+    ///   sdr_cai.shape  == (T, input_bits), dtype u8   (0/1 mask)
+    ///   cols_cai.shape == (T, n_columns),  dtype u8   (written)
+    ///   anom_cai.shape == (T,),            dtype f32  (written)
+    /// All three tensors must live on the SAME CUDA device as this region.
+    ///
+    /// The torch tensors still own their memory — this method only wraps
+    /// them as borrowed CudaSlice views (via ManuallyDrop) so cudarc's Drop
+    /// impl can't free pytorch's allocator.
+    #[pyo3(signature = (sdr_cai, cols_cai, anom_cai, learn=true))]
+    fn step_many_cuda(
+        &mut self,
+        py: Python<'_>,
+        sdr_cai: &Bound<'_, PyDict>,
+        cols_cai: &Bound<'_, PyDict>,
+        anom_cai: &Bound<'_, PyDict>,
+        learn: bool,
+    ) -> PyResult<()> {
+        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(sdr_cai)?;
+        let (cols_ptr, cols_shape, cols_type) = cai_parse(cols_cai)?;
+        let (anom_ptr, anom_shape, anom_type) = cai_parse(anom_cai)?;
+
+        // typestr sanity. numpy u1 is what torch.uint8 exports.
+        if sdr_type != "|u1" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "sdr_cai typestr must be '|u1' (uint8), got {sdr_type}",
+            )));
+        }
+        if cols_type != "|u1" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "cols_cai typestr must be '|u1' (uint8), got {cols_type}",
+            )));
+        }
+        if anom_type != "<f4" && anom_type != "=f4" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "anom_cai typestr must be '<f4' (float32), got {anom_type}",
+            )));
+        }
+
+        // Shape validation.
+        if sdr_shape.len() != 2 || sdr_shape[1] != self.input_bits {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "sdr_cai shape {sdr_shape:?} != (T, {})",
+                self.input_bits,
+            )));
+        }
+        let t = sdr_shape[0];
+        if cols_shape != [t, self.n_columns] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "cols_cai shape {cols_shape:?} != ({t}, {})",
+                self.n_columns,
+            )));
+        }
+        if anom_shape != [t] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "anom_cai shape {anom_shape:?} != ({t},)",
+            )));
+        }
+
+        let dev = self.sp_gpu.dev_ref().clone();
+        let n_cols = self.n_columns;
+        let input_bits = self.input_bits;
+
+        let result = py.allow_threads(|| -> Result<(), String> {
+            // SAFETY:
+            // - ptrs came from torch CUDA tensors validated non-null by the
+            //   __cuda_array_interface__ contract.
+            // - lens computed from validated shapes.
+            // - We wrap the returned CudaSlice in ManuallyDrop so cudarc's
+            //   Drop (which calls cuMemFree) never runs against torch memory.
+            //   The underlying allocation is owned+freed by torch.
+            // - The slices are used only for the duration of this call;
+            //   torch guarantees the backing tensors are live across it
+            //   (Python holds refs on the wrapping tensors).
+            let inputs_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<u8>(sdr_ptr, t * input_bits)
+            });
+            let mut cols_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols)
+            });
+            let mut anom_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<f32>(anom_ptr, t)
+            });
+
+            self.sp_gpu.step_batch_with_tm(
+                &inputs_dev,
+                t,
+                input_bits,
+                learn,
+                &mut cols_dev,
+                &mut anom_dev,
+                &mut self.tm_gpu,
+            ).map_err(|e| format!("step_batch_with_tm: {e:?}"))?;
+
+            // Synchronize: kernel writes must be visible to the next torch
+            // op that reads cols/anom. Pytorch's default stream is stream 0,
+            // and cudarc launches on its own stream — a full device sync
+            // is the simplest correct barrier. (Could narrow to a stream
+            // wait event in PR 2.)
+            // No dev.synchronize() here: caller must explicitly sync via the
+            // `device_sync()` method (or PyTorch auto-syncs when the output
+            // tensor is next consumed). Removing the per-launch barrier lets
+            // subsequent GPU work (mamba3 fwd, etc.) overlap in time.
+            Ok(())
+        });
+
+        result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
+        Ok(())
+    }
+
+    /// Clear TM state on the GPU.
+    fn reset(&mut self) -> PyResult<()> {
+        self.tm_gpu.reset().map_err(|e| {
+            pyo3::exceptions::PyRuntimeError::new_err(format!("GPU TM reset: {e:?}"))
+        })?;
+        self.fused_state.reset().map_err(|e| {
+            pyo3::exceptions::PyRuntimeError::new_err(format!("GPU fused reset: {e:?}"))
+        })
+    }
+
+    /// FUSED MEGAKERNEL PATH: single CUDA launch for the entire T-step
+    /// forward (SP + TM all in one). Accepts torch CUDA tensors via
+    /// `__cuda_array_interface__` (zero-copy). Writes active-column mask +
+    /// anomaly directly into caller-allocated torch tensors.
+    ///
+    /// Semantics diverge from `step_many_cuda` in one important way: column
+    /// activation uses per-column threshold inhibition instead of global
+    /// top-K. The threshold is EMA-adapted per column toward the sparsity
+    /// target. See `docs/GPU_HTM.md` §Fused Kernel.
+    #[pyo3(signature = (sdr_cai, cols_cai, anom_cai, learn=true))]
+    fn step_many_fused_cuda(
+        &mut self,
+        py: Python<'_>,
+        sdr_cai: &Bound<'_, PyDict>,
+        cols_cai: &Bound<'_, PyDict>,
+        anom_cai: &Bound<'_, PyDict>,
+        learn: bool,
+    ) -> PyResult<()> {
+        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(sdr_cai)?;
+        let (cols_ptr, cols_shape, cols_type) = cai_parse(cols_cai)?;
+        let (anom_ptr, anom_shape, anom_type) = cai_parse(anom_cai)?;
+
+        if sdr_type != "|u1" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "sdr_cai typestr must be '|u1' (uint8), got {sdr_type}",
+            )));
+        }
+        if cols_type != "|u1" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "cols_cai typestr must be '|u1' (uint8), got {cols_type}",
+            )));
+        }
+        if anom_type != "<f4" && anom_type != "=f4" {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "anom_cai typestr must be '<f4' (float32), got {anom_type}",
+            )));
+        }
+
+        if sdr_shape.len() != 2 || sdr_shape[1] != self.input_bits {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "sdr_cai shape {sdr_shape:?} != (T, {})",
+                self.input_bits,
+            )));
+        }
+        let t = sdr_shape[0];
+        if cols_shape != [t, self.n_columns] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "cols_cai shape {cols_shape:?} != ({t}, {})",
+                self.n_columns,
+            )));
+        }
+        if anom_shape != [t] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "anom_cai shape {anom_shape:?} != ({t},)",
+            )));
+        }
+
+        let dev = self.sp_gpu.dev_ref().clone();
+        let n_cols = self.n_columns;
+        let input_bits = self.input_bits;
+
+        let result = py.allow_threads(|| -> Result<(), String> {
+            let inputs_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<u8>(sdr_ptr, t * input_bits)
+            });
+            let mut cols_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols)
+            });
+            let mut anom_dev = ManuallyDrop::new(unsafe {
+                dev.upgrade_device_ptr::<f32>(anom_ptr, t)
+            });
+
+            fused::launch_fused(
+                &mut self.sp_gpu,
+                &mut self.tm_gpu,
+                &mut self.fused_state,
+                &inputs_dev,
+                &mut cols_dev,
+                &mut anom_dev,
+                t,
+                input_bits,
+                learn,
+            ).map_err(|e| format!("launch_fused: {e:?}"))?;
+
+            // No dev.synchronize() here: caller must explicitly sync via the
+            // `device_sync()` method (or PyTorch auto-syncs when the output
+            // tensor is next consumed). Removing the per-launch barrier lets
+            // subsequent GPU work (mamba3 fwd, etc.) overlap in time.
+            Ok(())
+        });
+
+        result.map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
+        Ok(())
+    }
+
+    /// Explicit device synchronization — the caller must invoke this after
+    /// all batched `step_many_*_cuda` calls complete, before reading the
+    /// output tensors from a different CUDA stream. Equivalent to the old
+    /// per-call `dev.synchronize()` that was removed for overlap.
+    fn device_sync(&self) -> PyResult<()> {
+        let dev = self.sp_gpu.dev_ref();
+        dev.synchronize()
+            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("sync: {e:?}")))?;
+        Ok(())
+    }
+}
+
+/// Batch B regions into ONE cooperative kernel launch. Breaks through the
+/// CUDA cooperative-kernel device-level serialization: a single cooperative
+/// launch with grid.y=B processes all regions concurrently — ~B× speedup
+/// over B sequential launches.
+///
+/// All regions must have the same config (input_bits, n_columns,
+/// cells_per_column). Each region keeps its independent GPU state.
+/// Does NOT sync; caller must invoke `device_sync()` on any region
+/// afterwards (or rely on a downstream torch op to auto-sync).
+#[pyfunction]
+#[pyo3(signature = (regions, sdr_cais, cols_cais, anom_cais, learn=true))]
+fn step_batch_fused_cuda(
+    py: Python<'_>,
+    regions: Vec<Py<HTMRegionGpu>>,
+    sdr_cais: Vec<Bound<'_, PyDict>>,
+    cols_cais: Vec<Bound<'_, PyDict>>,
+    anom_cais: Vec<Bound<'_, PyDict>>,
+    learn: bool,
+) -> PyResult<()> {
+    let b = regions.len();
+    if b == 0 {
+        return Err(pyo3::exceptions::PyValueError::new_err("regions is empty"));
+    }
+    if sdr_cais.len() != b || cols_cais.len() != b || anom_cais.len() != b {
+        return Err(pyo3::exceptions::PyValueError::new_err(
+            "sdr_cais / cols_cais / anom_cais length must match regions",
+        ));
+    }
+
+    // Parse all CAI dicts; collect device pointers. Validate shapes/dtypes.
+    let mut sdr_ptrs = Vec::with_capacity(b);
+    let mut cols_ptrs = Vec::with_capacity(b);
+    let mut anom_ptrs = Vec::with_capacity(b);
+    let (input_bits, n_columns, t) = {
+        let r0 = regions[0].bind(py).borrow();
+        (r0.input_bits, r0.n_columns, {
+            let (_p, sh, _ty) = cai_parse(&sdr_cais[0])?;
+            if sh.len() != 2 {
+                return Err(pyo3::exceptions::PyValueError::new_err(
+                    format!("sdr_cai must be 2-D (T, input_bits), got {sh:?}"),
+                ));
+            }
+            sh[0]
+        })
+    };
+
+    for i in 0..b {
+        let (sdr_ptr, sdr_shape, sdr_type) = cai_parse(&sdr_cais[i])?;
+        let (cols_ptr, cols_shape, cols_type) = cai_parse(&cols_cais[i])?;
+        let (anom_ptr, anom_shape, anom_type) = cai_parse(&anom_cais[i])?;
+        if sdr_type != "|u1" || cols_type != "|u1" {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "sdr/cols typestr must be '|u1' (uint8)",
+            ));
+        }
+        if anom_type != "<f4" && anom_type != "=f4" {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "anom typestr must be '<f4' (float32)",
+            ));
+        }
+        if sdr_shape != [t, input_bits] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "sdr[{i}] shape {sdr_shape:?} != ({t}, {input_bits})"
+            )));
+        }
+        if cols_shape != [t, n_columns] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "cols[{i}] shape {cols_shape:?} != ({t}, {n_columns})"
+            )));
+        }
+        if anom_shape != [t] {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "anom[{i}] shape {anom_shape:?} != ({t},)"
+            )));
+        }
+        sdr_ptrs.push(sdr_ptr);
+        cols_ptrs.push(cols_ptr);
+        anom_ptrs.push(anom_ptr);
+    }
+
+    // Exclusively borrow each region. PyRefMut guarantees uniqueness.
+    let mut region_refs: Vec<pyo3::PyRefMut<HTMRegionGpu>> =
+        regions.iter().map(|p| p.bind(py).borrow_mut()).collect();
+    // Collect raw mutable pointers — each PyRefMut exclusively borrows its
+    // region for the lifetime of this call, so pointers stay valid and
+    // unique. launch_fused_batched_raw only dereferences one region at a
+    // time, not constructing an aliased slice.
+    let raw_ptrs: Vec<*mut HTMRegionGpu> = region_refs
+        .iter_mut()
+        .map(|r| &mut **r as *mut HTMRegionGpu)
+        .collect();
+
+    // No allow_threads: raw pointers aren't Send. The launch is GPU-queued
+    // and sync'd downstream; holding the GIL for the duration is cheap.
+    fused::launch_fused_batched_raw(
+        &raw_ptrs, &sdr_ptrs, &cols_ptrs, &anom_ptrs,
+        t, input_bits, learn,
+    )
+    .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("launch_fused_batched: {e:?}")))?;
+    Ok(())
+}
+
+pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<HTMRegionGpu>()?;
+    m.add_function(pyo3::wrap_pyfunction!(step_batch_fused_cuda, m)?)?;
+    Ok(())
+}

overlay/htm_rust/src/gpu/sp_gpu.rs

@@ -1,796 +1,796 @@
-//! GPU implementation of the Spatial Pooler.
-//!
-//! One `SpatialPoolerGpu` owns a set of persistent device buffers + 4 PTX
-//! kernels. `compute(input, learn)` performs one SP step and returns the
-//! sorted active-column indices (host `Vec<u32>`) — this is what the CPU
-//! TemporalMemory consumes.
-//!
-//! Persistent state on device (per region):
-//!   syn_bit     : u32  [n_columns × S]  (constant after init)
-//!   syn_perm    : f32  [n_columns × S]  (updated by sp_learn)
-//!   boost       : f32  [n_columns]
-//!   active_duty : f32  [n_columns]
-//!   overlap_duty: f32  [n_columns]
-//!
-//! Per-step transient state:
-//!   inp_dev     : u8   [input_bits]     (H2D copy each step)
-//!   raw         : u32  [n_columns]
-//!   boosted     : f32  [n_columns]
-//!   active_mask : u8   [n_columns]      (topk output, D2H at the end)
-
-use std::sync::Arc;
-
-use cudarc::driver::{CudaDevice, CudaSlice, DeviceSlice, DriverError, LaunchAsync, LaunchConfig};
-use cudarc::nvrtc::Ptx;
-
-use crate::sp::SpatialPooler;
-
-// Embed PTX at compile time. OUT_DIR is set by build.rs.
-const PTX_SP_OVERLAP: &str =
-    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_overlap.ptx"));
-const PTX_SP_TOPK: &str =
-    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_topk.ptx"));
-const PTX_SP_LEARN: &str =
-    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_learn.ptx"));
-const PTX_SP_DUTY: &str =
-    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_duty.ptx"));
-const PTX_SP_BOOST_FUSED: &str =
-    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_boost_fused.ptx"));
-
-pub struct SpatialPoolerGpu {
-    dev: Arc<CudaDevice>,
-
-    // Config mirror (we don't touch CPU SpatialPooler after init).
-    input_bits: usize,
-    n_columns: usize,
-    synapses_per_col: usize,
-    conn_thr: f32,
-    inc: f32,
-    dec: f32,
-    sparsity: f32,
-    duty_period: f32,
-    boost_strength: f32,
-
-    // Persistent device state.
-    syn_bit: CudaSlice<u32>,
-    syn_perm: CudaSlice<f32>,
-    boost: CudaSlice<f32>,
-    active_duty: CudaSlice<f32>,
-    overlap_duty: CudaSlice<f32>,
-
-    // Transient scratch (reused each step).
-    inp_dev: CudaSlice<u8>,
-    raw: CudaSlice<u32>,
-    boosted: CudaSlice<f32>,
-    active_mask: CudaSlice<u8>,
-
-    // Reusable host buffer for D2H of active_mask.
-    host_mask: Vec<u8>,
-
-    /// Strict bit-parity with CPU reference. Enabled for tests.
-    /// Forces host-side boost/exp computation and the overlap-duty bump check
-    /// every step. Default false for max throughput.
-    strict_parity: bool,
-}
-
-impl SpatialPoolerGpu {
-    /// Copy CPU SpatialPooler state onto the device. This preserves the
-    /// exact seeded proximal synapse layout + initial permanences, so the
-    /// GPU SP is a bit-identical parallel implementation of the CPU SP.
-    pub fn from_cpu(cpu: &SpatialPooler) -> Result<Self, DriverError> {
-        let dev = CudaDevice::new(0)?;
-        let cfg = &cpu.cfg;
-        let n = cfg.n_columns;
-        let s = cfg.potential_synapses;
-
-        // Flatten proximal dendrites into column-major arrays.
-        let mut syn_bit_h: Vec<u32> = Vec::with_capacity(n * s);
-        let mut syn_perm_h: Vec<f32> = Vec::with_capacity(n * s);
-        for col in &cpu.columns {
-            debug_assert_eq!(col.inputs.len(), s);
-            debug_assert_eq!(col.perms.len(), s);
-            syn_bit_h.extend_from_slice(&col.inputs);
-            syn_perm_h.extend_from_slice(&col.perms);
-        }
-
-        let syn_bit = dev.htod_sync_copy(&syn_bit_h)?;
-        let syn_perm = dev.htod_sync_copy(&syn_perm_h)?;
-        let boost = dev.htod_sync_copy(&cpu.boost)?;
-        let active_duty = dev.htod_sync_copy(&cpu.active_duty_cycle)?;
-        let overlap_duty = dev.htod_sync_copy(&cpu.overlap_duty_cycle)?;
-
-        let inp_dev: CudaSlice<u8> = dev.alloc_zeros(cfg.input_bits)?;
-        let raw: CudaSlice<u32> = dev.alloc_zeros(n)?;
-        let boosted: CudaSlice<f32> = dev.alloc_zeros(n)?;
-        let active_mask: CudaSlice<u8> = dev.alloc_zeros(n)?;
-
-        // Load PTX modules. Each .ptx is a module containing one `extern "C"`
-        // function; we tag them by unique module names so multiple SP instances
-        // don't collide (cudarc uses the (module, func) pair).
-        // Actually: CudaDevice::load_ptx stores under the given module name
-        // globally on the device, so we use a deterministic naming scheme.
-        let modules = [
-            ("htm_sp_overlap", PTX_SP_OVERLAP, "sp_overlap"),
-            ("htm_sp_topk", PTX_SP_TOPK, "sp_topk_select"),
-            ("htm_sp_learn", PTX_SP_LEARN, "sp_learn"),
-            ("htm_sp_duty", PTX_SP_DUTY, "sp_duty_update"),
-            ("htm_sp_boost_fused", PTX_SP_BOOST_FUSED, "sp_boost_from_duty"),
-        ];
-        for (modname, ptx, fnname) in modules {
-            // load_ptx is NOT idempotent — calling twice errors. For multi-region
-            // support we check-then-load.
-            if dev.get_func(modname, fnname).is_none() {
-                dev.load_ptx(Ptx::from_src(ptx), modname, &[fnname])?;
-            }
-        }
-
-        Ok(Self {
-            dev,
-            input_bits: cfg.input_bits,
-            n_columns: n,
-            synapses_per_col: s,
-            conn_thr: cfg.connected_threshold,
-            inc: cfg.syn_perm_active_inc,
-            dec: cfg.syn_perm_inactive_dec,
-            sparsity: cfg.sparsity,
-            duty_period: cfg.duty_cycle_period,
-            boost_strength: cfg.boost_strength,
-            syn_bit,
-            syn_perm,
-            boost,
-            active_duty,
-            overlap_duty,
-            inp_dev,
-            raw,
-            boosted,
-            active_mask,
-            host_mask: vec![0u8; n],
-            strict_parity: false,
-        })
-    }
-
-    /// Enable strict bit-parity mode. Parity tests use this.
-    pub fn set_strict_parity(&mut self, strict: bool) {
-        self.strict_parity = strict;
-    }
-
-    /// Access to the underlying CudaDevice for host-side orchestration.
-    pub fn dev_ref(&self) -> &Arc<CudaDevice> {
-        &self.dev
-    }
-
-    // --- Fused-path accessors (immutable state reads + pointer-grabs). ---
-    pub fn n_columns_accessor(&self) -> usize { self.n_columns }
-    #[allow(dead_code)]
-    pub fn input_bits_accessor(&self) -> usize { self.input_bits }
-    pub fn synapses_per_col_accessor(&self) -> usize { self.synapses_per_col }
-    pub fn conn_thr_accessor(&self) -> f32 { self.conn_thr }
-    pub fn inc_accessor(&self) -> f32 { self.inc }
-    pub fn dec_accessor(&self) -> f32 { self.dec }
-    pub fn sparsity_accessor(&self) -> f32 { self.sparsity }
-    pub fn duty_period_accessor(&self) -> f32 { self.duty_period }
-    #[allow(dead_code)]
-    pub fn boost_strength_accessor(&self) -> f32 { self.boost_strength }
-
-    pub fn syn_bit_accessor(&self) -> &CudaSlice<u32> { &self.syn_bit }
-    pub fn syn_perm_accessor(&self) -> &CudaSlice<f32> { &self.syn_perm }
-    pub fn boost_accessor(&self) -> &CudaSlice<f32> { &self.boost }
-    pub fn active_duty_accessor(&self) -> &CudaSlice<f32> { &self.active_duty }
-
-    /// Compute the 95th-percentile-like initial threshold from raw overlaps
-    /// after a short warmup pass. Used to seed `inhibition_threshold` such
-    /// that activation rate starts near the sparsity target.
-    /// Placeholder (returns a conservative constant); real warmup pass
-    /// happens on the Rust orchestrator side.
-    pub fn initial_threshold_estimate(&self) -> f32 {
-        // With conn_thr=0.5, init_perm around 0.5±0.1, S=40, sparse SDR at 2%:
-        // expected overlap ~ 40 * 0.02 = 0.8 connected hits → boosted ~ 0.8.
-        // Top-K selects top 2%, so threshold for top 2% is roughly the
-        // 98th-percentile of boosted. Conservative start: 2.0.
-        // The per-column adaptation will quickly steer each column's thr.
-        2.0f32
-    }
-
-    /// Batched multi-step SP on the GPU. Processes T timesteps from a
-    /// pre-uploaded device input buffer. Emits `(T, n_cols)` u8 active-column
-    /// mask to `cols_dev_out` and `(T,)` active column index list (in a
-    /// per-step window of size k, padded with u32::MAX).
-    ///
-    /// For each step, this runs the same 5-kernel pipeline as `compute`, but
-    /// skips the per-step boost/duty D2H→exp→H2D round-trip: instead it
-    /// accumulates to a host scratch once every `boost_interval` steps.
-    ///
-    /// This is the fast path used by `HTMRegionGpu.step_many_gpu`.
-    #[allow(clippy::too_many_arguments)]
-    pub fn step_batch(
-        &mut self,
-        inputs_flat_dev: &CudaSlice<u8>,
-        t: usize,
-        input_bits: usize,
-        learn: bool,
-        cols_out: &mut [u8],
-        active_indices_host: &mut Vec<u32>,
-    ) -> Result<(), DriverError> {
-        let n = self.n_columns;
-        let k = ((self.sparsity * n as f32).round() as usize).max(1);
-        debug_assert_eq!(cols_out.len(), t * n);
-
-        let overlap_fn = self.dev.get_func("htm_sp_overlap", "sp_overlap").unwrap();
-        let topk_fn = self.dev.get_func("htm_sp_topk", "sp_topk_select").unwrap();
-        let learn_fn = self.dev.get_func("htm_sp_learn", "sp_learn").unwrap();
-        let duty_fn = self.dev.get_func("htm_sp_duty", "sp_duty_update").unwrap();
-
-        let overlap_cfg = LaunchConfig {
-            grid_dim: (n as u32, 1, 1),
-            block_dim: (128, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        let topk_cfg = LaunchConfig {
-            grid_dim: (1, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
-        };
-        let learn_cfg = overlap_cfg;
-        let duty_cfg = LaunchConfig {
-            grid_dim: ((n as u32 + 255) / 256, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        let alpha = 1.0f32 / self.duty_period.max(1.0);
-
-        // Reusable host buffer for the per-step active_mask D2H.
-        self.host_mask.resize(n, 0);
-
-        active_indices_host.clear();
-
-        for ti in 0..t {
-            // Point overlap kernel at the ti-th slice of the pre-uploaded input.
-            // cudarc CudaSlice doesn't have a "view" per se, so we must copy the
-            // slice into the reusable inp_dev buffer. This is a D2D copy — much
-            // faster than H2D.
-            // (Alternative: rewrite kernel to accept an offset; deferred.)
-            let in_off = ti * input_bits;
-            // Use dtod_copy via raw slice indexing: cudarc exposes slice() for this.
-            let sub = inputs_flat_dev.slice(in_off..in_off + input_bits);
-            self.dev.dtod_copy(&sub, &mut self.inp_dev)?;
-
-            // 1. sp_overlap
-            unsafe {
-                overlap_fn.clone().launch(
-                    overlap_cfg,
-                    (
-                        &self.inp_dev,
-                        &self.syn_bit,
-                        &self.syn_perm,
-                        &self.boost,
-                        self.conn_thr,
-                        self.synapses_per_col as u32,
-                        n as u32,
-                        &mut self.raw,
-                        &mut self.boosted,
-                    ),
-                )?;
-            }
-
-            // 2. Clear active_mask, then sp_topk
-            self.dev.memset_zeros(&mut self.active_mask)?;
-            unsafe {
-                topk_fn.clone().launch(
-                    topk_cfg,
-                    (&self.boosted, n as u32, k as u32, &mut self.active_mask),
-                )?;
-            }
-
-            // 3. sp_learn
-            if learn {
-                unsafe {
-                    learn_fn.clone().launch(
-                        learn_cfg,
-                        (
-                            &self.active_mask,
-                            &self.inp_dev,
-                            &self.syn_bit,
-                            &mut self.syn_perm,
-                            self.inc,
-                            self.dec,
-                            self.synapses_per_col as u32,
-                            n as u32,
-                        ),
-                    )?;
-                }
-            }
-
-            // 4. duty update (device)
-            unsafe {
-                duty_fn.clone().launch(
-                    duty_cfg,
-                    (
-                        &self.active_mask,
-                        &self.raw,
-                        &mut self.active_duty,
-                        &mut self.overlap_duty,
-                        &mut self.boost,
-                        alpha,
-                        1.0f32,
-                        0.0f32,
-                        0.0f32,
-                        0u32,
-                        n as u32,
-                    ),
-                )?;
-            }
-
-            // 5. Boost update. Two modes:
-            //    * strict_parity (tests): host-side exp for bit-exact match.
-            //    * default (production): GPU expf is close enough and ~10x faster
-            //      since we skip the D2H/H2D round-trip.
-            if learn && self.boost_strength > 0.0 {
-                if self.strict_parity {
-                    let mut duty_host = vec![0f32; n];
-                    self.dev
-                        .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
-                    let sum: f32 = duty_host.iter().sum();
-                    let mean = sum / (n as f32);
-                    let mut boost_host = vec![0f32; n];
-                    for i in 0..n {
-                        boost_host[i] =
-                            (-self.boost_strength * (duty_host[i] - mean)).exp();
-                    }
-                    self.dev.htod_sync_copy_into(&boost_host, &mut self.boost)?;
-
-                    // Permanence bump (rare). Only evaluated in strict mode.
-                    let mut ov_host = vec![0f32; n];
-                    self.dev
-                        .dtoh_sync_copy_into(&self.overlap_duty, &mut ov_host)?;
-                    let max_ov = ov_host.iter().cloned().fold(0f32, f32::max);
-                    if max_ov > 0.0 {
-                        let thr = 0.001f32 * max_ov;
-                        let bump = self.inc * 0.1f32;
-                        let bump_cols: Vec<u32> = ov_host
-                            .iter()
-                            .enumerate()
-                            .filter_map(|(i, &o)| {
-                                if o < thr { Some(i as u32) } else { None }
-                            })
-                            .collect();
-                        if !bump_cols.is_empty() {
-                            let s = self.synapses_per_col;
-                            let mut perm_host = vec![0f32; n * s];
-                            self.dev
-                                .dtoh_sync_copy_into(&self.syn_perm, &mut perm_host)?;
-                            for &c in &bump_cols {
-                                let base = (c as usize) * s;
-                                for p in &mut perm_host[base..base + s] {
-                                    *p = (*p + bump).min(1.0);
-                                }
-                            }
-                            self.dev.htod_sync_copy_into(&perm_host, &mut self.syn_perm)?;
-                        }
-                    }
-                } else {
-                    // Fast path: fused mean + boost = expf(-strength*(ad-mean))
-                    // in a single GPU block. Zero D2H, zero H2D — fully async.
-                    let boost_fn = self
-                        .dev
-                        .get_func("htm_sp_boost_fused", "sp_boost_from_duty")
-                        .expect("sp_boost_fused not loaded");
-                    let boost_cfg = LaunchConfig {
-                        grid_dim: (1, 1, 1),
-                        block_dim: (1024, 1, 1),
-                        shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
-                    };
-                    unsafe {
-                        boost_fn.launch(
-                            boost_cfg,
-                            (
-                                &self.active_duty,
-                                &mut self.boost,
-                                self.boost_strength,
-                                n as u32,
-                            ),
-                        )?;
-                    }
-                }
-            }
-
-            // D2H the active_mask for this step. This is the single
-            // unavoidable sync point per step — CPU TM needs the active
-            // indices for its next state update. At 2048 bytes / step this
-            // is tiny in bandwidth but costs a full syncronize (~5-10μs).
-            self.dev
-                .dtoh_sync_copy_into(&self.active_mask, &mut self.host_mask)?;
-            let co = ti * n;
-            cols_out[co..co + n].copy_from_slice(&self.host_mask);
-            // Extract active indices.
-            for (i, &b) in self.host_mask.iter().enumerate() {
-                if b != 0 {
-                    active_indices_host.push(i as u32);
-                }
-            }
-            // Insert separator (u32::MAX) between steps to demarcate step boundaries.
-            active_indices_host.push(u32::MAX);
-        }
-
-        Ok(())
-    }
-
-    /// Fully-on-GPU batched SP + TM. Zero per-step host sync.
-    ///
-    /// Inputs:
-    ///   inputs_flat_dev : (T * input_bits) u8 already uploaded
-    ///   cols_dev        : (T * n_cols) u8 output — active-column mask per step
-    ///   anom_dev        : (T,) f32 output — anomaly score per step
-    ///   tm              : persistent GPU TemporalMemory for this region
-    #[allow(clippy::too_many_arguments)]
-    pub fn step_batch_with_tm(
-        &mut self,
-        inputs_flat_dev: &CudaSlice<u8>,
-        t: usize,
-        input_bits: usize,
-        learn: bool,
-        cols_dev: &mut CudaSlice<u8>,
-        anom_dev: &mut CudaSlice<f32>,
-        tm: &mut crate::gpu::tm_gpu::TemporalMemoryGpu,
-    ) -> Result<(), DriverError> {
-        let n = self.n_columns;
-        let k = ((self.sparsity * n as f32).round() as usize).max(1);
-        debug_assert_eq!(cols_dev.len(), t * n);
-        debug_assert_eq!(anom_dev.len(), t);
-
-        let overlap_fn = self.dev.get_func("htm_sp_overlap", "sp_overlap").unwrap();
-        let topk_fn    = self.dev.get_func("htm_sp_topk", "sp_topk_select").unwrap();
-        let learn_fn   = self.dev.get_func("htm_sp_learn", "sp_learn").unwrap();
-        let duty_fn    = self.dev.get_func("htm_sp_duty", "sp_duty_update").unwrap();
-
-        let overlap_cfg = LaunchConfig {
-            grid_dim: (n as u32, 1, 1),
-            block_dim: (128, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        let topk_cfg = LaunchConfig {
-            grid_dim: (1, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
-        };
-        let learn_cfg = overlap_cfg;
-        let duty_cfg = LaunchConfig {
-            grid_dim: ((n as u32 + 255) / 256, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        let alpha = 1.0f32 / self.duty_period.max(1.0);
-
-        for ti in 0..t {
-            let in_off = ti * input_bits;
-            let sub = inputs_flat_dev.slice(in_off..in_off + input_bits);
-            self.dev.dtod_copy(&sub, &mut self.inp_dev)?;
-
-            // 1. sp_overlap
-            unsafe {
-                overlap_fn.clone().launch(
-                    overlap_cfg,
-                    (
-                        &self.inp_dev,
-                        &self.syn_bit,
-                        &self.syn_perm,
-                        &self.boost,
-                        self.conn_thr,
-                        self.synapses_per_col as u32,
-                        n as u32,
-                        &mut self.raw,
-                        &mut self.boosted,
-                    ),
-                )?;
-            }
-
-            // 2. clear + sp_topk
-            self.dev.memset_zeros(&mut self.active_mask)?;
-            unsafe {
-                topk_fn.clone().launch(
-                    topk_cfg,
-                    (&self.boosted, n as u32, k as u32, &mut self.active_mask),
-                )?;
-            }
-
-            // 3. sp_learn
-            if learn {
-                unsafe {
-                    learn_fn.clone().launch(
-                        learn_cfg,
-                        (
-                            &self.active_mask,
-                            &self.inp_dev,
-                            &self.syn_bit,
-                            &mut self.syn_perm,
-                            self.inc,
-                            self.dec,
-                            self.synapses_per_col as u32,
-                            n as u32,
-                        ),
-                    )?;
-                }
-            }
-
-            // 4. duty update (stage 1: no-boost write)
-            unsafe {
-                duty_fn.clone().launch(
-                    duty_cfg,
-                    (
-                        &self.active_mask,
-                        &self.raw,
-                        &mut self.active_duty,
-                        &mut self.overlap_duty,
-                        &mut self.boost,
-                        alpha,
-                        1.0f32,
-                        0.0f32,
-                        0.0f32,
-                        0u32,
-                        n as u32,
-                    ),
-                )?;
-            }
-
-            // 5. Boost update: fused GPU kernel (no D2H).
-            if learn && self.boost_strength > 0.0 {
-                let boost_fn = self.dev
-                    .get_func("htm_sp_boost_fused", "sp_boost_from_duty")
-                    .expect("sp_boost_fused not loaded");
-                let boost_cfg = LaunchConfig {
-                    grid_dim: (1, 1, 1),
-                    block_dim: (1024, 1, 1),
-                    shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
-                };
-                unsafe {
-                    boost_fn.launch(
-                        boost_cfg,
-                        (
-                            &self.active_duty,
-                            &mut self.boost,
-                            self.boost_strength,
-                            n as u32,
-                        ),
-                    )?;
-                }
-            }
-
-            // 6. Copy active_mask slice into cols_dev[ti*n .. (ti+1)*n].
-            let mut dst_slice = cols_dev.slice_mut(ti * n..(ti + 1) * n);
-            self.dev.dtod_copy(&self.active_mask, &mut dst_slice)?;
-
-            // 7. GPU TM step: predict + activate + anomaly + learn, all on device.
-            tm.step(&self.active_mask, anom_dev, ti as u32, learn)?;
-        }
-
-        Ok(())
-    }
-
-    /// One SP step on the GPU. Returns sorted active-column indices.
-    pub fn compute(&mut self, input: &[u8], learn: bool) -> Result<Vec<u32>, DriverError> {
-        debug_assert_eq!(input.len(), self.input_bits);
-        let n = self.n_columns;
-        let k = ((self.sparsity * n as f32).round() as usize).max(1);
-
-        // 1. H2D input SDR.
-        self.dev.htod_sync_copy_into(input, &mut self.inp_dev)?;
-
-        // 2. Launch sp_overlap: grid=n_columns, block=128.
-        let overlap_fn = self
-            .dev
-            .get_func("htm_sp_overlap", "sp_overlap")
-            .expect("sp_overlap not loaded");
-        let overlap_cfg = LaunchConfig {
-            grid_dim: (n as u32, 1, 1),
-            block_dim: (128, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        unsafe {
-            overlap_fn.launch(
-                overlap_cfg,
-                (
-                    &self.inp_dev,
-                    &self.syn_bit,
-                    &self.syn_perm,
-                    &self.boost,
-                    self.conn_thr,
-                    self.synapses_per_col as u32,
-                    n as u32,
-                    &mut self.raw,
-                    &mut self.boosted,
-                ),
-            )?;
-        }
-
-        // 3. Launch sp_topk: single block, shared mem = n_columns * f32.
-        let topk_fn = self
-            .dev
-            .get_func("htm_sp_topk", "sp_topk_select")
-            .expect("sp_topk not loaded");
-        let topk_cfg = LaunchConfig {
-            grid_dim: (1, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
-        };
-        // Clear active_mask first. memset_zeros avoids an H2D of a host
-        // zeroes vector every step.
-        self.dev.memset_zeros(&mut self.active_mask)?;
-        unsafe {
-            topk_fn.launch(
-                topk_cfg,
-                (
-                    &self.boosted,
-                    n as u32,
-                    k as u32,
-                    &mut self.active_mask,
-                ),
-            )?;
-        }
-
-        // 4. Optional: sp_learn on active columns.
-        if learn {
-            let learn_fn = self
-                .dev
-                .get_func("htm_sp_learn", "sp_learn")
-                .expect("sp_learn not loaded");
-            let learn_cfg = LaunchConfig {
-                grid_dim: (n as u32, 1, 1),
-                block_dim: (128, 1, 1),
-                shared_mem_bytes: 0,
-            };
-            unsafe {
-                learn_fn.launch(
-                    learn_cfg,
-                    (
-                        &self.active_mask,
-                        &self.inp_dev,
-                        &self.syn_bit,
-                        &mut self.syn_perm,
-                        self.inc,
-                        self.dec,
-                        self.synapses_per_col as u32,
-                        n as u32,
-                    ),
-                )?;
-            }
-        }
-
-        // 5. Duty cycle + boost update. Always runs (matches CPU).
-        //    We need mean_duty on the host — compute BEFORE the update (matches
-        //    CPU sp.rs line 200-205 where mean is computed then written).
-        //    Actually CPU computes mean of the PRE-update duty cycles too? Re-read:
-        //      sp.rs lines 186-196 update duty cycles (pre-mean).
-        //      Line 202: mean = sum(active_duty_cycle) / n  ← after update.
-        //      Line 204: boost[i] = exp(-strength*(active_duty[i] - mean)).
-        //    So mean is on POST-update values.
-        //    Easiest: 1) run duty update with boost_strength=0 (skip boost calc),
-        //             2) D2H active_duty, compute mean, 3) run a boost-only kernel
-        //                OR inline the exp() in a second launch with mean passed.
-        //
-        //    For simplicity and correctness we fuse: run the duty kernel with
-        //    mean=0 and boost_strength=0 (disables boost write), then D2H to
-        //    compute mean, then re-launch with the true mean. Two launches, one
-        //    tiny D2H (n × f32). At n=2048 this is 8KB per step — negligible.
-        let alpha = 1.0f32 / self.duty_period.max(1.0);
-        let duty_fn = self
-            .dev
-            .get_func("htm_sp_duty", "sp_duty_update")
-            .expect("sp_duty not loaded");
-        let duty_cfg = LaunchConfig {
-            grid_dim: ((n as u32 + 255) / 256, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        // Stage 1: update duty cycles (boost_strength=0 -> no write).
-        unsafe {
-            duty_fn.launch(
-                duty_cfg,
-                (
-                    &self.active_mask,
-                    &self.raw,
-                    &mut self.active_duty,
-                    &mut self.overlap_duty,
-                    &mut self.boost,
-                    alpha,
-                    1.0f32,   // stim_thr
-                    0.0f32,   // boost_strength = 0 -> skip write
-                    0.0f32,   // mean_duty (unused)
-                    0u32,     // learn_flag = 0
-                    n as u32,
-                ),
-            )?;
-        }
-
-        if learn && self.boost_strength > 0.0 && self.strict_parity {
-            // Boost update must bit-match CPU `f32::exp`, so we compute it on
-            // the host and copy back. Cost per step: 8KB D2H + 8KB H2D at n=2048.
-            // Critical for learning parity — CUDA expf (even without fast-math)
-            // uses different rounding for some inputs than host libm.
-            let mut duty_host = vec![0f32; n];
-            self.dev
-                .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
-            let sum: f32 = duty_host.iter().sum();
-            let mean = sum / (n as f32);
-            let mut boost_host = vec![0f32; n];
-            for i in 0..n {
-                boost_host[i] = (-self.boost_strength * (duty_host[i] - mean)).exp();
-            }
-            self.dev.htod_sync_copy_into(&boost_host, &mut self.boost)?;
-
-            // CPU sp.rs 210-226: permanence bump for chronically under-stimulated
-            // columns. If overlap_duty_cycle[i] < 0.001 * max(overlap_duty_cycle),
-            // add inc*0.1 to every synapse of column i (clamped to 1.0).
-            // This runs only once per step and only for the rare cases, but we
-            // need it for bit-exact parity with CPU learn.
-            let mut ov_host = vec![0f32; n];
-            self.dev
-                .dtoh_sync_copy_into(&self.overlap_duty, &mut ov_host)?;
-            let max_ov = ov_host.iter().cloned().fold(0f32, f32::max);
-            if max_ov > 0.0 {
-                let thr = 0.001f32 * max_ov;
-                let bump = self.inc * 0.1f32;
-                // Find columns needing a bump. Usually empty. Rare → D2H/H2D
-                // of syn_perm is cheap (n*S*4 = 320KB at n=2048,S=40).
-                let bump_cols: Vec<u32> = ov_host
-                    .iter()
-                    .enumerate()
-                    .filter_map(|(i, &o)| if o < thr { Some(i as u32) } else { None })
-                    .collect();
-                if !bump_cols.is_empty() {
-                    // Download, bump, upload. (Keeps implementation simple and
-                    // bit-exact. Could kernelize later.)
-                    let s = self.synapses_per_col;
-                    let mut perm_host = vec![0f32; n * s];
-                    self.dev.dtoh_sync_copy_into(&self.syn_perm, &mut perm_host)?;
-                    for &c in &bump_cols {
-                        let base = (c as usize) * s;
-                        for p in &mut perm_host[base..base + s] {
-                            *p = (*p + bump).min(1.0);
-                        }
-                    }
-                    self.dev.htod_sync_copy_into(&perm_host, &mut self.syn_perm)?;
-                }
-            }
-        } else if learn && self.boost_strength > 0.0 {
-            // Fast path: GPU-side boost using the already-loaded duty kernel.
-            let mut duty_host = vec![0f32; n];
-            self.dev
-                .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
-            let sum: f32 = duty_host.iter().sum();
-            let mean = sum / (n as f32);
-            let boost_fn = self
-                .dev
-                .get_func("htm_sp_duty", "sp_duty_update")
-                .expect("sp_duty not loaded");
-            unsafe {
-                boost_fn.launch(
-                    duty_cfg,
-                    (
-                        &self.active_mask,
-                        &self.raw,
-                        &mut self.active_duty,
-                        &mut self.overlap_duty,
-                        &mut self.boost,
-                        0.0f32,
-                        1.0f32,
-                        self.boost_strength,
-                        mean,
-                        1u32,
-                        n as u32,
-                    ),
-                )?;
-            }
-        }
-
-        // 6. D2H active_mask and convert to sorted index list.
-        self.dev
-            .dtoh_sync_copy_into(&self.active_mask, &mut self.host_mask)?;
-        let mut active: Vec<u32> = Vec::with_capacity(k);
-        for (i, &b) in self.host_mask.iter().enumerate() {
-            if b != 0 {
-                active.push(i as u32);
-            }
-        }
-        debug_assert_eq!(active.len(), k, "SP must emit exactly k winners");
-        Ok(active)
-    }
-}
+//! GPU implementation of the Spatial Pooler.
+//!
+//! One `SpatialPoolerGpu` owns a set of persistent device buffers + 4 PTX
+//! kernels. `compute(input, learn)` performs one SP step and returns the
+//! sorted active-column indices (host `Vec<u32>`) — this is what the CPU
+//! TemporalMemory consumes.
+//!
+//! Persistent state on device (per region):
+//!   syn_bit     : u32  [n_columns × S]  (constant after init)
+//!   syn_perm    : f32  [n_columns × S]  (updated by sp_learn)
+//!   boost       : f32  [n_columns]
+//!   active_duty : f32  [n_columns]
+//!   overlap_duty: f32  [n_columns]
+//!
+//! Per-step transient state:
+//!   inp_dev     : u8   [input_bits]     (H2D copy each step)
+//!   raw         : u32  [n_columns]
+//!   boosted     : f32  [n_columns]
+//!   active_mask : u8   [n_columns]      (topk output, D2H at the end)
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaDevice, CudaSlice, DeviceSlice, DriverError, LaunchAsync, LaunchConfig};
+use cudarc::nvrtc::Ptx;
+
+use crate::sp::SpatialPooler;
+
+// Embed PTX at compile time. OUT_DIR is set by build.rs.
+const PTX_SP_OVERLAP: &str =
+    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_overlap.ptx"));
+const PTX_SP_TOPK: &str =
+    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_topk.ptx"));
+const PTX_SP_LEARN: &str =
+    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_learn.ptx"));
+const PTX_SP_DUTY: &str =
+    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_duty.ptx"));
+const PTX_SP_BOOST_FUSED: &str =
+    include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/sp_boost_fused.ptx"));
+
+pub struct SpatialPoolerGpu {
+    dev: Arc<CudaDevice>,
+
+    // Config mirror (we don't touch CPU SpatialPooler after init).
+    input_bits: usize,
+    n_columns: usize,
+    synapses_per_col: usize,
+    conn_thr: f32,
+    inc: f32,
+    dec: f32,
+    sparsity: f32,
+    duty_period: f32,
+    boost_strength: f32,
+
+    // Persistent device state.
+    syn_bit: CudaSlice<u32>,
+    syn_perm: CudaSlice<f32>,
+    boost: CudaSlice<f32>,
+    active_duty: CudaSlice<f32>,
+    overlap_duty: CudaSlice<f32>,
+
+    // Transient scratch (reused each step).
+    inp_dev: CudaSlice<u8>,
+    raw: CudaSlice<u32>,
+    boosted: CudaSlice<f32>,
+    active_mask: CudaSlice<u8>,
+
+    // Reusable host buffer for D2H of active_mask.
+    host_mask: Vec<u8>,
+
+    /// Strict bit-parity with CPU reference. Enabled for tests.
+    /// Forces host-side boost/exp computation and the overlap-duty bump check
+    /// every step. Default false for max throughput.
+    strict_parity: bool,
+}
+
+impl SpatialPoolerGpu {
+    /// Copy CPU SpatialPooler state onto the device. This preserves the
+    /// exact seeded proximal synapse layout + initial permanences, so the
+    /// GPU SP is a bit-identical parallel implementation of the CPU SP.
+    pub fn from_cpu(cpu: &SpatialPooler) -> Result<Self, DriverError> {
+        let dev = CudaDevice::new(0)?;
+        let cfg = &cpu.cfg;
+        let n = cfg.n_columns;
+        let s = cfg.potential_synapses;
+
+        // Flatten proximal dendrites into column-major arrays.
+        let mut syn_bit_h: Vec<u32> = Vec::with_capacity(n * s);
+        let mut syn_perm_h: Vec<f32> = Vec::with_capacity(n * s);
+        for col in &cpu.columns {
+            debug_assert_eq!(col.inputs.len(), s);
+            debug_assert_eq!(col.perms.len(), s);
+            syn_bit_h.extend_from_slice(&col.inputs);
+            syn_perm_h.extend_from_slice(&col.perms);
+        }
+
+        let syn_bit = dev.htod_sync_copy(&syn_bit_h)?;
+        let syn_perm = dev.htod_sync_copy(&syn_perm_h)?;
+        let boost = dev.htod_sync_copy(&cpu.boost)?;
+        let active_duty = dev.htod_sync_copy(&cpu.active_duty_cycle)?;
+        let overlap_duty = dev.htod_sync_copy(&cpu.overlap_duty_cycle)?;
+
+        let inp_dev: CudaSlice<u8> = dev.alloc_zeros(cfg.input_bits)?;
+        let raw: CudaSlice<u32> = dev.alloc_zeros(n)?;
+        let boosted: CudaSlice<f32> = dev.alloc_zeros(n)?;
+        let active_mask: CudaSlice<u8> = dev.alloc_zeros(n)?;
+
+        // Load PTX modules. Each .ptx is a module containing one `extern "C"`
+        // function; we tag them by unique module names so multiple SP instances
+        // don't collide (cudarc uses the (module, func) pair).
+        // Actually: CudaDevice::load_ptx stores under the given module name
+        // globally on the device, so we use a deterministic naming scheme.
+        let modules = [
+            ("htm_sp_overlap", PTX_SP_OVERLAP, "sp_overlap"),
+            ("htm_sp_topk", PTX_SP_TOPK, "sp_topk_select"),
+            ("htm_sp_learn", PTX_SP_LEARN, "sp_learn"),
+            ("htm_sp_duty", PTX_SP_DUTY, "sp_duty_update"),
+            ("htm_sp_boost_fused", PTX_SP_BOOST_FUSED, "sp_boost_from_duty"),
+        ];
+        for (modname, ptx, fnname) in modules {
+            // load_ptx is NOT idempotent — calling twice errors. For multi-region
+            // support we check-then-load.
+            if dev.get_func(modname, fnname).is_none() {
+                dev.load_ptx(Ptx::from_src(ptx), modname, &[fnname])?;
+            }
+        }
+
+        Ok(Self {
+            dev,
+            input_bits: cfg.input_bits,
+            n_columns: n,
+            synapses_per_col: s,
+            conn_thr: cfg.connected_threshold,
+            inc: cfg.syn_perm_active_inc,
+            dec: cfg.syn_perm_inactive_dec,
+            sparsity: cfg.sparsity,
+            duty_period: cfg.duty_cycle_period,
+            boost_strength: cfg.boost_strength,
+            syn_bit,
+            syn_perm,
+            boost,
+            active_duty,
+            overlap_duty,
+            inp_dev,
+            raw,
+            boosted,
+            active_mask,
+            host_mask: vec![0u8; n],
+            strict_parity: false,
+        })
+    }
+
+    /// Enable strict bit-parity mode. Parity tests use this.
+    pub fn set_strict_parity(&mut self, strict: bool) {
+        self.strict_parity = strict;
+    }
+
+    /// Access to the underlying CudaDevice for host-side orchestration.
+    pub fn dev_ref(&self) -> &Arc<CudaDevice> {
+        &self.dev
+    }
+
+    // --- Fused-path accessors (immutable state reads + pointer-grabs). ---
+    pub fn n_columns_accessor(&self) -> usize { self.n_columns }
+    #[allow(dead_code)]
+    pub fn input_bits_accessor(&self) -> usize { self.input_bits }
+    pub fn synapses_per_col_accessor(&self) -> usize { self.synapses_per_col }
+    pub fn conn_thr_accessor(&self) -> f32 { self.conn_thr }
+    pub fn inc_accessor(&self) -> f32 { self.inc }
+    pub fn dec_accessor(&self) -> f32 { self.dec }
+    pub fn sparsity_accessor(&self) -> f32 { self.sparsity }
+    pub fn duty_period_accessor(&self) -> f32 { self.duty_period }
+    #[allow(dead_code)]
+    pub fn boost_strength_accessor(&self) -> f32 { self.boost_strength }
+
+    pub fn syn_bit_accessor(&self) -> &CudaSlice<u32> { &self.syn_bit }
+    pub fn syn_perm_accessor(&self) -> &CudaSlice<f32> { &self.syn_perm }
+    pub fn boost_accessor(&self) -> &CudaSlice<f32> { &self.boost }
+    pub fn active_duty_accessor(&self) -> &CudaSlice<f32> { &self.active_duty }
+
+    /// Compute the 95th-percentile-like initial threshold from raw overlaps
+    /// after a short warmup pass. Used to seed `inhibition_threshold` such
+    /// that activation rate starts near the sparsity target.
+    /// Placeholder (returns a conservative constant); real warmup pass
+    /// happens on the Rust orchestrator side.
+    pub fn initial_threshold_estimate(&self) -> f32 {
+        // With conn_thr=0.5, init_perm around 0.5±0.1, S=40, sparse SDR at 2%:
+        // expected overlap ~ 40 * 0.02 = 0.8 connected hits → boosted ~ 0.8.
+        // Top-K selects top 2%, so threshold for top 2% is roughly the
+        // 98th-percentile of boosted. Conservative start: 2.0.
+        // The per-column adaptation will quickly steer each column's thr.
+        2.0f32
+    }
+
+    /// Batched multi-step SP on the GPU. Processes T timesteps from a
+    /// pre-uploaded device input buffer. Emits `(T, n_cols)` u8 active-column
+    /// mask to `cols_dev_out` and `(T,)` active column index list (in a
+    /// per-step window of size k, padded with u32::MAX).
+    ///
+    /// For each step, this runs the same 5-kernel pipeline as `compute`, but
+    /// skips the per-step boost/duty D2H→exp→H2D round-trip: instead it
+    /// accumulates to a host scratch once every `boost_interval` steps.
+    ///
+    /// This is the fast path used by `HTMRegionGpu.step_many_gpu`.
+    #[allow(clippy::too_many_arguments)]
+    pub fn step_batch(
+        &mut self,
+        inputs_flat_dev: &CudaSlice<u8>,
+        t: usize,
+        input_bits: usize,
+        learn: bool,
+        cols_out: &mut [u8],
+        active_indices_host: &mut Vec<u32>,
+    ) -> Result<(), DriverError> {
+        let n = self.n_columns;
+        let k = ((self.sparsity * n as f32).round() as usize).max(1);
+        debug_assert_eq!(cols_out.len(), t * n);
+
+        let overlap_fn = self.dev.get_func("htm_sp_overlap", "sp_overlap").unwrap();
+        let topk_fn = self.dev.get_func("htm_sp_topk", "sp_topk_select").unwrap();
+        let learn_fn = self.dev.get_func("htm_sp_learn", "sp_learn").unwrap();
+        let duty_fn = self.dev.get_func("htm_sp_duty", "sp_duty_update").unwrap();
+
+        let overlap_cfg = LaunchConfig {
+            grid_dim: (n as u32, 1, 1),
+            block_dim: (128, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let topk_cfg = LaunchConfig {
+            grid_dim: (1, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
+        };
+        let learn_cfg = overlap_cfg;
+        let duty_cfg = LaunchConfig {
+            grid_dim: ((n as u32 + 255) / 256, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let alpha = 1.0f32 / self.duty_period.max(1.0);
+
+        // Reusable host buffer for the per-step active_mask D2H.
+        self.host_mask.resize(n, 0);
+
+        active_indices_host.clear();
+
+        for ti in 0..t {
+            // Point overlap kernel at the ti-th slice of the pre-uploaded input.
+            // cudarc CudaSlice doesn't have a "view" per se, so we must copy the
+            // slice into the reusable inp_dev buffer. This is a D2D copy — much
+            // faster than H2D.
+            // (Alternative: rewrite kernel to accept an offset; deferred.)
+            let in_off = ti * input_bits;
+            // Use dtod_copy via raw slice indexing: cudarc exposes slice() for this.
+            let sub = inputs_flat_dev.slice(in_off..in_off + input_bits);
+            self.dev.dtod_copy(&sub, &mut self.inp_dev)?;
+
+            // 1. sp_overlap
+            unsafe {
+                overlap_fn.clone().launch(
+                    overlap_cfg,
+                    (
+                        &self.inp_dev,
+                        &self.syn_bit,
+                        &self.syn_perm,
+                        &self.boost,
+                        self.conn_thr,
+                        self.synapses_per_col as u32,
+                        n as u32,
+                        &mut self.raw,
+                        &mut self.boosted,
+                    ),
+                )?;
+            }
+
+            // 2. Clear active_mask, then sp_topk
+            self.dev.memset_zeros(&mut self.active_mask)?;
+            unsafe {
+                topk_fn.clone().launch(
+                    topk_cfg,
+                    (&self.boosted, n as u32, k as u32, &mut self.active_mask),
+                )?;
+            }
+
+            // 3. sp_learn
+            if learn {
+                unsafe {
+                    learn_fn.clone().launch(
+                        learn_cfg,
+                        (
+                            &self.active_mask,
+                            &self.inp_dev,
+                            &self.syn_bit,
+                            &mut self.syn_perm,
+                            self.inc,
+                            self.dec,
+                            self.synapses_per_col as u32,
+                            n as u32,
+                        ),
+                    )?;
+                }
+            }
+
+            // 4. duty update (device)
+            unsafe {
+                duty_fn.clone().launch(
+                    duty_cfg,
+                    (
+                        &self.active_mask,
+                        &self.raw,
+                        &mut self.active_duty,
+                        &mut self.overlap_duty,
+                        &mut self.boost,
+                        alpha,
+                        1.0f32,
+                        0.0f32,
+                        0.0f32,
+                        0u32,
+                        n as u32,
+                    ),
+                )?;
+            }
+
+            // 5. Boost update. Two modes:
+            //    * strict_parity (tests): host-side exp for bit-exact match.
+            //    * default (production): GPU expf is close enough and ~10x faster
+            //      since we skip the D2H/H2D round-trip.
+            if learn && self.boost_strength > 0.0 {
+                if self.strict_parity {
+                    let mut duty_host = vec![0f32; n];
+                    self.dev
+                        .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
+                    let sum: f32 = duty_host.iter().sum();
+                    let mean = sum / (n as f32);
+                    let mut boost_host = vec![0f32; n];
+                    for i in 0..n {
+                        boost_host[i] =
+                            (-self.boost_strength * (duty_host[i] - mean)).exp();
+                    }
+                    self.dev.htod_sync_copy_into(&boost_host, &mut self.boost)?;
+
+                    // Permanence bump (rare). Only evaluated in strict mode.
+                    let mut ov_host = vec![0f32; n];
+                    self.dev
+                        .dtoh_sync_copy_into(&self.overlap_duty, &mut ov_host)?;
+                    let max_ov = ov_host.iter().cloned().fold(0f32, f32::max);
+                    if max_ov > 0.0 {
+                        let thr = 0.001f32 * max_ov;
+                        let bump = self.inc * 0.1f32;
+                        let bump_cols: Vec<u32> = ov_host
+                            .iter()
+                            .enumerate()
+                            .filter_map(|(i, &o)| {
+                                if o < thr { Some(i as u32) } else { None }
+                            })
+                            .collect();
+                        if !bump_cols.is_empty() {
+                            let s = self.synapses_per_col;
+                            let mut perm_host = vec![0f32; n * s];
+                            self.dev
+                                .dtoh_sync_copy_into(&self.syn_perm, &mut perm_host)?;
+                            for &c in &bump_cols {
+                                let base = (c as usize) * s;
+                                for p in &mut perm_host[base..base + s] {
+                                    *p = (*p + bump).min(1.0);
+                                }
+                            }
+                            self.dev.htod_sync_copy_into(&perm_host, &mut self.syn_perm)?;
+                        }
+                    }
+                } else {
+                    // Fast path: fused mean + boost = expf(-strength*(ad-mean))
+                    // in a single GPU block. Zero D2H, zero H2D — fully async.
+                    let boost_fn = self
+                        .dev
+                        .get_func("htm_sp_boost_fused", "sp_boost_from_duty")
+                        .expect("sp_boost_fused not loaded");
+                    let boost_cfg = LaunchConfig {
+                        grid_dim: (1, 1, 1),
+                        block_dim: (1024, 1, 1),
+                        shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
+                    };
+                    unsafe {
+                        boost_fn.launch(
+                            boost_cfg,
+                            (
+                                &self.active_duty,
+                                &mut self.boost,
+                                self.boost_strength,
+                                n as u32,
+                            ),
+                        )?;
+                    }
+                }
+            }
+
+            // D2H the active_mask for this step. This is the single
+            // unavoidable sync point per step — CPU TM needs the active
+            // indices for its next state update. At 2048 bytes / step this
+            // is tiny in bandwidth but costs a full syncronize (~5-10μs).
+            self.dev
+                .dtoh_sync_copy_into(&self.active_mask, &mut self.host_mask)?;
+            let co = ti * n;
+            cols_out[co..co + n].copy_from_slice(&self.host_mask);
+            // Extract active indices.
+            for (i, &b) in self.host_mask.iter().enumerate() {
+                if b != 0 {
+                    active_indices_host.push(i as u32);
+                }
+            }
+            // Insert separator (u32::MAX) between steps to demarcate step boundaries.
+            active_indices_host.push(u32::MAX);
+        }
+
+        Ok(())
+    }
+
+    /// Fully-on-GPU batched SP + TM. Zero per-step host sync.
+    ///
+    /// Inputs:
+    ///   inputs_flat_dev : (T * input_bits) u8 already uploaded
+    ///   cols_dev        : (T * n_cols) u8 output — active-column mask per step
+    ///   anom_dev        : (T,) f32 output — anomaly score per step
+    ///   tm              : persistent GPU TemporalMemory for this region
+    #[allow(clippy::too_many_arguments)]
+    pub fn step_batch_with_tm(
+        &mut self,
+        inputs_flat_dev: &CudaSlice<u8>,
+        t: usize,
+        input_bits: usize,
+        learn: bool,
+        cols_dev: &mut CudaSlice<u8>,
+        anom_dev: &mut CudaSlice<f32>,
+        tm: &mut crate::gpu::tm_gpu::TemporalMemoryGpu,
+    ) -> Result<(), DriverError> {
+        let n = self.n_columns;
+        let k = ((self.sparsity * n as f32).round() as usize).max(1);
+        debug_assert_eq!(cols_dev.len(), t * n);
+        debug_assert_eq!(anom_dev.len(), t);
+
+        let overlap_fn = self.dev.get_func("htm_sp_overlap", "sp_overlap").unwrap();
+        let topk_fn    = self.dev.get_func("htm_sp_topk", "sp_topk_select").unwrap();
+        let learn_fn   = self.dev.get_func("htm_sp_learn", "sp_learn").unwrap();
+        let duty_fn    = self.dev.get_func("htm_sp_duty", "sp_duty_update").unwrap();
+
+        let overlap_cfg = LaunchConfig {
+            grid_dim: (n as u32, 1, 1),
+            block_dim: (128, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let topk_cfg = LaunchConfig {
+            grid_dim: (1, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
+        };
+        let learn_cfg = overlap_cfg;
+        let duty_cfg = LaunchConfig {
+            grid_dim: ((n as u32 + 255) / 256, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let alpha = 1.0f32 / self.duty_period.max(1.0);
+
+        for ti in 0..t {
+            let in_off = ti * input_bits;
+            let sub = inputs_flat_dev.slice(in_off..in_off + input_bits);
+            self.dev.dtod_copy(&sub, &mut self.inp_dev)?;
+
+            // 1. sp_overlap
+            unsafe {
+                overlap_fn.clone().launch(
+                    overlap_cfg,
+                    (
+                        &self.inp_dev,
+                        &self.syn_bit,
+                        &self.syn_perm,
+                        &self.boost,
+                        self.conn_thr,
+                        self.synapses_per_col as u32,
+                        n as u32,
+                        &mut self.raw,
+                        &mut self.boosted,
+                    ),
+                )?;
+            }
+
+            // 2. clear + sp_topk
+            self.dev.memset_zeros(&mut self.active_mask)?;
+            unsafe {
+                topk_fn.clone().launch(
+                    topk_cfg,
+                    (&self.boosted, n as u32, k as u32, &mut self.active_mask),
+                )?;
+            }
+
+            // 3. sp_learn
+            if learn {
+                unsafe {
+                    learn_fn.clone().launch(
+                        learn_cfg,
+                        (
+                            &self.active_mask,
+                            &self.inp_dev,
+                            &self.syn_bit,
+                            &mut self.syn_perm,
+                            self.inc,
+                            self.dec,
+                            self.synapses_per_col as u32,
+                            n as u32,
+                        ),
+                    )?;
+                }
+            }
+
+            // 4. duty update (stage 1: no-boost write)
+            unsafe {
+                duty_fn.clone().launch(
+                    duty_cfg,
+                    (
+                        &self.active_mask,
+                        &self.raw,
+                        &mut self.active_duty,
+                        &mut self.overlap_duty,
+                        &mut self.boost,
+                        alpha,
+                        1.0f32,
+                        0.0f32,
+                        0.0f32,
+                        0u32,
+                        n as u32,
+                    ),
+                )?;
+            }
+
+            // 5. Boost update: fused GPU kernel (no D2H).
+            if learn && self.boost_strength > 0.0 {
+                let boost_fn = self.dev
+                    .get_func("htm_sp_boost_fused", "sp_boost_from_duty")
+                    .expect("sp_boost_fused not loaded");
+                let boost_cfg = LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (1024, 1, 1),
+                    shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
+                };
+                unsafe {
+                    boost_fn.launch(
+                        boost_cfg,
+                        (
+                            &self.active_duty,
+                            &mut self.boost,
+                            self.boost_strength,
+                            n as u32,
+                        ),
+                    )?;
+                }
+            }
+
+            // 6. Copy active_mask slice into cols_dev[ti*n .. (ti+1)*n].
+            let mut dst_slice = cols_dev.slice_mut(ti * n..(ti + 1) * n);
+            self.dev.dtod_copy(&self.active_mask, &mut dst_slice)?;
+
+            // 7. GPU TM step: predict + activate + anomaly + learn, all on device.
+            tm.step(&self.active_mask, anom_dev, ti as u32, learn)?;
+        }
+
+        Ok(())
+    }
+
+    /// One SP step on the GPU. Returns sorted active-column indices.
+    pub fn compute(&mut self, input: &[u8], learn: bool) -> Result<Vec<u32>, DriverError> {
+        debug_assert_eq!(input.len(), self.input_bits);
+        let n = self.n_columns;
+        let k = ((self.sparsity * n as f32).round() as usize).max(1);
+
+        // 1. H2D input SDR.
+        self.dev.htod_sync_copy_into(input, &mut self.inp_dev)?;
+
+        // 2. Launch sp_overlap: grid=n_columns, block=128.
+        let overlap_fn = self
+            .dev
+            .get_func("htm_sp_overlap", "sp_overlap")
+            .expect("sp_overlap not loaded");
+        let overlap_cfg = LaunchConfig {
+            grid_dim: (n as u32, 1, 1),
+            block_dim: (128, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        unsafe {
+            overlap_fn.launch(
+                overlap_cfg,
+                (
+                    &self.inp_dev,
+                    &self.syn_bit,
+                    &self.syn_perm,
+                    &self.boost,
+                    self.conn_thr,
+                    self.synapses_per_col as u32,
+                    n as u32,
+                    &mut self.raw,
+                    &mut self.boosted,
+                ),
+            )?;
+        }
+
+        // 3. Launch sp_topk: single block, shared mem = n_columns * f32.
+        let topk_fn = self
+            .dev
+            .get_func("htm_sp_topk", "sp_topk_select")
+            .expect("sp_topk not loaded");
+        let topk_cfg = LaunchConfig {
+            grid_dim: (1, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: (n * std::mem::size_of::<f32>()) as u32,
+        };
+        // Clear active_mask first. memset_zeros avoids an H2D of a host
+        // zeroes vector every step.
+        self.dev.memset_zeros(&mut self.active_mask)?;
+        unsafe {
+            topk_fn.launch(
+                topk_cfg,
+                (
+                    &self.boosted,
+                    n as u32,
+                    k as u32,
+                    &mut self.active_mask,
+                ),
+            )?;
+        }
+
+        // 4. Optional: sp_learn on active columns.
+        if learn {
+            let learn_fn = self
+                .dev
+                .get_func("htm_sp_learn", "sp_learn")
+                .expect("sp_learn not loaded");
+            let learn_cfg = LaunchConfig {
+                grid_dim: (n as u32, 1, 1),
+                block_dim: (128, 1, 1),
+                shared_mem_bytes: 0,
+            };
+            unsafe {
+                learn_fn.launch(
+                    learn_cfg,
+                    (
+                        &self.active_mask,
+                        &self.inp_dev,
+                        &self.syn_bit,
+                        &mut self.syn_perm,
+                        self.inc,
+                        self.dec,
+                        self.synapses_per_col as u32,
+                        n as u32,
+                    ),
+                )?;
+            }
+        }
+
+        // 5. Duty cycle + boost update. Always runs (matches CPU).
+        //    We need mean_duty on the host — compute BEFORE the update (matches
+        //    CPU sp.rs line 200-205 where mean is computed then written).
+        //    Actually CPU computes mean of the PRE-update duty cycles too? Re-read:
+        //      sp.rs lines 186-196 update duty cycles (pre-mean).
+        //      Line 202: mean = sum(active_duty_cycle) / n  ← after update.
+        //      Line 204: boost[i] = exp(-strength*(active_duty[i] - mean)).
+        //    So mean is on POST-update values.
+        //    Easiest: 1) run duty update with boost_strength=0 (skip boost calc),
+        //             2) D2H active_duty, compute mean, 3) run a boost-only kernel
+        //                OR inline the exp() in a second launch with mean passed.
+        //
+        //    For simplicity and correctness we fuse: run the duty kernel with
+        //    mean=0 and boost_strength=0 (disables boost write), then D2H to
+        //    compute mean, then re-launch with the true mean. Two launches, one
+        //    tiny D2H (n × f32). At n=2048 this is 8KB per step — negligible.
+        let alpha = 1.0f32 / self.duty_period.max(1.0);
+        let duty_fn = self
+            .dev
+            .get_func("htm_sp_duty", "sp_duty_update")
+            .expect("sp_duty not loaded");
+        let duty_cfg = LaunchConfig {
+            grid_dim: ((n as u32 + 255) / 256, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        // Stage 1: update duty cycles (boost_strength=0 -> no write).
+        unsafe {
+            duty_fn.launch(
+                duty_cfg,
+                (
+                    &self.active_mask,
+                    &self.raw,
+                    &mut self.active_duty,
+                    &mut self.overlap_duty,
+                    &mut self.boost,
+                    alpha,
+                    1.0f32,   // stim_thr
+                    0.0f32,   // boost_strength = 0 -> skip write
+                    0.0f32,   // mean_duty (unused)
+                    0u32,     // learn_flag = 0
+                    n as u32,
+                ),
+            )?;
+        }
+
+        if learn && self.boost_strength > 0.0 && self.strict_parity {
+            // Boost update must bit-match CPU `f32::exp`, so we compute it on
+            // the host and copy back. Cost per step: 8KB D2H + 8KB H2D at n=2048.
+            // Critical for learning parity — CUDA expf (even without fast-math)
+            // uses different rounding for some inputs than host libm.
+            let mut duty_host = vec![0f32; n];
+            self.dev
+                .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
+            let sum: f32 = duty_host.iter().sum();
+            let mean = sum / (n as f32);
+            let mut boost_host = vec![0f32; n];
+            for i in 0..n {
+                boost_host[i] = (-self.boost_strength * (duty_host[i] - mean)).exp();
+            }
+            self.dev.htod_sync_copy_into(&boost_host, &mut self.boost)?;
+
+            // CPU sp.rs 210-226: permanence bump for chronically under-stimulated
+            // columns. If overlap_duty_cycle[i] < 0.001 * max(overlap_duty_cycle),
+            // add inc*0.1 to every synapse of column i (clamped to 1.0).
+            // This runs only once per step and only for the rare cases, but we
+            // need it for bit-exact parity with CPU learn.
+            let mut ov_host = vec![0f32; n];
+            self.dev
+                .dtoh_sync_copy_into(&self.overlap_duty, &mut ov_host)?;
+            let max_ov = ov_host.iter().cloned().fold(0f32, f32::max);
+            if max_ov > 0.0 {
+                let thr = 0.001f32 * max_ov;
+                let bump = self.inc * 0.1f32;
+                // Find columns needing a bump. Usually empty. Rare → D2H/H2D
+                // of syn_perm is cheap (n*S*4 = 320KB at n=2048,S=40).
+                let bump_cols: Vec<u32> = ov_host
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(i, &o)| if o < thr { Some(i as u32) } else { None })
+                    .collect();
+                if !bump_cols.is_empty() {
+                    // Download, bump, upload. (Keeps implementation simple and
+                    // bit-exact. Could kernelize later.)
+                    let s = self.synapses_per_col;
+                    let mut perm_host = vec![0f32; n * s];
+                    self.dev.dtoh_sync_copy_into(&self.syn_perm, &mut perm_host)?;
+                    for &c in &bump_cols {
+                        let base = (c as usize) * s;
+                        for p in &mut perm_host[base..base + s] {
+                            *p = (*p + bump).min(1.0);
+                        }
+                    }
+                    self.dev.htod_sync_copy_into(&perm_host, &mut self.syn_perm)?;
+                }
+            }
+        } else if learn && self.boost_strength > 0.0 {
+            // Fast path: GPU-side boost using the already-loaded duty kernel.
+            let mut duty_host = vec![0f32; n];
+            self.dev
+                .dtoh_sync_copy_into(&self.active_duty, &mut duty_host)?;
+            let sum: f32 = duty_host.iter().sum();
+            let mean = sum / (n as f32);
+            let boost_fn = self
+                .dev
+                .get_func("htm_sp_duty", "sp_duty_update")
+                .expect("sp_duty not loaded");
+            unsafe {
+                boost_fn.launch(
+                    duty_cfg,
+                    (
+                        &self.active_mask,
+                        &self.raw,
+                        &mut self.active_duty,
+                        &mut self.overlap_duty,
+                        &mut self.boost,
+                        0.0f32,
+                        1.0f32,
+                        self.boost_strength,
+                        mean,
+                        1u32,
+                        n as u32,
+                    ),
+                )?;
+            }
+        }
+
+        // 6. D2H active_mask and convert to sorted index list.
+        self.dev
+            .dtoh_sync_copy_into(&self.active_mask, &mut self.host_mask)?;
+        let mut active: Vec<u32> = Vec::with_capacity(k);
+        for (i, &b) in self.host_mask.iter().enumerate() {
+            if b != 0 {
+                active.push(i as u32);
+            }
+        }
+        debug_assert_eq!(active.len(), k, "SP must emit exactly k winners");
+        Ok(active)
+    }
+}

overlay/htm_rust/src/gpu/tests.rs

@@ -1,663 +1,663 @@
-//! Parity tests: GPU SP vs CPU SP reference.
-//!
-//! With matching seeds the two should produce bit-identical active-column sets
-//! when `learn=false`, and remain bit-identical over repeated `learn=true`
-//! steps because the Hebbian update is deterministic (no RNG once initialised).
-//!
-//! Run with:  cargo test --release --features gpu
-
-#![cfg(test)]
-#![cfg(feature = "gpu")]
-
-use crate::sp::{SpatialPooler, SpatialPoolerConfig};
-use crate::gpu::sp_gpu::SpatialPoolerGpu;
-use crate::gpu::tm_gpu::TemporalMemoryGpu;
-use crate::gpu::fused::{
-    launch_fused, plan_batched_grid_dim, plan_fused_launch, FusedState,
-};
-use cudarc::driver::CudaSlice;
-use rand::{Rng, SeedableRng};
-use rand_xoshiro::Xoshiro256PlusPlus;
-
-fn make_sdr(rng: &mut Xoshiro256PlusPlus, bits: usize, sparsity: f32) -> Vec<u8> {
-    let on = ((sparsity * bits as f32) as usize).max(1);
-    let mut v = vec![0u8; bits];
-    let mut placed = 0;
-    while placed < on {
-        let i = rng.gen_range(0..bits);
-        if v[i] == 0 {
-            v[i] = 1;
-            placed += 1;
-        }
-    }
-    v
-}
-
-#[test]
-fn gpu_sp_matches_cpu_no_learn() {
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let mut cpu = SpatialPooler::new(
-        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
-        1234,
-    );
-    let cpu_for_gpu = SpatialPooler::new(
-        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
-        1234,
-    );
-    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu_for_gpu)
-        .expect("gpu init (CUDA device available)");
-    gpu.set_strict_parity(true);
-
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(99);
-    for step in 0..20 {
-        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
-        let sdr_bool: Vec<bool> = sdr_u8.iter().map(|&x| x != 0).collect();
-
-        let cpu_active: Vec<u32> = cpu.compute(&sdr_bool, false);
-        let gpu_active: Vec<u32> = gpu.compute(&sdr_u8, false).expect("gpu compute");
-
-        assert_eq!(
-            cpu_active, gpu_active,
-            "mismatch at step {step}: len cpu={} gpu={}",
-            cpu_active.len(), gpu_active.len()
-        );
-    }
-}
-
-#[test]
-fn gpu_sp_matches_cpu_with_learn() {
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let mut cpu = SpatialPooler::new(
-        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
-        5678,
-    );
-    let cpu_for_gpu = SpatialPooler::new(
-        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
-        5678,
-    );
-    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu_for_gpu).expect("gpu init");
-    gpu.set_strict_parity(true);
-
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(42);
-    for step in 0..50 {
-        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
-        let sdr_bool: Vec<bool> = sdr_u8.iter().map(|&x| x != 0).collect();
-
-        let cpu_active = cpu.compute(&sdr_bool, true);
-        let gpu_active = gpu.compute(&sdr_u8, true).expect("gpu compute");
-
-        assert_eq!(
-            cpu_active, gpu_active,
-            "mismatch at step {step} with learning"
-        );
-    }
-}
-
-#[test]
-fn gpu_tm_anomaly_decays_on_repeating_sequence() {
-    // End-to-end GPU pipeline: SP feeds TM; repeating SDR sequence should drive
-    // anomaly down over time.
-    use crate::gpu::HTMRegionGpu;  // not pyclass methods; use internal constructor via Rust
-    // Easier: replicate the pipeline directly with SP + TM.
-
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let n_cols = cfg.n_columns;
-    let cells_per_col = 32usize;
-
-    let cpu_for_gpu = SpatialPooler::new(SpatialPoolerConfig::default(), 314);
-    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_for_gpu).expect("gpu init");
-    let dev = sp.dev_ref().clone();
-    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col)
-        .expect("gpu tm init");
-    tm.reset().expect("tm reset");
-
-    // Build 3 fixed SDRs, feed them in a repeating sequence.
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
-    let make = |rng: &mut Xoshiro256PlusPlus| make_sdr(rng, bits, 0.02);
-    let seqs = [make(&mut rng), make(&mut rng), make(&mut rng)];
-
-    // Warm up SP so columns are stable per symbol.
-    for _ in 0..100 {
-        for s in &seqs {
-            let _ = sp.compute(s, true).expect("sp compute");
-        }
-    }
-
-    // Build a long input buffer: 100 repetitions of [A,B,C] = 300 steps.
-    let repeats = 100usize;
-    let t = repeats * 3;
-    let mut inputs_flat = vec![0u8; t * bits];
-    for r in 0..repeats {
-        for (i, s) in seqs.iter().enumerate() {
-            let off = (r * 3 + i) * bits;
-            inputs_flat[off..off + bits].copy_from_slice(s);
-        }
-    }
-    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs_flat).expect("htod");
-
-    let mut cols_dev = dev.alloc_zeros::<u8>(t * n_cols).expect("alloc cols");
-    let mut anom_dev = dev.alloc_zeros::<f32>(t).expect("alloc anom");
-
-    sp.step_batch_with_tm(
-        &inputs_dev,
-        t,
-        bits,
-        true,
-        &mut cols_dev,
-        &mut anom_dev,
-        &mut tm,
-    ).expect("step_batch_with_tm");
-
-    let anom: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
-    let cols: Vec<u8> = dev.dtoh_sync_copy(&cols_dev).expect("d2h cols");
-
-    // Active column count per step must equal k for every step.
-    let k = ((cfg.sparsity * n_cols as f32).round() as usize).max(1);
-    for ti in 0..t {
-        let step_slice = &cols[ti * n_cols..(ti + 1) * n_cols];
-        let n_on = step_slice.iter().filter(|&&b| b != 0).count();
-        assert_eq!(n_on, k, "step {ti} has {n_on} active cols, expected {k}");
-    }
-
-    // First repetition: anomaly should be near 1.0 (nothing predicted).
-    let early_avg: f32 = anom[3..9].iter().sum::<f32>() / 6.0;
-    // Last repetitions: anomaly should be noticeably lower.
-    let late_avg: f32 = anom[(t - 9)..t].iter().sum::<f32>() / 9.0;
-    eprintln!("gpu tm: early anomaly = {early_avg:.3}, late = {late_avg:.3}");
-    assert!(
-        late_avg < early_avg,
-        "GPU TM should reduce anomaly on repeating sequence: early={early_avg:.3}, late={late_avg:.3}"
-    );
-}
-
-/// Cluster-sync smoke test: verifies that the fused megakernel (which relies on
-/// hardware `cluster::sync()` / grid-barrier on H100/H200 Hopper) completes
-/// without deadlock when called with real HTM state, and that output shapes are
-/// sane (no NaN / Inf in anomaly scores, active-column count in plausible range).
-///
-/// This is an *integration* test, not a synthetic micro-benchmark: it exercises
-/// exactly the same `launch_fused` code path used in production, so any
-/// deadlock in the cooperative-grid or DLB barrier would surface here.
-///
-/// Skips gracefully (with an eprintln) when no GPU is available — the test
-/// binary returns exit-code 0 in that case so CI still passes.
-#[test]
-fn cluster_sync_smoke_test() {
-    // Build a tiny HTM region (1024 inputs, 256 columns, 4 cells/column).
-    // This keeps VRAM usage minimal while still exercising all kernel paths.
-    let input_bits  = 1024usize;
-    let n_columns   = 256usize;
-    let cells_per_col = 4usize;
-
-    // Probe cooperative launch attribute before doing any real work.
-    // CU_DEVICE_ATTRIBUTE_CLUSTER_LAUNCH = 223 (added in CUDA 11.8 for Hopper).
-    // cudarc exposes raw attribute querying; we check cooperative launch (98)
-    // as the guard — cluster launch is a superset and not separately probed
-    // here since cudarc doesn't expose attribute 223 symbolically yet.
-    // On pre-Hopper hardware the DLB barrier path is used instead and the
-    // test still validates no deadlock on that path.
-
-    let make_cfg = || SpatialPoolerConfig {
-        input_bits,
-        n_columns,
-        sparsity: 0.04,  // ~10 active cols out of 256
-        ..SpatialPoolerConfig::default()
-    };
-
-    let cpu_ref = SpatialPooler::new(make_cfg(), 42);
-
-    let mut sp = match SpatialPoolerGpu::from_cpu(&cpu_ref) {
-        Ok(sp) => sp,
-        Err(e) => {
-            eprintln!("[cluster_sync_smoke_test] No GPU available ({e:?}) — skipping");
-            return;
-        }
-    };
-
-    let dev = sp.dev_ref().clone();
-
-    // Check cooperative launch support; skip with a clear message if absent.
-    let cooperative_ok = matches!(
-        dev.attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH),
-        Ok(v) if v > 0
-    );
-    if !cooperative_ok {
-        eprintln!("[cluster_sync_smoke_test] CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH=0 — DLB path only, still running test");
-        // We continue — the DLB path is the production fallback and must not deadlock either.
-    }
-
-    let mut tm = match TemporalMemoryGpu::new(dev.clone(), n_columns, cells_per_col) {
-        Ok(tm) => tm,
-        Err(e) => {
-            eprintln!("[cluster_sync_smoke_test] TemporalMemoryGpu::new failed ({e:?}) — skipping");
-            return;
-        }
-    };
-    tm.reset().expect("tm reset");
-
-    let mut fused_st: FusedState = match FusedState::new(
-        dev.clone(),
-        n_columns,
-        cells_per_col,
-        sp.initial_threshold_estimate(),
-    ) {
-        Ok(f) => f,
-        Err(e) => {
-            eprintln!("[cluster_sync_smoke_test] FusedState::new failed ({e:?}) — skipping");
-            return;
-        }
-    };
-    fused_st.reset().expect("fused reset");
-
-    // Build T=4 timesteps of all-zero input SDRs.
-    let t = 4usize;
-    let inputs_flat = vec![0u8; t * input_bits];
-    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs_flat).expect("htod inputs");
-
-    let mut cols_dev  = dev.alloc_zeros::<u8>(t * n_columns).expect("alloc cols");
-    let mut anom_dev  = dev.alloc_zeros::<f32>(t).expect("alloc anom");
-
-    // Execute with a 2-second timeout guard via a thread. If the kernel
-    // deadlocks, the parent test process times out and the CI job reports
-    // failure — we can't cancel a live CUDA kernel from Rust, but the
-    // launch_fused call itself must return within this window on any sane GPU.
-    //
-    // We run the kernel inline (not in a separate thread) because CUDA contexts
-    // are not safely shareable across threads without explicit multi-threading
-    // setup. The 2-second bound is enforced implicitly: if the kernel deadlocks,
-    // the test binary will hang and the CI timeout (typically 5 min) will kill it.
-    // For local dev, the deadlock would be immediately obvious.
-
-    launch_fused(
-        &mut sp,
-        &mut tm,
-        &mut fused_st,
-        &inputs_dev,
-        &mut cols_dev,
-        &mut anom_dev,
-        t,
-        input_bits,
-        false, // learn=false for determinism
-    ).expect("launch_fused (cluster_sync_smoke_test): deadlock or CUDA error");
-
-    dev.synchronize().expect("device sync after launch_fused");
-
-    // --- Correctness assertions ---
-
-    let cols_host: Vec<u8>  = dev.dtoh_sync_copy(&cols_dev).expect("d2h cols");
-    let anom_host: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
-
-    // Output buffers must be exactly the right size.
-    assert_eq!(cols_host.len(), t * n_columns, "cols buffer size mismatch");
-    assert_eq!(anom_host.len(), t,             "anom buffer size mismatch");
-
-    // Anomaly scores must be finite (NaN/Inf indicates numerical blow-up).
-    for (i, &a) in anom_host.iter().enumerate() {
-        assert!(a.is_finite(), "anomaly[{i}] is not finite: {a}");
-        assert!(a >= 0.0 && a <= 1.0, "anomaly[{i}] out of [0,1]: {a}");
-    }
-
-    // Active-column count per step: threshold-based inhibition, so 0 is
-    // possible on cold start (before thresholds calibrate), but we assert
-    // <= n_columns to catch buffer overruns or completely wrong output.
-    for ti in 0..t {
-        let n_on = cols_host[ti * n_columns..(ti + 1) * n_columns]
-            .iter()
-            .filter(|&&b| b != 0)
-            .count();
-        assert!(
-            n_on <= n_columns,
-            "step {ti}: active columns {n_on} > n_columns {n_columns} (buffer overrun?)"
-        );
-    }
-
-    eprintln!(
-        "[cluster_sync_smoke_test] PASSED: T={t}, n_cols={n_columns}, \
-         input_bits={input_bits}, cooperative_supported={cooperative_ok}, \
-         anom={anom_host:?}"
-    );
-}
-
-/// Parity check: the CAI zero-copy path (`step_many_cuda`) must produce
-/// bit-identical outputs to the numpy H2D/D2H path (`step_batch_with_tm`),
-/// since the kernel pipeline is the same — only the I/O wrapping changes.
-/// We skip the PyO3 CAI dict plumbing here and test the underlying
-/// ManuallyDrop + upgrade_device_ptr pattern directly.
-#[test]
-fn gpu_cuda_vs_numpy_parity() {
-    use std::mem::ManuallyDrop;
-
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let n_cols = cfg.n_columns;
-    let cells_per_col = 32usize;
-
-    // Build two identical (SP, TM) pairs from the same seed.
-    let build = || -> (SpatialPoolerGpu, TemporalMemoryGpu) {
-        let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 271828);
-        let sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu init");
-        let dev = sp.dev_ref().clone();
-        let mut tm = TemporalMemoryGpu::new(dev, n_cols, cells_per_col).expect("tm init");
-        tm.reset().expect("tm reset");
-        (sp, tm)
-    };
-
-    // Deterministic SDR sequence.
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(31337);
-    let t = 32usize;
-    let mut inputs_flat = vec![0u8; t * bits];
-    for i in 0..t {
-        let sdr = make_sdr(&mut rng, bits, 0.02);
-        inputs_flat[i * bits..(i + 1) * bits].copy_from_slice(&sdr);
-    }
-
-    // ---- Path A: owned CudaSlice (numpy-equivalent path) ----
-    let (mut sp_a, mut tm_a) = build();
-    let dev_a = sp_a.dev_ref().clone();
-    let inputs_a: CudaSlice<u8> = dev_a.htod_sync_copy(&inputs_flat).expect("htod");
-    let mut cols_a = dev_a.alloc_zeros::<u8>(t * n_cols).expect("alloc cols_a");
-    let mut anom_a = dev_a.alloc_zeros::<f32>(t).expect("alloc anom_a");
-    sp_a.step_batch_with_tm(&inputs_a, t, bits, false, &mut cols_a, &mut anom_a, &mut tm_a)
-        .expect("owned step_batch_with_tm");
-    dev_a.synchronize().expect("sync a");
-    let cols_a_host: Vec<u8> = dev_a.dtoh_sync_copy(&cols_a).expect("d2h cols_a");
-    let anom_a_host: Vec<f32> = dev_a.dtoh_sync_copy(&anom_a).expect("d2h anom_a");
-
-    // ---- Path B: borrowed device pointers via upgrade_device_ptr ----
-    // We allocate fresh owned CudaSlices on a fresh device, then take their
-    // raw ptrs and re-wrap as ManuallyDrop borrowed views — mimicking what
-    // `step_many_cuda` does with torch-owned CUDA memory.
-    let (mut sp_b, mut tm_b) = build();
-    let dev_b = sp_b.dev_ref().clone();
-    let inputs_b_owned: CudaSlice<u8> = dev_b.htod_sync_copy(&inputs_flat).expect("htod");
-    let cols_b_owned = dev_b.alloc_zeros::<u8>(t * n_cols).expect("alloc cols_b");
-    let anom_b_owned = dev_b.alloc_zeros::<f32>(t).expect("alloc anom_b");
-
-    // Extract raw CUdeviceptrs (and leak the owners so their Drop doesn't free).
-    let inputs_ptr = inputs_b_owned.leak();
-    let cols_ptr = cols_b_owned.leak();
-    let anom_ptr = anom_b_owned.leak();
-
-    // Re-wrap as borrowed views.
-    let inputs_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<u8>(inputs_ptr, t * bits) });
-    let mut cols_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols) });
-    let mut anom_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<f32>(anom_ptr, t) });
-
-    sp_b.step_batch_with_tm(&inputs_b, t, bits, false, &mut cols_b, &mut anom_b, &mut tm_b)
-        .expect("borrowed step_batch_with_tm");
-    dev_b.synchronize().expect("sync b");
-    // `ManuallyDrop` doesn't auto-coerce to `&CudaSlice<T>` for the DevicePtr
-    // trait bound on `dtoh_sync_copy`; explicit deref.
-    let cols_b_host: Vec<u8> = dev_b.dtoh_sync_copy(&*cols_b).expect("d2h cols_b");
-    let anom_b_host: Vec<f32> = dev_b.dtoh_sync_copy(&*anom_b).expect("d2h anom_b");
-
-    // Re-own so Drop actually frees (we leaked above).
-    let _inputs_owned_again = unsafe { dev_b.upgrade_device_ptr::<u8>(inputs_ptr, t * bits) };
-    let _cols_owned_again = unsafe { dev_b.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols) };
-    let _anom_owned_again = unsafe { dev_b.upgrade_device_ptr::<f32>(anom_ptr, t) };
-
-    assert_eq!(cols_a_host, cols_b_host, "active-column mask diverges between numpy and CAI paths");
-    assert_eq!(anom_a_host.len(), anom_b_host.len());
-    for (i, (a, b)) in anom_a_host.iter().zip(anom_b_host.iter()).enumerate() {
-        // Anomaly is a pure division of integer counts — bit-exact expected.
-        assert!((a - b).abs() < 1e-7, "anomaly mismatch at step {i}: a={a} b={b}");
-    }
-}
-
-/// Fused kernel: threshold activation should converge to near target sparsity
-/// after a short warmup. Acceptance: mean activation rate per step lands in
-/// [0.3*target, 2.5*target] after 500-step warmup. Because the threshold
-/// starts conservative (=2.0) and the per-column adaptation rate is slow
-/// (0.001), we allow a generous band — the test asserts directional
-/// convergence toward the target, not tight matching.
-#[test]
-fn gpu_threshold_converges_to_sparsity() {
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let n_cols = cfg.n_columns;
-    let cells_per_col = 32usize;
-    let target = cfg.sparsity;  // 0.02 = 40 cols expected
-
-    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 111);
-    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
-    let dev = sp.dev_ref().clone();
-    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col).expect("tm init");
-    let mut fused = FusedState::new(
-        dev.clone(),
-        n_cols,
-        cells_per_col,
-        sp.initial_threshold_estimate(),
-    ).expect("fused init");
-    tm.reset().expect("tm reset");
-    fused.reset().expect("fused reset");
-
-    // Warmup: 1000 random 2%-sparse SDRs.
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(31337);
-    let t_warm = 1000usize;
-    let mut inputs = vec![0u8; t_warm * bits];
-    for ti in 0..t_warm {
-        let sdr = make_sdr(&mut rng, bits, 0.02);
-        inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
-    }
-    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs).expect("htod");
-    let mut cols_dev = dev.alloc_zeros::<u8>(t_warm * n_cols).expect("alloc cols");
-    let mut anom_dev = dev.alloc_zeros::<f32>(t_warm).expect("alloc anom");
-    launch_fused(
-        &mut sp, &mut tm, &mut fused,
-        &inputs_dev, &mut cols_dev, &mut anom_dev,
-        t_warm, bits, true,
-    ).expect("warmup launch");
-    dev.synchronize().expect("sync");
-
-    // Measurement pass: another 200 steps, measure mean activation.
-    let t_meas = 200usize;
-    let mut meas_inputs = vec![0u8; t_meas * bits];
-    for ti in 0..t_meas {
-        let sdr = make_sdr(&mut rng, bits, 0.02);
-        meas_inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
-    }
-    let meas_dev: CudaSlice<u8> = dev.htod_sync_copy(&meas_inputs).expect("htod meas");
-    let mut meas_cols = dev.alloc_zeros::<u8>(t_meas * n_cols).expect("alloc meas cols");
-    let mut meas_anom = dev.alloc_zeros::<f32>(t_meas).expect("alloc meas anom");
-    launch_fused(
-        &mut sp, &mut tm, &mut fused,
-        &meas_dev, &mut meas_cols, &mut meas_anom,
-        t_meas, bits, true,
-    ).expect("meas launch");
-    dev.synchronize().expect("sync meas");
-
-    let cols_host: Vec<u8> = dev.dtoh_sync_copy(&meas_cols).expect("d2h");
-    let mut step_counts = Vec::with_capacity(t_meas);
-    for ti in 0..t_meas {
-        let n_on = cols_host[ti*n_cols..(ti+1)*n_cols]
-            .iter().filter(|&&b| b != 0).count();
-        step_counts.push(n_on);
-    }
-    let mean_active: f64 = step_counts.iter().map(|&c| c as f64).sum::<f64>()
-        / (t_meas as f64);
-    let target_active = target as f64 * n_cols as f64;
-    eprintln!(
-        "threshold-activation convergence: mean_active/step = {mean_active:.1} \
-         (target = {target_active:.1})"
-    );
-    // Very generous band — we just want to confirm the threshold loop is
-    // functioning (not diverged to 0 or to all-active).
-    assert!(
-        mean_active >= 0.25 * target_active && mean_active <= 4.0 * target_active,
-        "mean active {mean_active:.1} outside [0.25x, 4x] of target {target_active:.1}"
-    );
-}
-
-/// Fused kernel: TM should learn a repeating sequence — anomaly decays.
-#[test]
-fn gpu_fused_tm_anomaly_decays_on_repeating_sequence() {
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let n_cols = cfg.n_columns;
-    let cells_per_col = 32usize;
-
-    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 271);
-    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
-    let dev = sp.dev_ref().clone();
-    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col).expect("tm init");
-    let mut fused = FusedState::new(
-        dev.clone(),
-        n_cols,
-        cells_per_col,
-        sp.initial_threshold_estimate(),
-    ).expect("fused init");
-    tm.reset().expect("tm reset");
-    fused.reset().expect("fused reset");
-
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
-    let make = |rng: &mut Xoshiro256PlusPlus| make_sdr(rng, bits, 0.02);
-    let seqs = [make(&mut rng), make(&mut rng), make(&mut rng)];
-
-    // Warmup SP threshold calibration with random SDRs first.
-    let warm = 300usize;
-    let mut warm_inputs = vec![0u8; warm * bits];
-    for ti in 0..warm {
-        let sdr = make_sdr(&mut rng, bits, 0.02);
-        warm_inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
-    }
-    let warm_dev: CudaSlice<u8> = dev.htod_sync_copy(&warm_inputs).expect("htod warm");
-    let mut warm_cols = dev.alloc_zeros::<u8>(warm * n_cols).expect("alloc warm cols");
-    let mut warm_anom = dev.alloc_zeros::<f32>(warm).expect("alloc warm anom");
-    launch_fused(
-        &mut sp, &mut tm, &mut fused,
-        &warm_dev, &mut warm_cols, &mut warm_anom,
-        warm, bits, true,
-    ).expect("warm launch");
-    dev.synchronize().expect("sync warm");
-
-    // Feed repeating A,B,C sequence for 100 reps.
-    let repeats = 100usize;
-    let t = repeats * 3;
-    let mut inputs = vec![0u8; t * bits];
-    for r in 0..repeats {
-        for (i, s) in seqs.iter().enumerate() {
-            let off = (r*3 + i) * bits;
-            inputs[off..off+bits].copy_from_slice(s);
-        }
-    }
-    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs).expect("htod rep");
-    let mut cols_dev = dev.alloc_zeros::<u8>(t * n_cols).expect("alloc rep cols");
-    let mut anom_dev = dev.alloc_zeros::<f32>(t).expect("alloc rep anom");
-    launch_fused(
-        &mut sp, &mut tm, &mut fused,
-        &inputs_dev, &mut cols_dev, &mut anom_dev,
-        t, bits, true,
-    ).expect("rep launch");
-    dev.synchronize().expect("sync rep");
-
-    let anom: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
-    let early_avg: f32 = anom[3..12].iter().sum::<f32>() / 9.0;
-    let late_avg: f32 = anom[(t-9)..t].iter().sum::<f32>() / 9.0;
-    eprintln!("fused TM anomaly: early={early_avg:.3} late={late_avg:.3}");
-    assert!(
-        late_avg < early_avg,
-        "anomaly must decay: early={early_avg:.3} late={late_avg:.3}"
-    );
-    assert!(
-        late_avg < 0.5,
-        "late anomaly must be < 0.5 (got {late_avg:.3})"
-    );
-}
-
-#[test]
-fn gpu_sp_yields_k_winners() {
-    let cfg = SpatialPoolerConfig::default();
-    let bits = cfg.input_bits;
-    let n = cfg.n_columns;
-    let expected_k = ((cfg.sparsity * n as f32).round() as usize).max(1);
-    let cpu = SpatialPooler::new(SpatialPoolerConfig::default(), 7);
-    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu).expect("gpu init");
-
-    let mut rng = Xoshiro256PlusPlus::seed_from_u64(1);
-    for _ in 0..10 {
-        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
-        let active = gpu.compute(&sdr_u8, false).expect("gpu compute");
-        assert_eq!(active.len(), expected_k);
-        // Ensure sorted + unique.
-        for w in active.windows(2) {
-            assert!(w[0] < w[1], "duplicate or out-of-order winner indices");
-        }
-    }
-}
-
-#[test]
-fn fused_launch_plan_uses_cooperative_grid_sync() {
-    let plan = plan_fused_launch(30, true, 30, None).expect("cooperative supported");
-    assert_eq!(plan.grid_dim_x, 16);
-    assert_eq!(plan.cooperative_grid_limit, 30);
-}
-
-#[test]
-fn fused_launch_plan_scales_to_big_gpu() {
-    // H200-like: 132 SMs, high cooperative_grid_limit. Cap still applies.
-    let plan = plan_fused_launch(132, true, 1000, None).expect("cooperative supported");
-    assert_eq!(plan.grid_dim_x, 16); // capped by default override
-    let plan = plan_fused_launch(132, true, 1000, Some(64)).expect("cooperative supported");
-    assert_eq!(plan.grid_dim_x, 64); // override raises the cap
-}
-
-#[test]
-fn fused_launch_plan_refuses_non_cooperative_devices() {
-    // The slow path was removed. Devices without cooperative launch fail fast.
-    let err = plan_fused_launch(30, false, 0, None).unwrap_err();
-    assert!(err.contains("cooperative launch"));
-}
-
-#[test]
-fn fused_grid_cap_env_override_is_honored() {
-    let cfg = SpatialPoolerConfig::default();
-    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 5252);
-    let sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
-    let dev = sp.dev_ref().clone();
-
-    unsafe { std::env::set_var("HTM_FUSED_GRID_CAP", "12"); }
-    let fused = FusedState::new(
-        dev.clone(),
-        cfg.n_columns,
-        32usize,
-        sp.initial_threshold_estimate(),
-    ).expect("fused init");
-    unsafe { std::env::remove_var("HTM_FUSED_GRID_CAP"); }
-
-    let sm_count = match dev.attribute(
-        cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
-    ) {
-        Ok(v) => v as u32,
-        Err(_) => 16u32,
-    };
-    let expected = sm_count.max(1).min(12);
-    assert_eq!(
-        fused.grid_dim_x,
-        expected,
-        "fused grid cap env override ignored: expected min(sm_count, 12) = {expected}, got {}",
-        fused.grid_dim_x,
-    );
-}
-
-#[test]
-fn batched_grid_plan_clamps_a10g_batch32_under_cooperative_limit() {
-    // A10G observed in HF Jobs: cooperative_grid_limit=400, B=32.
-    // grid_x=16 requests 512 cooperative blocks and fails; clamp to 12.
-    let grid_x = plan_batched_grid_dim(16, 400, 32, false).expect("fits after clamp");
-    assert_eq!(grid_x, 12);
-}
-
-#[test]
-fn batched_grid_plan_reports_oversized_batch() {
-    let err = plan_batched_grid_dim(16, 31, 32, false).unwrap_err();
-    assert!(err.contains("COOPERATIVE_LAUNCH_TOO_LARGE"));
-}
-
-#[test]
-fn batched_grid_plan_does_not_clamp_cluster_launches() {
-    let grid_x = plan_batched_grid_dim(16, 31, 32, true).expect("cluster path bypasses cooperative limit");
-    assert_eq!(grid_x, 16);
-}
+//! Parity tests: GPU SP vs CPU SP reference.
+//!
+//! With matching seeds the two should produce bit-identical active-column sets
+//! when `learn=false`, and remain bit-identical over repeated `learn=true`
+//! steps because the Hebbian update is deterministic (no RNG once initialised).
+//!
+//! Run with:  cargo test --release --features gpu
+
+#![cfg(test)]
+#![cfg(feature = "gpu")]
+
+use crate::sp::{SpatialPooler, SpatialPoolerConfig};
+use crate::gpu::sp_gpu::SpatialPoolerGpu;
+use crate::gpu::tm_gpu::TemporalMemoryGpu;
+use crate::gpu::fused::{
+    launch_fused, plan_batched_grid_dim, plan_fused_launch, FusedState,
+};
+use cudarc::driver::CudaSlice;
+use rand::{Rng, SeedableRng};
+use rand_xoshiro::Xoshiro256PlusPlus;
+
+fn make_sdr(rng: &mut Xoshiro256PlusPlus, bits: usize, sparsity: f32) -> Vec<u8> {
+    let on = ((sparsity * bits as f32) as usize).max(1);
+    let mut v = vec![0u8; bits];
+    let mut placed = 0;
+    while placed < on {
+        let i = rng.gen_range(0..bits);
+        if v[i] == 0 {
+            v[i] = 1;
+            placed += 1;
+        }
+    }
+    v
+}
+
+#[test]
+fn gpu_sp_matches_cpu_no_learn() {
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let mut cpu = SpatialPooler::new(
+        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
+        1234,
+    );
+    let cpu_for_gpu = SpatialPooler::new(
+        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
+        1234,
+    );
+    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu_for_gpu)
+        .expect("gpu init (CUDA device available)");
+    gpu.set_strict_parity(true);
+
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(99);
+    for step in 0..20 {
+        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
+        let sdr_bool: Vec<bool> = sdr_u8.iter().map(|&x| x != 0).collect();
+
+        let cpu_active: Vec<u32> = cpu.compute(&sdr_bool, false);
+        let gpu_active: Vec<u32> = gpu.compute(&sdr_u8, false).expect("gpu compute");
+
+        assert_eq!(
+            cpu_active, gpu_active,
+            "mismatch at step {step}: len cpu={} gpu={}",
+            cpu_active.len(), gpu_active.len()
+        );
+    }
+}
+
+#[test]
+fn gpu_sp_matches_cpu_with_learn() {
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let mut cpu = SpatialPooler::new(
+        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
+        5678,
+    );
+    let cpu_for_gpu = SpatialPooler::new(
+        SpatialPoolerConfig { ..SpatialPoolerConfig::default() },
+        5678,
+    );
+    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu_for_gpu).expect("gpu init");
+    gpu.set_strict_parity(true);
+
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(42);
+    for step in 0..50 {
+        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
+        let sdr_bool: Vec<bool> = sdr_u8.iter().map(|&x| x != 0).collect();
+
+        let cpu_active = cpu.compute(&sdr_bool, true);
+        let gpu_active = gpu.compute(&sdr_u8, true).expect("gpu compute");
+
+        assert_eq!(
+            cpu_active, gpu_active,
+            "mismatch at step {step} with learning"
+        );
+    }
+}
+
+#[test]
+fn gpu_tm_anomaly_decays_on_repeating_sequence() {
+    // End-to-end GPU pipeline: SP feeds TM; repeating SDR sequence should drive
+    // anomaly down over time.
+    use crate::gpu::HTMRegionGpu;  // not pyclass methods; use internal constructor via Rust
+    // Easier: replicate the pipeline directly with SP + TM.
+
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let n_cols = cfg.n_columns;
+    let cells_per_col = 32usize;
+
+    let cpu_for_gpu = SpatialPooler::new(SpatialPoolerConfig::default(), 314);
+    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_for_gpu).expect("gpu init");
+    let dev = sp.dev_ref().clone();
+    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col)
+        .expect("gpu tm init");
+    tm.reset().expect("tm reset");
+
+    // Build 3 fixed SDRs, feed them in a repeating sequence.
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
+    let make = |rng: &mut Xoshiro256PlusPlus| make_sdr(rng, bits, 0.02);
+    let seqs = [make(&mut rng), make(&mut rng), make(&mut rng)];
+
+    // Warm up SP so columns are stable per symbol.
+    for _ in 0..100 {
+        for s in &seqs {
+            let _ = sp.compute(s, true).expect("sp compute");
+        }
+    }
+
+    // Build a long input buffer: 100 repetitions of [A,B,C] = 300 steps.
+    let repeats = 100usize;
+    let t = repeats * 3;
+    let mut inputs_flat = vec![0u8; t * bits];
+    for r in 0..repeats {
+        for (i, s) in seqs.iter().enumerate() {
+            let off = (r * 3 + i) * bits;
+            inputs_flat[off..off + bits].copy_from_slice(s);
+        }
+    }
+    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs_flat).expect("htod");
+
+    let mut cols_dev = dev.alloc_zeros::<u8>(t * n_cols).expect("alloc cols");
+    let mut anom_dev = dev.alloc_zeros::<f32>(t).expect("alloc anom");
+
+    sp.step_batch_with_tm(
+        &inputs_dev,
+        t,
+        bits,
+        true,
+        &mut cols_dev,
+        &mut anom_dev,
+        &mut tm,
+    ).expect("step_batch_with_tm");
+
+    let anom: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
+    let cols: Vec<u8> = dev.dtoh_sync_copy(&cols_dev).expect("d2h cols");
+
+    // Active column count per step must equal k for every step.
+    let k = ((cfg.sparsity * n_cols as f32).round() as usize).max(1);
+    for ti in 0..t {
+        let step_slice = &cols[ti * n_cols..(ti + 1) * n_cols];
+        let n_on = step_slice.iter().filter(|&&b| b != 0).count();
+        assert_eq!(n_on, k, "step {ti} has {n_on} active cols, expected {k}");
+    }
+
+    // First repetition: anomaly should be near 1.0 (nothing predicted).
+    let early_avg: f32 = anom[3..9].iter().sum::<f32>() / 6.0;
+    // Last repetitions: anomaly should be noticeably lower.
+    let late_avg: f32 = anom[(t - 9)..t].iter().sum::<f32>() / 9.0;
+    eprintln!("gpu tm: early anomaly = {early_avg:.3}, late = {late_avg:.3}");
+    assert!(
+        late_avg < early_avg,
+        "GPU TM should reduce anomaly on repeating sequence: early={early_avg:.3}, late={late_avg:.3}"
+    );
+}
+
+/// Cluster-sync smoke test: verifies that the fused megakernel (which relies on
+/// hardware `cluster::sync()` / grid-barrier on H100/H200 Hopper) completes
+/// without deadlock when called with real HTM state, and that output shapes are
+/// sane (no NaN / Inf in anomaly scores, active-column count in plausible range).
+///
+/// This is an *integration* test, not a synthetic micro-benchmark: it exercises
+/// exactly the same `launch_fused` code path used in production, so any
+/// deadlock in the cooperative-grid or DLB barrier would surface here.
+///
+/// Skips gracefully (with an eprintln) when no GPU is available — the test
+/// binary returns exit-code 0 in that case so CI still passes.
+#[test]
+fn cluster_sync_smoke_test() {
+    // Build a tiny HTM region (1024 inputs, 256 columns, 4 cells/column).
+    // This keeps VRAM usage minimal while still exercising all kernel paths.
+    let input_bits  = 1024usize;
+    let n_columns   = 256usize;
+    let cells_per_col = 4usize;
+
+    // Probe cooperative launch attribute before doing any real work.
+    // CU_DEVICE_ATTRIBUTE_CLUSTER_LAUNCH = 223 (added in CUDA 11.8 for Hopper).
+    // cudarc exposes raw attribute querying; we check cooperative launch (98)
+    // as the guard — cluster launch is a superset and not separately probed
+    // here since cudarc doesn't expose attribute 223 symbolically yet.
+    // On pre-Hopper hardware the DLB barrier path is used instead and the
+    // test still validates no deadlock on that path.
+
+    let make_cfg = || SpatialPoolerConfig {
+        input_bits,
+        n_columns,
+        sparsity: 0.04,  // ~10 active cols out of 256
+        ..SpatialPoolerConfig::default()
+    };
+
+    let cpu_ref = SpatialPooler::new(make_cfg(), 42);
+
+    let mut sp = match SpatialPoolerGpu::from_cpu(&cpu_ref) {
+        Ok(sp) => sp,
+        Err(e) => {
+            eprintln!("[cluster_sync_smoke_test] No GPU available ({e:?}) — skipping");
+            return;
+        }
+    };
+
+    let dev = sp.dev_ref().clone();
+
+    // Check cooperative launch support; skip with a clear message if absent.
+    let cooperative_ok = matches!(
+        dev.attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH),
+        Ok(v) if v > 0
+    );
+    if !cooperative_ok {
+        eprintln!("[cluster_sync_smoke_test] CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH=0 — DLB path only, still running test");
+        // We continue — the DLB path is the production fallback and must not deadlock either.
+    }
+
+    let mut tm = match TemporalMemoryGpu::new(dev.clone(), n_columns, cells_per_col) {
+        Ok(tm) => tm,
+        Err(e) => {
+            eprintln!("[cluster_sync_smoke_test] TemporalMemoryGpu::new failed ({e:?}) — skipping");
+            return;
+        }
+    };
+    tm.reset().expect("tm reset");
+
+    let mut fused_st: FusedState = match FusedState::new(
+        dev.clone(),
+        n_columns,
+        cells_per_col,
+        sp.initial_threshold_estimate(),
+    ) {
+        Ok(f) => f,
+        Err(e) => {
+            eprintln!("[cluster_sync_smoke_test] FusedState::new failed ({e:?}) — skipping");
+            return;
+        }
+    };
+    fused_st.reset().expect("fused reset");
+
+    // Build T=4 timesteps of all-zero input SDRs.
+    let t = 4usize;
+    let inputs_flat = vec![0u8; t * input_bits];
+    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs_flat).expect("htod inputs");
+
+    let mut cols_dev  = dev.alloc_zeros::<u8>(t * n_columns).expect("alloc cols");
+    let mut anom_dev  = dev.alloc_zeros::<f32>(t).expect("alloc anom");
+
+    // Execute with a 2-second timeout guard via a thread. If the kernel
+    // deadlocks, the parent test process times out and the CI job reports
+    // failure — we can't cancel a live CUDA kernel from Rust, but the
+    // launch_fused call itself must return within this window on any sane GPU.
+    //
+    // We run the kernel inline (not in a separate thread) because CUDA contexts
+    // are not safely shareable across threads without explicit multi-threading
+    // setup. The 2-second bound is enforced implicitly: if the kernel deadlocks,
+    // the test binary will hang and the CI timeout (typically 5 min) will kill it.
+    // For local dev, the deadlock would be immediately obvious.
+
+    launch_fused(
+        &mut sp,
+        &mut tm,
+        &mut fused_st,
+        &inputs_dev,
+        &mut cols_dev,
+        &mut anom_dev,
+        t,
+        input_bits,
+        false, // learn=false for determinism
+    ).expect("launch_fused (cluster_sync_smoke_test): deadlock or CUDA error");
+
+    dev.synchronize().expect("device sync after launch_fused");
+
+    // --- Correctness assertions ---
+
+    let cols_host: Vec<u8>  = dev.dtoh_sync_copy(&cols_dev).expect("d2h cols");
+    let anom_host: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
+
+    // Output buffers must be exactly the right size.
+    assert_eq!(cols_host.len(), t * n_columns, "cols buffer size mismatch");
+    assert_eq!(anom_host.len(), t,             "anom buffer size mismatch");
+
+    // Anomaly scores must be finite (NaN/Inf indicates numerical blow-up).
+    for (i, &a) in anom_host.iter().enumerate() {
+        assert!(a.is_finite(), "anomaly[{i}] is not finite: {a}");
+        assert!(a >= 0.0 && a <= 1.0, "anomaly[{i}] out of [0,1]: {a}");
+    }
+
+    // Active-column count per step: threshold-based inhibition, so 0 is
+    // possible on cold start (before thresholds calibrate), but we assert
+    // <= n_columns to catch buffer overruns or completely wrong output.
+    for ti in 0..t {
+        let n_on = cols_host[ti * n_columns..(ti + 1) * n_columns]
+            .iter()
+            .filter(|&&b| b != 0)
+            .count();
+        assert!(
+            n_on <= n_columns,
+            "step {ti}: active columns {n_on} > n_columns {n_columns} (buffer overrun?)"
+        );
+    }
+
+    eprintln!(
+        "[cluster_sync_smoke_test] PASSED: T={t}, n_cols={n_columns}, \
+         input_bits={input_bits}, cooperative_supported={cooperative_ok}, \
+         anom={anom_host:?}"
+    );
+}
+
+/// Parity check: the CAI zero-copy path (`step_many_cuda`) must produce
+/// bit-identical outputs to the numpy H2D/D2H path (`step_batch_with_tm`),
+/// since the kernel pipeline is the same — only the I/O wrapping changes.
+/// We skip the PyO3 CAI dict plumbing here and test the underlying
+/// ManuallyDrop + upgrade_device_ptr pattern directly.
+#[test]
+fn gpu_cuda_vs_numpy_parity() {
+    use std::mem::ManuallyDrop;
+
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let n_cols = cfg.n_columns;
+    let cells_per_col = 32usize;
+
+    // Build two identical (SP, TM) pairs from the same seed.
+    let build = || -> (SpatialPoolerGpu, TemporalMemoryGpu) {
+        let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 271828);
+        let sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu init");
+        let dev = sp.dev_ref().clone();
+        let mut tm = TemporalMemoryGpu::new(dev, n_cols, cells_per_col).expect("tm init");
+        tm.reset().expect("tm reset");
+        (sp, tm)
+    };
+
+    // Deterministic SDR sequence.
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(31337);
+    let t = 32usize;
+    let mut inputs_flat = vec![0u8; t * bits];
+    for i in 0..t {
+        let sdr = make_sdr(&mut rng, bits, 0.02);
+        inputs_flat[i * bits..(i + 1) * bits].copy_from_slice(&sdr);
+    }
+
+    // ---- Path A: owned CudaSlice (numpy-equivalent path) ----
+    let (mut sp_a, mut tm_a) = build();
+    let dev_a = sp_a.dev_ref().clone();
+    let inputs_a: CudaSlice<u8> = dev_a.htod_sync_copy(&inputs_flat).expect("htod");
+    let mut cols_a = dev_a.alloc_zeros::<u8>(t * n_cols).expect("alloc cols_a");
+    let mut anom_a = dev_a.alloc_zeros::<f32>(t).expect("alloc anom_a");
+    sp_a.step_batch_with_tm(&inputs_a, t, bits, false, &mut cols_a, &mut anom_a, &mut tm_a)
+        .expect("owned step_batch_with_tm");
+    dev_a.synchronize().expect("sync a");
+    let cols_a_host: Vec<u8> = dev_a.dtoh_sync_copy(&cols_a).expect("d2h cols_a");
+    let anom_a_host: Vec<f32> = dev_a.dtoh_sync_copy(&anom_a).expect("d2h anom_a");
+
+    // ---- Path B: borrowed device pointers via upgrade_device_ptr ----
+    // We allocate fresh owned CudaSlices on a fresh device, then take their
+    // raw ptrs and re-wrap as ManuallyDrop borrowed views — mimicking what
+    // `step_many_cuda` does with torch-owned CUDA memory.
+    let (mut sp_b, mut tm_b) = build();
+    let dev_b = sp_b.dev_ref().clone();
+    let inputs_b_owned: CudaSlice<u8> = dev_b.htod_sync_copy(&inputs_flat).expect("htod");
+    let cols_b_owned = dev_b.alloc_zeros::<u8>(t * n_cols).expect("alloc cols_b");
+    let anom_b_owned = dev_b.alloc_zeros::<f32>(t).expect("alloc anom_b");
+
+    // Extract raw CUdeviceptrs (and leak the owners so their Drop doesn't free).
+    let inputs_ptr = inputs_b_owned.leak();
+    let cols_ptr = cols_b_owned.leak();
+    let anom_ptr = anom_b_owned.leak();
+
+    // Re-wrap as borrowed views.
+    let inputs_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<u8>(inputs_ptr, t * bits) });
+    let mut cols_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols) });
+    let mut anom_b = ManuallyDrop::new(unsafe { dev_b.upgrade_device_ptr::<f32>(anom_ptr, t) });
+
+    sp_b.step_batch_with_tm(&inputs_b, t, bits, false, &mut cols_b, &mut anom_b, &mut tm_b)
+        .expect("borrowed step_batch_with_tm");
+    dev_b.synchronize().expect("sync b");
+    // `ManuallyDrop` doesn't auto-coerce to `&CudaSlice<T>` for the DevicePtr
+    // trait bound on `dtoh_sync_copy`; explicit deref.
+    let cols_b_host: Vec<u8> = dev_b.dtoh_sync_copy(&*cols_b).expect("d2h cols_b");
+    let anom_b_host: Vec<f32> = dev_b.dtoh_sync_copy(&*anom_b).expect("d2h anom_b");
+
+    // Re-own so Drop actually frees (we leaked above).
+    let _inputs_owned_again = unsafe { dev_b.upgrade_device_ptr::<u8>(inputs_ptr, t * bits) };
+    let _cols_owned_again = unsafe { dev_b.upgrade_device_ptr::<u8>(cols_ptr, t * n_cols) };
+    let _anom_owned_again = unsafe { dev_b.upgrade_device_ptr::<f32>(anom_ptr, t) };
+
+    assert_eq!(cols_a_host, cols_b_host, "active-column mask diverges between numpy and CAI paths");
+    assert_eq!(anom_a_host.len(), anom_b_host.len());
+    for (i, (a, b)) in anom_a_host.iter().zip(anom_b_host.iter()).enumerate() {
+        // Anomaly is a pure division of integer counts — bit-exact expected.
+        assert!((a - b).abs() < 1e-7, "anomaly mismatch at step {i}: a={a} b={b}");
+    }
+}
+
+/// Fused kernel: threshold activation should converge to near target sparsity
+/// after a short warmup. Acceptance: mean activation rate per step lands in
+/// [0.3*target, 2.5*target] after 500-step warmup. Because the threshold
+/// starts conservative (=2.0) and the per-column adaptation rate is slow
+/// (0.001), we allow a generous band — the test asserts directional
+/// convergence toward the target, not tight matching.
+#[test]
+fn gpu_threshold_converges_to_sparsity() {
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let n_cols = cfg.n_columns;
+    let cells_per_col = 32usize;
+    let target = cfg.sparsity;  // 0.02 = 40 cols expected
+
+    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 111);
+    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
+    let dev = sp.dev_ref().clone();
+    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col).expect("tm init");
+    let mut fused = FusedState::new(
+        dev.clone(),
+        n_cols,
+        cells_per_col,
+        sp.initial_threshold_estimate(),
+    ).expect("fused init");
+    tm.reset().expect("tm reset");
+    fused.reset().expect("fused reset");
+
+    // Warmup: 1000 random 2%-sparse SDRs.
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(31337);
+    let t_warm = 1000usize;
+    let mut inputs = vec![0u8; t_warm * bits];
+    for ti in 0..t_warm {
+        let sdr = make_sdr(&mut rng, bits, 0.02);
+        inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
+    }
+    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs).expect("htod");
+    let mut cols_dev = dev.alloc_zeros::<u8>(t_warm * n_cols).expect("alloc cols");
+    let mut anom_dev = dev.alloc_zeros::<f32>(t_warm).expect("alloc anom");
+    launch_fused(
+        &mut sp, &mut tm, &mut fused,
+        &inputs_dev, &mut cols_dev, &mut anom_dev,
+        t_warm, bits, true,
+    ).expect("warmup launch");
+    dev.synchronize().expect("sync");
+
+    // Measurement pass: another 200 steps, measure mean activation.
+    let t_meas = 200usize;
+    let mut meas_inputs = vec![0u8; t_meas * bits];
+    for ti in 0..t_meas {
+        let sdr = make_sdr(&mut rng, bits, 0.02);
+        meas_inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
+    }
+    let meas_dev: CudaSlice<u8> = dev.htod_sync_copy(&meas_inputs).expect("htod meas");
+    let mut meas_cols = dev.alloc_zeros::<u8>(t_meas * n_cols).expect("alloc meas cols");
+    let mut meas_anom = dev.alloc_zeros::<f32>(t_meas).expect("alloc meas anom");
+    launch_fused(
+        &mut sp, &mut tm, &mut fused,
+        &meas_dev, &mut meas_cols, &mut meas_anom,
+        t_meas, bits, true,
+    ).expect("meas launch");
+    dev.synchronize().expect("sync meas");
+
+    let cols_host: Vec<u8> = dev.dtoh_sync_copy(&meas_cols).expect("d2h");
+    let mut step_counts = Vec::with_capacity(t_meas);
+    for ti in 0..t_meas {
+        let n_on = cols_host[ti*n_cols..(ti+1)*n_cols]
+            .iter().filter(|&&b| b != 0).count();
+        step_counts.push(n_on);
+    }
+    let mean_active: f64 = step_counts.iter().map(|&c| c as f64).sum::<f64>()
+        / (t_meas as f64);
+    let target_active = target as f64 * n_cols as f64;
+    eprintln!(
+        "threshold-activation convergence: mean_active/step = {mean_active:.1} \
+         (target = {target_active:.1})"
+    );
+    // Very generous band — we just want to confirm the threshold loop is
+    // functioning (not diverged to 0 or to all-active).
+    assert!(
+        mean_active >= 0.25 * target_active && mean_active <= 4.0 * target_active,
+        "mean active {mean_active:.1} outside [0.25x, 4x] of target {target_active:.1}"
+    );
+}
+
+/// Fused kernel: TM should learn a repeating sequence — anomaly decays.
+#[test]
+fn gpu_fused_tm_anomaly_decays_on_repeating_sequence() {
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let n_cols = cfg.n_columns;
+    let cells_per_col = 32usize;
+
+    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 271);
+    let mut sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
+    let dev = sp.dev_ref().clone();
+    let mut tm = TemporalMemoryGpu::new(dev.clone(), n_cols, cells_per_col).expect("tm init");
+    let mut fused = FusedState::new(
+        dev.clone(),
+        n_cols,
+        cells_per_col,
+        sp.initial_threshold_estimate(),
+    ).expect("fused init");
+    tm.reset().expect("tm reset");
+    fused.reset().expect("fused reset");
+
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
+    let make = |rng: &mut Xoshiro256PlusPlus| make_sdr(rng, bits, 0.02);
+    let seqs = [make(&mut rng), make(&mut rng), make(&mut rng)];
+
+    // Warmup SP threshold calibration with random SDRs first.
+    let warm = 300usize;
+    let mut warm_inputs = vec![0u8; warm * bits];
+    for ti in 0..warm {
+        let sdr = make_sdr(&mut rng, bits, 0.02);
+        warm_inputs[ti*bits..(ti+1)*bits].copy_from_slice(&sdr);
+    }
+    let warm_dev: CudaSlice<u8> = dev.htod_sync_copy(&warm_inputs).expect("htod warm");
+    let mut warm_cols = dev.alloc_zeros::<u8>(warm * n_cols).expect("alloc warm cols");
+    let mut warm_anom = dev.alloc_zeros::<f32>(warm).expect("alloc warm anom");
+    launch_fused(
+        &mut sp, &mut tm, &mut fused,
+        &warm_dev, &mut warm_cols, &mut warm_anom,
+        warm, bits, true,
+    ).expect("warm launch");
+    dev.synchronize().expect("sync warm");
+
+    // Feed repeating A,B,C sequence for 100 reps.
+    let repeats = 100usize;
+    let t = repeats * 3;
+    let mut inputs = vec![0u8; t * bits];
+    for r in 0..repeats {
+        for (i, s) in seqs.iter().enumerate() {
+            let off = (r*3 + i) * bits;
+            inputs[off..off+bits].copy_from_slice(s);
+        }
+    }
+    let inputs_dev: CudaSlice<u8> = dev.htod_sync_copy(&inputs).expect("htod rep");
+    let mut cols_dev = dev.alloc_zeros::<u8>(t * n_cols).expect("alloc rep cols");
+    let mut anom_dev = dev.alloc_zeros::<f32>(t).expect("alloc rep anom");
+    launch_fused(
+        &mut sp, &mut tm, &mut fused,
+        &inputs_dev, &mut cols_dev, &mut anom_dev,
+        t, bits, true,
+    ).expect("rep launch");
+    dev.synchronize().expect("sync rep");
+
+    let anom: Vec<f32> = dev.dtoh_sync_copy(&anom_dev).expect("d2h anom");
+    let early_avg: f32 = anom[3..12].iter().sum::<f32>() / 9.0;
+    let late_avg: f32 = anom[(t-9)..t].iter().sum::<f32>() / 9.0;
+    eprintln!("fused TM anomaly: early={early_avg:.3} late={late_avg:.3}");
+    assert!(
+        late_avg < early_avg,
+        "anomaly must decay: early={early_avg:.3} late={late_avg:.3}"
+    );
+    assert!(
+        late_avg < 0.5,
+        "late anomaly must be < 0.5 (got {late_avg:.3})"
+    );
+}
+
+#[test]
+fn gpu_sp_yields_k_winners() {
+    let cfg = SpatialPoolerConfig::default();
+    let bits = cfg.input_bits;
+    let n = cfg.n_columns;
+    let expected_k = ((cfg.sparsity * n as f32).round() as usize).max(1);
+    let cpu = SpatialPooler::new(SpatialPoolerConfig::default(), 7);
+    let mut gpu = SpatialPoolerGpu::from_cpu(&cpu).expect("gpu init");
+
+    let mut rng = Xoshiro256PlusPlus::seed_from_u64(1);
+    for _ in 0..10 {
+        let sdr_u8 = make_sdr(&mut rng, bits, 0.02);
+        let active = gpu.compute(&sdr_u8, false).expect("gpu compute");
+        assert_eq!(active.len(), expected_k);
+        // Ensure sorted + unique.
+        for w in active.windows(2) {
+            assert!(w[0] < w[1], "duplicate or out-of-order winner indices");
+        }
+    }
+}
+
+#[test]
+fn fused_launch_plan_uses_cooperative_grid_sync() {
+    let plan = plan_fused_launch(30, true, 30, None).expect("cooperative supported");
+    assert_eq!(plan.grid_dim_x, 16);
+    assert_eq!(plan.cooperative_grid_limit, 30);
+}
+
+#[test]
+fn fused_launch_plan_scales_to_big_gpu() {
+    // H200-like: 132 SMs, high cooperative_grid_limit. Cap still applies.
+    let plan = plan_fused_launch(132, true, 1000, None).expect("cooperative supported");
+    assert_eq!(plan.grid_dim_x, 16); // capped by default override
+    let plan = plan_fused_launch(132, true, 1000, Some(64)).expect("cooperative supported");
+    assert_eq!(plan.grid_dim_x, 64); // override raises the cap
+}
+
+#[test]
+fn fused_launch_plan_refuses_non_cooperative_devices() {
+    // The slow path was removed. Devices without cooperative launch fail fast.
+    let err = plan_fused_launch(30, false, 0, None).unwrap_err();
+    assert!(err.contains("cooperative launch"));
+}
+
+#[test]
+fn fused_grid_cap_env_override_is_honored() {
+    let cfg = SpatialPoolerConfig::default();
+    let cpu_ref = SpatialPooler::new(SpatialPoolerConfig::default(), 5252);
+    let sp = SpatialPoolerGpu::from_cpu(&cpu_ref).expect("gpu sp init");
+    let dev = sp.dev_ref().clone();
+
+    unsafe { std::env::set_var("HTM_FUSED_GRID_CAP", "12"); }
+    let fused = FusedState::new(
+        dev.clone(),
+        cfg.n_columns,
+        32usize,
+        sp.initial_threshold_estimate(),
+    ).expect("fused init");
+    unsafe { std::env::remove_var("HTM_FUSED_GRID_CAP"); }
+
+    let sm_count = match dev.attribute(
+        cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
+    ) {
+        Ok(v) => v as u32,
+        Err(_) => 16u32,
+    };
+    let expected = sm_count.max(1).min(12);
+    assert_eq!(
+        fused.grid_dim_x,
+        expected,
+        "fused grid cap env override ignored: expected min(sm_count, 12) = {expected}, got {}",
+        fused.grid_dim_x,
+    );
+}
+
+#[test]
+fn batched_grid_plan_clamps_a10g_batch32_under_cooperative_limit() {
+    // A10G observed in HF Jobs: cooperative_grid_limit=400, B=32.
+    // grid_x=16 requests 512 cooperative blocks and fails; clamp to 12.
+    let grid_x = plan_batched_grid_dim(16, 400, 32, false).expect("fits after clamp");
+    assert_eq!(grid_x, 12);
+}
+
+#[test]
+fn batched_grid_plan_reports_oversized_batch() {
+    let err = plan_batched_grid_dim(16, 31, 32, false).unwrap_err();
+    assert!(err.contains("COOPERATIVE_LAUNCH_TOO_LARGE"));
+}
+
+#[test]
+fn batched_grid_plan_does_not_clamp_cluster_launches() {
+    let grid_x = plan_batched_grid_dim(16, 31, 32, true).expect("cluster path bypasses cooperative limit");
+    assert_eq!(grid_x, 16);
+}

overlay/htm_rust/src/gpu/tm_gpu.rs

@@ -1,460 +1,460 @@
-//! GPU Temporal Memory.
-//!
-//! Flat device storage. Pre-allocated segment slab:
-//!   n_cells = n_columns * cells_per_column
-//!   n_segments_max = n_cells * MAX_SEGMENTS_PER_CELL
-//!   n_synapses_max = n_segments_max * MAX_SYN_PER_SEGMENT
-//!
-//! Defaults (CPU parity targets relaxed on GPU to keep memory tractable):
-//!   MAX_SEGMENTS_PER_CELL = 16
-//!   MAX_SYN_PER_SEGMENT   = 32
-//!
-//! At n_cells = 65536:
-//!   n_segments_max =  1_048_576   (~1M)
-//!   n_synapses_max = 33_554_432   (~33M)
-//! Storage:
-//!   syn_presyn : u32  × 33M = 128 MB
-//!   syn_perm   : i16  × 33M =  64 MB
-//!   seg_cell   : u32  ×  1M =   4 MB
-//!   seg_syn_n  : u32  ×  1M =   4 MB
-//!   misc bitsets etc        ~ <1 MB
-//!   -------------------------------
-//!   Total per region        ~200 MB
-//!
-//! Permanences are stored as i16 scaled by 32767 (→ [0, 32767] represents
-//! [0.0, 1.0]). inc/dec are provided pre-scaled.
-
-use std::sync::Arc;
-
-use cudarc::driver::{CudaDevice, CudaSlice, DriverError, DeviceRepr, LaunchAsync, LaunchConfig};
-use cudarc::nvrtc::Ptx;
-
-/// Packed config struct passed by value to TM kernels to stay under
-/// cudarc's 12-tuple launch limit. Layout must match the C-side
-/// `TmConfig` struct declared in each kernel.
-#[repr(C)]
-#[derive(Clone, Copy)]
-pub struct TmConfig {
-    pub activation_threshold: u32,
-    pub learning_threshold: u32,
-    pub cells_per_column: u32,
-    pub synapses_per_segment: u32,
-    pub n_segments: u32,
-    pub n_cells: u32,
-    pub max_segments_per_cell: u32,
-    pub max_new_synapses: u32,
-    pub conn_thr_i16: i32,        // i16 widened to i32 for alignment
-    pub perm_inc_i16: i32,
-    pub perm_dec_i16: i32,
-    pub predicted_seg_dec_i16: i32,
-    pub initial_perm_i16: i32,
-    pub iter_seed: u32,
-    pub n_cols: u32,
-    pub bits_words: u32,
-}
-
-unsafe impl DeviceRepr for TmConfig {}
-
-// Embedded PTX.
-const PTX_TM_PREDICT:  &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_predict.ptx"));
-const PTX_TM_ACTIVATE: &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_activate.ptx"));
-const PTX_TM_LEARN:    &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_learn.ptx"));
-const PTX_TM_PUNISH:   &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_punish.ptx"));
-const PTX_TM_GROW:     &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_grow.ptx"));
-const PTX_TM_ANOMALY:  &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_anomaly.ptx"));
-const PTX_TM_RESET:    &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_reset.ptx"));
-
-/// Capacity trade-offs for 6 GB VRAM (RTX 3060) shared with the model:
-///   n_cells            = 2048 × 32 = 65_536
-///   n_segments_max     = n_cells × MAX_SEGMENTS_PER_CELL
-///   n_synapses_max     = n_segments_max × MAX_SYN_PER_SEGMENT
-///
-/// At 4/20 these are 262_144 segments and ~5.2M synapses (~50 MB per region).
-/// The training loop runs with `reset_each_forward=True`, so segment counts
-/// per window stay well below 32K (typical: ~n_cols new segs per step until
-/// the first matching segment is reused; in a 2048-step window that plateaus
-/// around ~5K total live segments). The 262K ceiling is generous headroom.
-pub const MAX_SEGMENTS_PER_CELL: usize = 4;
-pub const MAX_SYN_PER_SEGMENT:   usize = 20;
-
-const PERM_SCALE: f32 = 32767.0;
-
-fn perm_f32_to_i16(x: f32) -> i16 {
-    let clamped = x.clamp(0.0, 1.0);
-    (clamped * PERM_SCALE).round() as i16
-}
-
-pub struct TemporalMemoryGpu {
-    dev: Arc<CudaDevice>,
-
-    // Config mirror
-    pub n_columns: usize,
-    pub cells_per_column: usize,
-    pub activation_threshold: u32,
-    pub learning_threshold: u32,
-    pub initial_perm_i16: i16,
-    pub conn_thr_i16: i16,
-    pub perm_inc_i16: i16,
-    pub perm_dec_i16: i16,
-    pub predicted_seg_dec_i16: i16,
-    pub max_new_synapse_count: u32,
-
-    // Sizes
-    pub n_cells: usize,
-    pub n_segments_max: usize,
-    pub bits_words: usize,  // n_cells / 32
-
-    // Persistent device buffers
-    seg_cell_id:       CudaSlice<u32>,
-    seg_syn_count:     CudaSlice<u32>,
-    syn_presyn:        CudaSlice<u32>,
-    syn_perm:          CudaSlice<i16>,
-    cell_seg_count:    CudaSlice<u32>,
-
-    cell_active_bits:     CudaSlice<u32>,
-    cell_winner_bits:     CudaSlice<u32>,
-    cell_predictive_bits: CudaSlice<u32>,
-    prev_active_bits:     CudaSlice<u32>,
-    prev_winner_bits:     CudaSlice<u32>,
-
-    col_predicted:         CudaSlice<u8>,
-    seg_num_active_conn:   CudaSlice<u32>,
-    seg_num_active_pot:    CudaSlice<u32>,
-    unpredicted_count:     CudaSlice<u32>,
-    burst_cols_flat:       CudaSlice<u32>,
-    burst_cols_count:      CudaSlice<u32>,
-    col_best_match:        CudaSlice<u32>,
-
-    iter_counter: u32,
-}
-
-impl TemporalMemoryGpu {
-    pub fn new(
-        dev: Arc<CudaDevice>,
-        n_columns: usize,
-        cells_per_column: usize,
-    ) -> Result<Self, DriverError> {
-        let n_cells = n_columns * cells_per_column;
-        assert!(n_cells % 32 == 0, "n_cells must be divisible by 32 for bitsets");
-        let n_segments_max = n_cells * MAX_SEGMENTS_PER_CELL;
-        let bits_words = n_cells / 32;
-
-        // Numenta defaults.
-        let activation_threshold = 15u32;
-        let learning_threshold   = 13u32;
-        let initial_perm_i16        = perm_f32_to_i16(0.21);
-        let conn_thr_i16            = perm_f32_to_i16(0.50);
-        let perm_inc_i16            = perm_f32_to_i16(0.10);
-        let perm_dec_i16            = perm_f32_to_i16(0.10);
-        let predicted_seg_dec_i16   = perm_f32_to_i16(0.10);
-        let max_new_synapse_count   = 20u32;
-
-        // Allocate buffers.
-        let seg_cell_id_host: Vec<u32> = vec![u32::MAX; n_segments_max];
-        let seg_cell_id = dev.htod_sync_copy(&seg_cell_id_host)?;
-        let seg_syn_count = dev.alloc_zeros::<u32>(n_segments_max)?;
-        let syn_presyn = dev.alloc_zeros::<u32>(n_segments_max * MAX_SYN_PER_SEGMENT)?;
-        let syn_perm = dev.alloc_zeros::<i16>(n_segments_max * MAX_SYN_PER_SEGMENT)?;
-        let cell_seg_count = dev.alloc_zeros::<u32>(n_cells)?;
-
-        let cell_active_bits     = dev.alloc_zeros::<u32>(bits_words)?;
-        let cell_winner_bits     = dev.alloc_zeros::<u32>(bits_words)?;
-        let cell_predictive_bits = dev.alloc_zeros::<u32>(bits_words)?;
-        let prev_active_bits     = dev.alloc_zeros::<u32>(bits_words)?;
-        let prev_winner_bits     = dev.alloc_zeros::<u32>(bits_words)?;
-
-        let col_predicted       = dev.alloc_zeros::<u8>(n_columns)?;
-        let seg_num_active_conn = dev.alloc_zeros::<u32>(n_segments_max)?;
-        let seg_num_active_pot  = dev.alloc_zeros::<u32>(n_segments_max)?;
-        let unpredicted_count   = dev.alloc_zeros::<u32>(1)?;
-        // Bursting columns for one step bounded by n_columns.
-        let burst_cols_flat   = dev.alloc_zeros::<u32>(n_columns)?;
-        let burst_cols_count  = dev.alloc_zeros::<u32>(1)?;
-        let col_best_match    = dev.alloc_zeros::<u32>(n_columns)?;
-
-        // Load PTX modules.
-        let modules = [
-            ("htm_tm_predict",  PTX_TM_PREDICT,  "tm_predict"),
-            ("htm_tm_activate", PTX_TM_ACTIVATE, "tm_activate"),
-            ("htm_tm_learn",    PTX_TM_LEARN,    "tm_learn_reinforce"),
-            ("htm_tm_punish",   PTX_TM_PUNISH,   "tm_punish"),
-            ("htm_tm_grow",     PTX_TM_GROW,     "tm_grow"),
-            ("htm_tm_anomaly",  PTX_TM_ANOMALY,  "tm_anomaly"),
-            ("htm_tm_reset",    PTX_TM_RESET,    "tm_reset_step"),
-        ];
-        for (modname, ptx, fnname) in modules {
-            if dev.get_func(modname, fnname).is_none() {
-                dev.load_ptx(Ptx::from_src(ptx), modname, &[fnname])?;
-            }
-        }
-
-        Ok(Self {
-            dev,
-            n_columns,
-            cells_per_column,
-            activation_threshold,
-            learning_threshold,
-            initial_perm_i16,
-            conn_thr_i16,
-            perm_inc_i16,
-            perm_dec_i16,
-            predicted_seg_dec_i16,
-            max_new_synapse_count,
-            n_cells,
-            n_segments_max,
-            bits_words,
-            seg_cell_id,
-            seg_syn_count,
-            syn_presyn,
-            syn_perm,
-            cell_seg_count,
-            cell_active_bits,
-            cell_winner_bits,
-            cell_predictive_bits,
-            prev_active_bits,
-            prev_winner_bits,
-            col_predicted,
-            seg_num_active_conn,
-            seg_num_active_pot,
-            unpredicted_count,
-            burst_cols_flat,
-            burst_cols_count,
-            col_best_match,
-            iter_counter: 0,
-        })
-    }
-
-    // --- Fused-path accessors ---
-    pub fn seg_cell_id_accessor(&self) -> &CudaSlice<u32> { &self.seg_cell_id }
-    pub fn seg_syn_count_accessor(&self) -> &CudaSlice<u32> { &self.seg_syn_count }
-    pub fn syn_presyn_accessor(&self) -> &CudaSlice<u32> { &self.syn_presyn }
-    pub fn syn_perm_accessor(&self) -> &CudaSlice<i16> { &self.syn_perm }
-    pub fn cell_seg_count_accessor(&self) -> &CudaSlice<u32> { &self.cell_seg_count }
-
-    /// Hard reset — clear everything (predictive + active + segments).
-    pub fn reset(&mut self) -> Result<(), DriverError> {
-        // Restore "unused" sentinel in seg_cell_id.
-        let unused_host: Vec<u32> = vec![u32::MAX; self.n_segments_max];
-        self.dev.htod_sync_copy_into(&unused_host, &mut self.seg_cell_id)?;
-        self.dev.memset_zeros(&mut self.seg_syn_count)?;
-        self.dev.memset_zeros(&mut self.cell_seg_count)?;
-        self.dev.memset_zeros(&mut self.cell_active_bits)?;
-        self.dev.memset_zeros(&mut self.cell_winner_bits)?;
-        self.dev.memset_zeros(&mut self.cell_predictive_bits)?;
-        self.dev.memset_zeros(&mut self.prev_active_bits)?;
-        self.dev.memset_zeros(&mut self.prev_winner_bits)?;
-        self.dev.memset_zeros(&mut self.col_best_match)?;
-        self.iter_counter = 0;
-        Ok(())
-    }
-
-    fn build_cfg(&self) -> TmConfig {
-        TmConfig {
-            activation_threshold: self.activation_threshold,
-            learning_threshold: self.learning_threshold,
-            cells_per_column: self.cells_per_column as u32,
-            synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
-            n_segments: self.n_segments_max as u32,
-            n_cells: self.n_cells as u32,
-            max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
-            max_new_synapses: self.max_new_synapse_count,
-            conn_thr_i16: self.conn_thr_i16 as i32,
-            perm_inc_i16: self.perm_inc_i16 as i32,
-            perm_dec_i16: self.perm_dec_i16 as i32,
-            predicted_seg_dec_i16: self.predicted_seg_dec_i16 as i32,
-            initial_perm_i16: self.initial_perm_i16 as i32,
-            iter_seed: self.iter_counter,
-            n_cols: self.n_columns as u32,
-            bits_words: self.bits_words as u32,
-        }
-    }
-
-    /// Run one TM step on the GPU. Takes the SP active-column mask (u8, already
-    /// on device) and writes `anomaly_out[t_slot]`.
-    pub fn step(
-        &mut self,
-        sp_active_mask: &CudaSlice<u8>,
-        anomaly_out: &mut CudaSlice<f32>,
-        t_slot: u32,
-        learn: bool,
-    ) -> Result<(), DriverError> {
-        let n_cells = self.n_cells;
-        let n_cols = self.n_columns;
-
-        let predict_fn  = self.dev.get_func("htm_tm_predict",  "tm_predict").unwrap();
-        let activate_fn = self.dev.get_func("htm_tm_activate", "tm_activate").unwrap();
-        let learn_fn    = self.dev.get_func("htm_tm_learn",    "tm_learn_reinforce").unwrap();
-        let punish_fn   = self.dev.get_func("htm_tm_punish",   "tm_punish").unwrap();
-        let grow_fn     = self.dev.get_func("htm_tm_grow",     "tm_grow").unwrap();
-        let anom_fn     = self.dev.get_func("htm_tm_anomaly",  "tm_anomaly").unwrap();
-        let reset_fn    = self.dev.get_func("htm_tm_reset",    "tm_reset_step").unwrap();
-
-        self.iter_counter = self.iter_counter.wrapping_add(1);
-        let cfg_val = self.build_cfg();
-
-        // 0. Per-step reset.
-        let reset_words = self.bits_words.max(n_cols);
-        let reset_cfg = LaunchConfig {
-            grid_dim: (((reset_words + 255) / 256) as u32, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        unsafe {
-            reset_fn.clone().launch(
-                reset_cfg,
-                (
-                    &mut self.cell_active_bits,
-                    &mut self.cell_winner_bits,
-                    &mut self.cell_predictive_bits,
-                    &mut self.prev_active_bits,
-                    &mut self.prev_winner_bits,
-                    &mut self.col_predicted,
-                    &mut self.unpredicted_count,
-                    &mut self.burst_cols_count,
-                    &mut self.col_best_match,
-                    self.bits_words as u32,
-                    n_cols as u32,
-                ),
-            )?;
-        }
-
-        // 1. Predict (grid = n_cells; each block iterates its cell's segments).
-        let predict_cfg = LaunchConfig {
-            grid_dim: (n_cells as u32, 1, 1),
-            block_dim: (32, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        unsafe {
-            predict_fn.clone().launch(
-                predict_cfg,
-                (
-                    &self.seg_cell_id,
-                    &self.seg_syn_count,
-                    &self.syn_presyn,
-                    &self.syn_perm,
-                    &self.prev_active_bits,
-                    &mut self.cell_predictive_bits,
-                    &mut self.col_predicted,
-                    &mut self.seg_num_active_conn,
-                    &mut self.seg_num_active_pot,
-                    &mut self.col_best_match,
-                    &self.cell_seg_count,
-                    cfg_val,
-                ),
-            )?;
-        }
-
-        // 2. Activate.
-        let activate_cfg = LaunchConfig {
-            grid_dim: (((n_cols + 255) / 256) as u32, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        unsafe {
-            activate_fn.clone().launch(
-                activate_cfg,
-                (
-                    sp_active_mask,
-                    &self.col_predicted,
-                    &self.cell_predictive_bits,
-                    &mut self.cell_active_bits,
-                    &mut self.cell_winner_bits,
-                    &mut self.unpredicted_count,
-                    &mut self.burst_cols_flat,
-                    &mut self.burst_cols_count,
-                    cfg_val,
-                ),
-            )?;
-        }
-
-        // 3. Anomaly.
-        let anom_cfg = LaunchConfig {
-            grid_dim: (1, 1, 1),
-            block_dim: (256, 1, 1),
-            shared_mem_bytes: 0,
-        };
-        unsafe {
-            anom_fn.clone().launch(
-                anom_cfg,
-                (
-                    sp_active_mask,
-                    &self.unpredicted_count,
-                    anomaly_out,
-                    t_slot,
-                    n_cols as u32,
-                ),
-            )?;
-        }
-
-        if learn {
-            // 4. Reinforce (grid = n_cells).
-            let learn_cfg = LaunchConfig {
-                grid_dim: (n_cells as u32, 1, 1),
-                block_dim: (32, 1, 1),
-                shared_mem_bytes: 0,
-            };
-            unsafe {
-                learn_fn.clone().launch(
-                    learn_cfg,
-                    (
-                        &self.seg_cell_id,
-                        &self.seg_syn_count,
-                        &self.syn_presyn,
-                        &mut self.syn_perm,
-                        &self.seg_num_active_conn,
-                        &self.prev_active_bits,
-                        sp_active_mask,
-                        &self.col_predicted,
-                        &self.cell_seg_count,
-                        cfg_val,
-                    ),
-                )?;
-            }
-
-            // 5. Punish.
-            unsafe {
-                punish_fn.clone().launch(
-                    learn_cfg,
-                    (
-                        &self.seg_cell_id,
-                        &self.seg_syn_count,
-                        &self.syn_presyn,
-                        &mut self.syn_perm,
-                        &self.seg_num_active_pot,
-                        &self.prev_active_bits,
-                        sp_active_mask,
-                        &self.cell_seg_count,
-                        cfg_val,
-                    ),
-                )?;
-            }
-
-            // 6. Grow.
-            let grow_cfg = LaunchConfig {
-                grid_dim: (n_cols as u32, 1, 1),
-                block_dim: (32, 1, 1),
-                shared_mem_bytes: 0,
-            };
-            unsafe {
-                grow_fn.clone().launch(
-                    grow_cfg,
-                    (
-                        &mut self.seg_cell_id,
-                        &mut self.seg_syn_count,
-                        &mut self.syn_presyn,
-                        &mut self.syn_perm,
-                        &mut self.cell_seg_count,
-                        &self.burst_cols_flat,
-                        &self.burst_cols_count,
-                        &self.prev_winner_bits,
-                        &self.prev_active_bits,
-                        &self.col_best_match,
-                        cfg_val,
-                    ),
-                )?;
-            }
-        }
-
-        Ok(())
-    }
-}
+//! GPU Temporal Memory.
+//!
+//! Flat device storage. Pre-allocated segment slab:
+//!   n_cells = n_columns * cells_per_column
+//!   n_segments_max = n_cells * MAX_SEGMENTS_PER_CELL
+//!   n_synapses_max = n_segments_max * MAX_SYN_PER_SEGMENT
+//!
+//! Defaults (CPU parity targets relaxed on GPU to keep memory tractable):
+//!   MAX_SEGMENTS_PER_CELL = 16
+//!   MAX_SYN_PER_SEGMENT   = 32
+//!
+//! At n_cells = 65536:
+//!   n_segments_max =  1_048_576   (~1M)
+//!   n_synapses_max = 33_554_432   (~33M)
+//! Storage:
+//!   syn_presyn : u32  × 33M = 128 MB
+//!   syn_perm   : i16  × 33M =  64 MB
+//!   seg_cell   : u32  ×  1M =   4 MB
+//!   seg_syn_n  : u32  ×  1M =   4 MB
+//!   misc bitsets etc        ~ <1 MB
+//!   -------------------------------
+//!   Total per region        ~200 MB
+//!
+//! Permanences are stored as i16 scaled by 32767 (→ [0, 32767] represents
+//! [0.0, 1.0]). inc/dec are provided pre-scaled.
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaDevice, CudaSlice, DriverError, DeviceRepr, LaunchAsync, LaunchConfig};
+use cudarc::nvrtc::Ptx;
+
+/// Packed config struct passed by value to TM kernels to stay under
+/// cudarc's 12-tuple launch limit. Layout must match the C-side
+/// `TmConfig` struct declared in each kernel.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct TmConfig {
+    pub activation_threshold: u32,
+    pub learning_threshold: u32,
+    pub cells_per_column: u32,
+    pub synapses_per_segment: u32,
+    pub n_segments: u32,
+    pub n_cells: u32,
+    pub max_segments_per_cell: u32,
+    pub max_new_synapses: u32,
+    pub conn_thr_i16: i32,        // i16 widened to i32 for alignment
+    pub perm_inc_i16: i32,
+    pub perm_dec_i16: i32,
+    pub predicted_seg_dec_i16: i32,
+    pub initial_perm_i16: i32,
+    pub iter_seed: u32,
+    pub n_cols: u32,
+    pub bits_words: u32,
+}
+
+unsafe impl DeviceRepr for TmConfig {}
+
+// Embedded PTX.
+const PTX_TM_PREDICT:  &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_predict.ptx"));
+const PTX_TM_ACTIVATE: &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_activate.ptx"));
+const PTX_TM_LEARN:    &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_learn.ptx"));
+const PTX_TM_PUNISH:   &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_punish.ptx"));
+const PTX_TM_GROW:     &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_grow.ptx"));
+const PTX_TM_ANOMALY:  &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_anomaly.ptx"));
+const PTX_TM_RESET:    &str = include_str!(concat!(env!("HTM_GPU_PTX_DIR"), "/tm_reset.ptx"));
+
+/// Capacity trade-offs for 6 GB VRAM (RTX 3060) shared with the model:
+///   n_cells            = 2048 × 32 = 65_536
+///   n_segments_max     = n_cells × MAX_SEGMENTS_PER_CELL
+///   n_synapses_max     = n_segments_max × MAX_SYN_PER_SEGMENT
+///
+/// At 4/20 these are 262_144 segments and ~5.2M synapses (~50 MB per region).
+/// The training loop runs with `reset_each_forward=True`, so segment counts
+/// per window stay well below 32K (typical: ~n_cols new segs per step until
+/// the first matching segment is reused; in a 2048-step window that plateaus
+/// around ~5K total live segments). The 262K ceiling is generous headroom.
+pub const MAX_SEGMENTS_PER_CELL: usize = 4;
+pub const MAX_SYN_PER_SEGMENT:   usize = 20;
+
+const PERM_SCALE: f32 = 32767.0;
+
+fn perm_f32_to_i16(x: f32) -> i16 {
+    let clamped = x.clamp(0.0, 1.0);
+    (clamped * PERM_SCALE).round() as i16
+}
+
+pub struct TemporalMemoryGpu {
+    dev: Arc<CudaDevice>,
+
+    // Config mirror
+    pub n_columns: usize,
+    pub cells_per_column: usize,
+    pub activation_threshold: u32,
+    pub learning_threshold: u32,
+    pub initial_perm_i16: i16,
+    pub conn_thr_i16: i16,
+    pub perm_inc_i16: i16,
+    pub perm_dec_i16: i16,
+    pub predicted_seg_dec_i16: i16,
+    pub max_new_synapse_count: u32,
+
+    // Sizes
+    pub n_cells: usize,
+    pub n_segments_max: usize,
+    pub bits_words: usize,  // n_cells / 32
+
+    // Persistent device buffers
+    seg_cell_id:       CudaSlice<u32>,
+    seg_syn_count:     CudaSlice<u32>,
+    syn_presyn:        CudaSlice<u32>,
+    syn_perm:          CudaSlice<i16>,
+    cell_seg_count:    CudaSlice<u32>,
+
+    cell_active_bits:     CudaSlice<u32>,
+    cell_winner_bits:     CudaSlice<u32>,
+    cell_predictive_bits: CudaSlice<u32>,
+    prev_active_bits:     CudaSlice<u32>,
+    prev_winner_bits:     CudaSlice<u32>,
+
+    col_predicted:         CudaSlice<u8>,
+    seg_num_active_conn:   CudaSlice<u32>,
+    seg_num_active_pot:    CudaSlice<u32>,
+    unpredicted_count:     CudaSlice<u32>,
+    burst_cols_flat:       CudaSlice<u32>,
+    burst_cols_count:      CudaSlice<u32>,
+    col_best_match:        CudaSlice<u32>,
+
+    iter_counter: u32,
+}
+
+impl TemporalMemoryGpu {
+    pub fn new(
+        dev: Arc<CudaDevice>,
+        n_columns: usize,
+        cells_per_column: usize,
+    ) -> Result<Self, DriverError> {
+        let n_cells = n_columns * cells_per_column;
+        assert!(n_cells % 32 == 0, "n_cells must be divisible by 32 for bitsets");
+        let n_segments_max = n_cells * MAX_SEGMENTS_PER_CELL;
+        let bits_words = n_cells / 32;
+
+        // Numenta defaults.
+        let activation_threshold = 15u32;
+        let learning_threshold   = 13u32;
+        let initial_perm_i16        = perm_f32_to_i16(0.21);
+        let conn_thr_i16            = perm_f32_to_i16(0.50);
+        let perm_inc_i16            = perm_f32_to_i16(0.10);
+        let perm_dec_i16            = perm_f32_to_i16(0.10);
+        let predicted_seg_dec_i16   = perm_f32_to_i16(0.10);
+        let max_new_synapse_count   = 20u32;
+
+        // Allocate buffers.
+        let seg_cell_id_host: Vec<u32> = vec![u32::MAX; n_segments_max];
+        let seg_cell_id = dev.htod_sync_copy(&seg_cell_id_host)?;
+        let seg_syn_count = dev.alloc_zeros::<u32>(n_segments_max)?;
+        let syn_presyn = dev.alloc_zeros::<u32>(n_segments_max * MAX_SYN_PER_SEGMENT)?;
+        let syn_perm = dev.alloc_zeros::<i16>(n_segments_max * MAX_SYN_PER_SEGMENT)?;
+        let cell_seg_count = dev.alloc_zeros::<u32>(n_cells)?;
+
+        let cell_active_bits     = dev.alloc_zeros::<u32>(bits_words)?;
+        let cell_winner_bits     = dev.alloc_zeros::<u32>(bits_words)?;
+        let cell_predictive_bits = dev.alloc_zeros::<u32>(bits_words)?;
+        let prev_active_bits     = dev.alloc_zeros::<u32>(bits_words)?;
+        let prev_winner_bits     = dev.alloc_zeros::<u32>(bits_words)?;
+
+        let col_predicted       = dev.alloc_zeros::<u8>(n_columns)?;
+        let seg_num_active_conn = dev.alloc_zeros::<u32>(n_segments_max)?;
+        let seg_num_active_pot  = dev.alloc_zeros::<u32>(n_segments_max)?;
+        let unpredicted_count   = dev.alloc_zeros::<u32>(1)?;
+        // Bursting columns for one step bounded by n_columns.
+        let burst_cols_flat   = dev.alloc_zeros::<u32>(n_columns)?;
+        let burst_cols_count  = dev.alloc_zeros::<u32>(1)?;
+        let col_best_match    = dev.alloc_zeros::<u32>(n_columns)?;
+
+        // Load PTX modules.
+        let modules = [
+            ("htm_tm_predict",  PTX_TM_PREDICT,  "tm_predict"),
+            ("htm_tm_activate", PTX_TM_ACTIVATE, "tm_activate"),
+            ("htm_tm_learn",    PTX_TM_LEARN,    "tm_learn_reinforce"),
+            ("htm_tm_punish",   PTX_TM_PUNISH,   "tm_punish"),
+            ("htm_tm_grow",     PTX_TM_GROW,     "tm_grow"),
+            ("htm_tm_anomaly",  PTX_TM_ANOMALY,  "tm_anomaly"),
+            ("htm_tm_reset",    PTX_TM_RESET,    "tm_reset_step"),
+        ];
+        for (modname, ptx, fnname) in modules {
+            if dev.get_func(modname, fnname).is_none() {
+                dev.load_ptx(Ptx::from_src(ptx), modname, &[fnname])?;
+            }
+        }
+
+        Ok(Self {
+            dev,
+            n_columns,
+            cells_per_column,
+            activation_threshold,
+            learning_threshold,
+            initial_perm_i16,
+            conn_thr_i16,
+            perm_inc_i16,
+            perm_dec_i16,
+            predicted_seg_dec_i16,
+            max_new_synapse_count,
+            n_cells,
+            n_segments_max,
+            bits_words,
+            seg_cell_id,
+            seg_syn_count,
+            syn_presyn,
+            syn_perm,
+            cell_seg_count,
+            cell_active_bits,
+            cell_winner_bits,
+            cell_predictive_bits,
+            prev_active_bits,
+            prev_winner_bits,
+            col_predicted,
+            seg_num_active_conn,
+            seg_num_active_pot,
+            unpredicted_count,
+            burst_cols_flat,
+            burst_cols_count,
+            col_best_match,
+            iter_counter: 0,
+        })
+    }
+
+    // --- Fused-path accessors ---
+    pub fn seg_cell_id_accessor(&self) -> &CudaSlice<u32> { &self.seg_cell_id }
+    pub fn seg_syn_count_accessor(&self) -> &CudaSlice<u32> { &self.seg_syn_count }
+    pub fn syn_presyn_accessor(&self) -> &CudaSlice<u32> { &self.syn_presyn }
+    pub fn syn_perm_accessor(&self) -> &CudaSlice<i16> { &self.syn_perm }
+    pub fn cell_seg_count_accessor(&self) -> &CudaSlice<u32> { &self.cell_seg_count }
+
+    /// Hard reset — clear everything (predictive + active + segments).
+    pub fn reset(&mut self) -> Result<(), DriverError> {
+        // Restore "unused" sentinel in seg_cell_id.
+        let unused_host: Vec<u32> = vec![u32::MAX; self.n_segments_max];
+        self.dev.htod_sync_copy_into(&unused_host, &mut self.seg_cell_id)?;
+        self.dev.memset_zeros(&mut self.seg_syn_count)?;
+        self.dev.memset_zeros(&mut self.cell_seg_count)?;
+        self.dev.memset_zeros(&mut self.cell_active_bits)?;
+        self.dev.memset_zeros(&mut self.cell_winner_bits)?;
+        self.dev.memset_zeros(&mut self.cell_predictive_bits)?;
+        self.dev.memset_zeros(&mut self.prev_active_bits)?;
+        self.dev.memset_zeros(&mut self.prev_winner_bits)?;
+        self.dev.memset_zeros(&mut self.col_best_match)?;
+        self.iter_counter = 0;
+        Ok(())
+    }
+
+    fn build_cfg(&self) -> TmConfig {
+        TmConfig {
+            activation_threshold: self.activation_threshold,
+            learning_threshold: self.learning_threshold,
+            cells_per_column: self.cells_per_column as u32,
+            synapses_per_segment: MAX_SYN_PER_SEGMENT as u32,
+            n_segments: self.n_segments_max as u32,
+            n_cells: self.n_cells as u32,
+            max_segments_per_cell: MAX_SEGMENTS_PER_CELL as u32,
+            max_new_synapses: self.max_new_synapse_count,
+            conn_thr_i16: self.conn_thr_i16 as i32,
+            perm_inc_i16: self.perm_inc_i16 as i32,
+            perm_dec_i16: self.perm_dec_i16 as i32,
+            predicted_seg_dec_i16: self.predicted_seg_dec_i16 as i32,
+            initial_perm_i16: self.initial_perm_i16 as i32,
+            iter_seed: self.iter_counter,
+            n_cols: self.n_columns as u32,
+            bits_words: self.bits_words as u32,
+        }
+    }
+
+    /// Run one TM step on the GPU. Takes the SP active-column mask (u8, already
+    /// on device) and writes `anomaly_out[t_slot]`.
+    pub fn step(
+        &mut self,
+        sp_active_mask: &CudaSlice<u8>,
+        anomaly_out: &mut CudaSlice<f32>,
+        t_slot: u32,
+        learn: bool,
+    ) -> Result<(), DriverError> {
+        let n_cells = self.n_cells;
+        let n_cols = self.n_columns;
+
+        let predict_fn  = self.dev.get_func("htm_tm_predict",  "tm_predict").unwrap();
+        let activate_fn = self.dev.get_func("htm_tm_activate", "tm_activate").unwrap();
+        let learn_fn    = self.dev.get_func("htm_tm_learn",    "tm_learn_reinforce").unwrap();
+        let punish_fn   = self.dev.get_func("htm_tm_punish",   "tm_punish").unwrap();
+        let grow_fn     = self.dev.get_func("htm_tm_grow",     "tm_grow").unwrap();
+        let anom_fn     = self.dev.get_func("htm_tm_anomaly",  "tm_anomaly").unwrap();
+        let reset_fn    = self.dev.get_func("htm_tm_reset",    "tm_reset_step").unwrap();
+
+        self.iter_counter = self.iter_counter.wrapping_add(1);
+        let cfg_val = self.build_cfg();
+
+        // 0. Per-step reset.
+        let reset_words = self.bits_words.max(n_cols);
+        let reset_cfg = LaunchConfig {
+            grid_dim: (((reset_words + 255) / 256) as u32, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        unsafe {
+            reset_fn.clone().launch(
+                reset_cfg,
+                (
+                    &mut self.cell_active_bits,
+                    &mut self.cell_winner_bits,
+                    &mut self.cell_predictive_bits,
+                    &mut self.prev_active_bits,
+                    &mut self.prev_winner_bits,
+                    &mut self.col_predicted,
+                    &mut self.unpredicted_count,
+                    &mut self.burst_cols_count,
+                    &mut self.col_best_match,
+                    self.bits_words as u32,
+                    n_cols as u32,
+                ),
+            )?;
+        }
+
+        // 1. Predict (grid = n_cells; each block iterates its cell's segments).
+        let predict_cfg = LaunchConfig {
+            grid_dim: (n_cells as u32, 1, 1),
+            block_dim: (32, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        unsafe {
+            predict_fn.clone().launch(
+                predict_cfg,
+                (
+                    &self.seg_cell_id,
+                    &self.seg_syn_count,
+                    &self.syn_presyn,
+                    &self.syn_perm,
+                    &self.prev_active_bits,
+                    &mut self.cell_predictive_bits,
+                    &mut self.col_predicted,
+                    &mut self.seg_num_active_conn,
+                    &mut self.seg_num_active_pot,
+                    &mut self.col_best_match,
+                    &self.cell_seg_count,
+                    cfg_val,
+                ),
+            )?;
+        }
+
+        // 2. Activate.
+        let activate_cfg = LaunchConfig {
+            grid_dim: (((n_cols + 255) / 256) as u32, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        unsafe {
+            activate_fn.clone().launch(
+                activate_cfg,
+                (
+                    sp_active_mask,
+                    &self.col_predicted,
+                    &self.cell_predictive_bits,
+                    &mut self.cell_active_bits,
+                    &mut self.cell_winner_bits,
+                    &mut self.unpredicted_count,
+                    &mut self.burst_cols_flat,
+                    &mut self.burst_cols_count,
+                    cfg_val,
+                ),
+            )?;
+        }
+
+        // 3. Anomaly.
+        let anom_cfg = LaunchConfig {
+            grid_dim: (1, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        unsafe {
+            anom_fn.clone().launch(
+                anom_cfg,
+                (
+                    sp_active_mask,
+                    &self.unpredicted_count,
+                    anomaly_out,
+                    t_slot,
+                    n_cols as u32,
+                ),
+            )?;
+        }
+
+        if learn {
+            // 4. Reinforce (grid = n_cells).
+            let learn_cfg = LaunchConfig {
+                grid_dim: (n_cells as u32, 1, 1),
+                block_dim: (32, 1, 1),
+                shared_mem_bytes: 0,
+            };
+            unsafe {
+                learn_fn.clone().launch(
+                    learn_cfg,
+                    (
+                        &self.seg_cell_id,
+                        &self.seg_syn_count,
+                        &self.syn_presyn,
+                        &mut self.syn_perm,
+                        &self.seg_num_active_conn,
+                        &self.prev_active_bits,
+                        sp_active_mask,
+                        &self.col_predicted,
+                        &self.cell_seg_count,
+                        cfg_val,
+                    ),
+                )?;
+            }
+
+            // 5. Punish.
+            unsafe {
+                punish_fn.clone().launch(
+                    learn_cfg,
+                    (
+                        &self.seg_cell_id,
+                        &self.seg_syn_count,
+                        &self.syn_presyn,
+                        &mut self.syn_perm,
+                        &self.seg_num_active_pot,
+                        &self.prev_active_bits,
+                        sp_active_mask,
+                        &self.cell_seg_count,
+                        cfg_val,
+                    ),
+                )?;
+            }
+
+            // 6. Grow.
+            let grow_cfg = LaunchConfig {
+                grid_dim: (n_cols as u32, 1, 1),
+                block_dim: (32, 1, 1),
+                shared_mem_bytes: 0,
+            };
+            unsafe {
+                grow_fn.clone().launch(
+                    grow_cfg,
+                    (
+                        &mut self.seg_cell_id,
+                        &mut self.seg_syn_count,
+                        &mut self.syn_presyn,
+                        &mut self.syn_perm,
+                        &mut self.cell_seg_count,
+                        &self.burst_cols_flat,
+                        &self.burst_cols_count,
+                        &self.prev_winner_bits,
+                        &self.prev_active_bits,
+                        &self.col_best_match,
+                        cfg_val,
+                    ),
+                )?;
+            }
+        }
+
+        Ok(())
+    }
+}

overlay/htm_rust/src/lib.rs

@@ -1,198 +1,198 @@
-//! pyo3 bindings for HTMRegion (Numenta BAMI-spec HTM).
-//!
-//! Exposed class:
-//!     HTMRegion(input_bits, n_columns, cells_per_column, seed) -> HTMRegion
-//!       .step(input_sdr: np.ndarray[bool; input_bits], learn: bool = True)
-//!           -> (active_columns: np.ndarray[bool; n_columns],
-//!               active_cells:   np.ndarray[bool; n_columns*cells_per_column],
-//!               predicted_cells:np.ndarray[bool; n_columns*cells_per_column],
-//!               anomaly: float)
-//!       .reset()
-//!       .n_columns -> int
-//!       .cells_per_column -> int
-//!       .input_bits -> int
-//!
-//! GIL is dropped during the heavy compute via `py.allow_threads(...)` so the
-//! region is effectively `Send` for Python-side threading.
-
-// pyo3 0.22 `#[pymethods]` expansion inserts an implicit `.into()` on the
-// returned `Result` to normalise the error type, which clippy reports as
-// `useless_conversion` when our methods already return `PyErr`. The emitted
-// code sits outside the user-written impl, so item-level allows don't reach
-// it; the module-wide allow is the documented workaround.
-#![allow(clippy::useless_conversion)]
-
-mod region;
-mod sp;
-mod tm;
-
-#[cfg(feature = "gpu")]
-mod gpu;
-
-use numpy::{
-    IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2,
-    PyUntypedArrayMethods,
-};
-use pyo3::prelude::*;
-
-use crate::region::HTMRegionCore;
-
-/// Result of one HTM step: (active_columns, active_cells, predicted_cells, anomaly).
-type StepOutput<'py> = (
-    Bound<'py, PyArray1<bool>>,
-    Bound<'py, PyArray1<bool>>,
-    Bound<'py, PyArray1<bool>>,
-    f32,
-);
-
-#[pyclass(module = "htm_rust")]
-pub struct HTMRegion {
-    core: HTMRegionCore,
-}
-
-#[pymethods]
-impl HTMRegion {
-    /// Create a new HTM region.
-    ///
-    /// Args:
-    ///     input_bits: length of binary input SDR
-    ///     n_columns: number of mini-columns in the SP (e.g. 2048)
-    ///     cells_per_column: cells per column in the TM (e.g. 32)
-    ///     seed: RNG seed for reproducibility
-    #[new]
-    #[pyo3(signature = (input_bits, n_columns, cells_per_column, seed=42))]
-    fn new(
-        input_bits: usize,
-        n_columns: usize,
-        cells_per_column: usize,
-        seed: u64,
-    ) -> PyResult<Self> {
-        if input_bits == 0 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "input_bits must be > 0",
-            ));
-        }
-        if n_columns == 0 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "n_columns must be > 0",
-            ));
-        }
-        if cells_per_column == 0 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "cells_per_column must be > 0",
-            ));
-        }
-        Ok(Self {
-            core: HTMRegionCore::new(input_bits, n_columns, cells_per_column, seed),
-        })
-    }
-
-    #[getter]
-    fn input_bits(&self) -> usize { self.core.sp.cfg.input_bits }
-
-    #[getter]
-    fn n_columns(&self) -> usize { self.core.sp.cfg.n_columns }
-
-    #[getter]
-    fn cells_per_column(&self) -> usize { self.core.tm.cfg.cells_per_column }
-
-    /// Process one timestep.
-    ///
-    /// Args:
-    ///     input_sdr: 1-D numpy boolean array of length `input_bits`.
-    ///     learn: if True, update SP permanences and TM synapses.
-    ///
-    /// Returns:
-    ///     (active_columns, active_cells, predicted_cells, anomaly)
-    #[pyo3(signature = (input_sdr, learn=true))]
-    fn step<'py>(
-        &mut self,
-        py: Python<'py>,
-        input_sdr: PyReadonlyArray1<'py, bool>,
-        learn: bool,
-    ) -> PyResult<StepOutput<'py>> {
-        let expected = self.core.sp.cfg.input_bits;
-        let slice = input_sdr.as_slice()?;
-        let got = slice.len();
-        if got != expected {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "input_sdr length {got} != expected input_bits {expected}",
-            )));
-        }
-
-        // Copy input to an owned Vec so we can drop the GIL.
-        let input_vec: Vec<bool> = slice.to_vec();
-
-        let (active_cols, active_cells, predicted_cells, anomaly) =
-            py.allow_threads(|| self.core.step(&input_vec, learn));
-
-        let a: Bound<'py, PyArray1<bool>> = active_cols.into_pyarray_bound(py);
-        let c: Bound<'py, PyArray1<bool>> = active_cells.into_pyarray_bound(py);
-        let p: Bound<'py, PyArray1<bool>> = predicted_cells.into_pyarray_bound(py);
-        Ok((a, c, p, anomaly))
-    }
-
-    /// Clear TM predictive state. Does NOT unlearn synapses.
-    fn reset(&mut self) { self.core.reset(); }
-
-    /// Process T timesteps from a `(T, input_bits)` bool ndarray.
-    ///
-    /// Returns:
-    ///     cols: (T, n_columns) float32 0/1 active-column mask
-    ///     anom: (T,) float32 anomaly scores
-    ///
-    /// Single GIL release for the whole pass, avoiding T × Python-call overhead.
-    #[pyo3(signature = (inputs, learn=true))]
-    fn step_many<'py>(
-        &mut self,
-        py: Python<'py>,
-        inputs: PyReadonlyArray2<'py, bool>,
-        learn: bool,
-    ) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray1<f32>>)> {
-        let shape = inputs.shape();
-        if shape.len() != 2 {
-            return Err(pyo3::exceptions::PyValueError::new_err(
-                "inputs must be 2-D (T, input_bits)",
-            ));
-        }
-        let t = shape[0];
-        let bits = shape[1];
-        let expected = self.core.sp.cfg.input_bits;
-        if bits != expected {
-            return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "inputs last dim {bits} != expected input_bits {expected}",
-            )));
-        }
-        let slice = inputs.as_slice()?;
-        let n_cols = self.core.sp.cfg.n_columns;
-
-        // Own the input buffer so we can drop the GIL.
-        let input_vec: Vec<bool> = slice.to_vec();
-
-        let (cols_u8, anom) =
-            py.allow_threads(|| self.core.step_many(&input_vec, bits, t, learn));
-
-        // Convert u8 mask to f32 for direct numpy consumption.
-        let cols_f32: Vec<f32> = cols_u8.iter().map(|&b| b as f32).collect();
-
-        // Build (T, n_cols) and (T,) arrays.
-        let cols_arr =
-            numpy::PyArray1::from_vec_bound(py, cols_f32)
-                .reshape([t, n_cols])
-                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
-        let anom_arr = numpy::PyArray1::from_vec_bound(py, anom);
-        Ok((cols_arr, anom_arr))
-    }
-}
-
-/// Python module entry point.
-#[pymodule]
-fn htm_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<HTMRegion>()?;
-    #[cfg(feature = "gpu")]
-    {
-        gpu::register(m)?;
-    }
-    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
-    Ok(())
-}
+//! pyo3 bindings for HTMRegion (Numenta BAMI-spec HTM).
+//!
+//! Exposed class:
+//!     HTMRegion(input_bits, n_columns, cells_per_column, seed) -> HTMRegion
+//!       .step(input_sdr: np.ndarray[bool; input_bits], learn: bool = True)
+//!           -> (active_columns: np.ndarray[bool; n_columns],
+//!               active_cells:   np.ndarray[bool; n_columns*cells_per_column],
+//!               predicted_cells:np.ndarray[bool; n_columns*cells_per_column],
+//!               anomaly: float)
+//!       .reset()
+//!       .n_columns -> int
+//!       .cells_per_column -> int
+//!       .input_bits -> int
+//!
+//! GIL is dropped during the heavy compute via `py.allow_threads(...)` so the
+//! region is effectively `Send` for Python-side threading.
+
+// pyo3 0.22 `#[pymethods]` expansion inserts an implicit `.into()` on the
+// returned `Result` to normalise the error type, which clippy reports as
+// `useless_conversion` when our methods already return `PyErr`. The emitted
+// code sits outside the user-written impl, so item-level allows don't reach
+// it; the module-wide allow is the documented workaround.
+#![allow(clippy::useless_conversion)]
+
+mod region;
+mod sp;
+mod tm;
+
+#[cfg(feature = "gpu")]
+mod gpu;
+
+use numpy::{
+    IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2,
+    PyUntypedArrayMethods,
+};
+use pyo3::prelude::*;
+
+use crate::region::HTMRegionCore;
+
+/// Result of one HTM step: (active_columns, active_cells, predicted_cells, anomaly).
+type StepOutput<'py> = (
+    Bound<'py, PyArray1<bool>>,
+    Bound<'py, PyArray1<bool>>,
+    Bound<'py, PyArray1<bool>>,
+    f32,
+);
+
+#[pyclass(module = "htm_rust")]
+pub struct HTMRegion {
+    core: HTMRegionCore,
+}
+
+#[pymethods]
+impl HTMRegion {
+    /// Create a new HTM region.
+    ///
+    /// Args:
+    ///     input_bits: length of binary input SDR
+    ///     n_columns: number of mini-columns in the SP (e.g. 2048)
+    ///     cells_per_column: cells per column in the TM (e.g. 32)
+    ///     seed: RNG seed for reproducibility
+    #[new]
+    #[pyo3(signature = (input_bits, n_columns, cells_per_column, seed=42))]
+    fn new(
+        input_bits: usize,
+        n_columns: usize,
+        cells_per_column: usize,
+        seed: u64,
+    ) -> PyResult<Self> {
+        if input_bits == 0 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "input_bits must be > 0",
+            ));
+        }
+        if n_columns == 0 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "n_columns must be > 0",
+            ));
+        }
+        if cells_per_column == 0 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "cells_per_column must be > 0",
+            ));
+        }
+        Ok(Self {
+            core: HTMRegionCore::new(input_bits, n_columns, cells_per_column, seed),
+        })
+    }
+
+    #[getter]
+    fn input_bits(&self) -> usize { self.core.sp.cfg.input_bits }
+
+    #[getter]
+    fn n_columns(&self) -> usize { self.core.sp.cfg.n_columns }
+
+    #[getter]
+    fn cells_per_column(&self) -> usize { self.core.tm.cfg.cells_per_column }
+
+    /// Process one timestep.
+    ///
+    /// Args:
+    ///     input_sdr: 1-D numpy boolean array of length `input_bits`.
+    ///     learn: if True, update SP permanences and TM synapses.
+    ///
+    /// Returns:
+    ///     (active_columns, active_cells, predicted_cells, anomaly)
+    #[pyo3(signature = (input_sdr, learn=true))]
+    fn step<'py>(
+        &mut self,
+        py: Python<'py>,
+        input_sdr: PyReadonlyArray1<'py, bool>,
+        learn: bool,
+    ) -> PyResult<StepOutput<'py>> {
+        let expected = self.core.sp.cfg.input_bits;
+        let slice = input_sdr.as_slice()?;
+        let got = slice.len();
+        if got != expected {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "input_sdr length {got} != expected input_bits {expected}",
+            )));
+        }
+
+        // Copy input to an owned Vec so we can drop the GIL.
+        let input_vec: Vec<bool> = slice.to_vec();
+
+        let (active_cols, active_cells, predicted_cells, anomaly) =
+            py.allow_threads(|| self.core.step(&input_vec, learn));
+
+        let a: Bound<'py, PyArray1<bool>> = active_cols.into_pyarray_bound(py);
+        let c: Bound<'py, PyArray1<bool>> = active_cells.into_pyarray_bound(py);
+        let p: Bound<'py, PyArray1<bool>> = predicted_cells.into_pyarray_bound(py);
+        Ok((a, c, p, anomaly))
+    }
+
+    /// Clear TM predictive state. Does NOT unlearn synapses.
+    fn reset(&mut self) { self.core.reset(); }
+
+    /// Process T timesteps from a `(T, input_bits)` bool ndarray.
+    ///
+    /// Returns:
+    ///     cols: (T, n_columns) float32 0/1 active-column mask
+    ///     anom: (T,) float32 anomaly scores
+    ///
+    /// Single GIL release for the whole pass, avoiding T × Python-call overhead.
+    #[pyo3(signature = (inputs, learn=true))]
+    fn step_many<'py>(
+        &mut self,
+        py: Python<'py>,
+        inputs: PyReadonlyArray2<'py, bool>,
+        learn: bool,
+    ) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray1<f32>>)> {
+        let shape = inputs.shape();
+        if shape.len() != 2 {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "inputs must be 2-D (T, input_bits)",
+            ));
+        }
+        let t = shape[0];
+        let bits = shape[1];
+        let expected = self.core.sp.cfg.input_bits;
+        if bits != expected {
+            return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                "inputs last dim {bits} != expected input_bits {expected}",
+            )));
+        }
+        let slice = inputs.as_slice()?;
+        let n_cols = self.core.sp.cfg.n_columns;
+
+        // Own the input buffer so we can drop the GIL.
+        let input_vec: Vec<bool> = slice.to_vec();
+
+        let (cols_u8, anom) =
+            py.allow_threads(|| self.core.step_many(&input_vec, bits, t, learn));
+
+        // Convert u8 mask to f32 for direct numpy consumption.
+        let cols_f32: Vec<f32> = cols_u8.iter().map(|&b| b as f32).collect();
+
+        // Build (T, n_cols) and (T,) arrays.
+        let cols_arr =
+            numpy::PyArray1::from_vec_bound(py, cols_f32)
+                .reshape([t, n_cols])
+                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
+        let anom_arr = numpy::PyArray1::from_vec_bound(py, anom);
+        Ok((cols_arr, anom_arr))
+    }
+}
+
+/// Python module entry point.
+#[pymodule]
+fn htm_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<HTMRegion>()?;
+    #[cfg(feature = "gpu")]
+    {
+        gpu::register(m)?;
+    }
+    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
+    Ok(())
+}

overlay/htm_rust/src/region.rs

@@ -1,94 +1,94 @@
-//! HTMRegion: compose SpatialPooler + TemporalMemory into a single step().
-
-use crate::sp::{SpatialPooler, SpatialPoolerConfig};
-use crate::tm::{TemporalMemory, TemporalMemoryConfig};
-
-pub struct HTMRegionCore {
-    pub sp: SpatialPooler,
-    pub tm: TemporalMemory,
-}
-
-impl HTMRegionCore {
-    pub fn new(
-        input_bits: usize,
-        n_columns: usize,
-        cells_per_column: usize,
-        seed: u64,
-    ) -> Self {
-        let defaults = SpatialPoolerConfig::default();
-        let sp_cfg = SpatialPoolerConfig {
-            input_bits,
-            n_columns,
-            // Scale potential_radius to at most the input size.
-            potential_radius: defaults.potential_radius.min(input_bits),
-            ..defaults
-        };
-
-        let tm_cfg = TemporalMemoryConfig {
-            n_columns,
-            cells_per_column,
-            ..TemporalMemoryConfig::default()
-        };
-
-        Self {
-            sp: SpatialPooler::new(sp_cfg, seed),
-            tm: TemporalMemory::new(tm_cfg, seed.wrapping_add(0x9E3779B97F4A7C15)),
-        }
-    }
-
-    /// Process one timestep. Returns (active_columns_mask,
-    /// active_cells_mask, predicted_cells_mask, anomaly).
-    pub fn step(
-        &mut self,
-        input_sdr: &[bool],
-        learn: bool,
-    ) -> (Vec<bool>, Vec<bool>, Vec<bool>, f32) {
-        let active_cols = self.sp.compute(input_sdr, learn);
-
-        let mut active_cols_mask = vec![false; self.sp.cfg.n_columns];
-        for &c in &active_cols {
-            active_cols_mask[c as usize] = true;
-        }
-
-        let anomaly = self.tm.compute(&active_cols, learn);
-
-        // active_cells and predictive_cells are stored as Vec<bool> already.
-        let active_cells_mask = self.tm.active_cells.clone();
-        let predicted_cells_mask = self.tm.predictive_cells.clone();
-
-        (active_cols_mask, active_cells_mask, predicted_cells_mask, anomaly)
-    }
-
-    pub fn reset(&mut self) {
-        self.tm.reset();
-    }
-
-    /// Process T timesteps in one call. Returns flat `(T*n_columns)` active-column
-    /// mask (u8 0/1) and `(T,)` anomaly scores.
-    ///
-    /// Amortises the per-step Python round-trip for training: one GIL release,
-    /// one copy-out. Used by `HTMLayer.step_many`.
-    pub fn step_many(
-        &mut self,
-        inputs_flat: &[bool],
-        input_bits: usize,
-        t: usize,
-        learn: bool,
-    ) -> (Vec<u8>, Vec<f32>) {
-        let n_cols = self.sp.cfg.n_columns;
-        debug_assert_eq!(inputs_flat.len(), t * input_bits);
-        let mut cols = vec![0u8; t * n_cols];
-        let mut anom = vec![0f32; t];
-        for ti in 0..t {
-            let off = ti * input_bits;
-            let input = &inputs_flat[off..off + input_bits];
-            let active_cols = self.sp.compute(input, learn);
-            let co = ti * n_cols;
-            for &c in &active_cols {
-                cols[co + c as usize] = 1;
-            }
-            anom[ti] = self.tm.compute(&active_cols, learn);
-        }
-        (cols, anom)
-    }
-}
+//! HTMRegion: compose SpatialPooler + TemporalMemory into a single step().
+
+use crate::sp::{SpatialPooler, SpatialPoolerConfig};
+use crate::tm::{TemporalMemory, TemporalMemoryConfig};
+
+pub struct HTMRegionCore {
+    pub sp: SpatialPooler,
+    pub tm: TemporalMemory,
+}
+
+impl HTMRegionCore {
+    pub fn new(
+        input_bits: usize,
+        n_columns: usize,
+        cells_per_column: usize,
+        seed: u64,
+    ) -> Self {
+        let defaults = SpatialPoolerConfig::default();
+        let sp_cfg = SpatialPoolerConfig {
+            input_bits,
+            n_columns,
+            // Scale potential_radius to at most the input size.
+            potential_radius: defaults.potential_radius.min(input_bits),
+            ..defaults
+        };
+
+        let tm_cfg = TemporalMemoryConfig {
+            n_columns,
+            cells_per_column,
+            ..TemporalMemoryConfig::default()
+        };
+
+        Self {
+            sp: SpatialPooler::new(sp_cfg, seed),
+            tm: TemporalMemory::new(tm_cfg, seed.wrapping_add(0x9E3779B97F4A7C15)),
+        }
+    }
+
+    /// Process one timestep. Returns (active_columns_mask,
+    /// active_cells_mask, predicted_cells_mask, anomaly).
+    pub fn step(
+        &mut self,
+        input_sdr: &[bool],
+        learn: bool,
+    ) -> (Vec<bool>, Vec<bool>, Vec<bool>, f32) {
+        let active_cols = self.sp.compute(input_sdr, learn);
+
+        let mut active_cols_mask = vec![false; self.sp.cfg.n_columns];
+        for &c in &active_cols {
+            active_cols_mask[c as usize] = true;
+        }
+
+        let anomaly = self.tm.compute(&active_cols, learn);
+
+        // active_cells and predictive_cells are stored as Vec<bool> already.
+        let active_cells_mask = self.tm.active_cells.clone();
+        let predicted_cells_mask = self.tm.predictive_cells.clone();
+
+        (active_cols_mask, active_cells_mask, predicted_cells_mask, anomaly)
+    }
+
+    pub fn reset(&mut self) {
+        self.tm.reset();
+    }
+
+    /// Process T timesteps in one call. Returns flat `(T*n_columns)` active-column
+    /// mask (u8 0/1) and `(T,)` anomaly scores.
+    ///
+    /// Amortises the per-step Python round-trip for training: one GIL release,
+    /// one copy-out. Used by `HTMLayer.step_many`.
+    pub fn step_many(
+        &mut self,
+        inputs_flat: &[bool],
+        input_bits: usize,
+        t: usize,
+        learn: bool,
+    ) -> (Vec<u8>, Vec<f32>) {
+        let n_cols = self.sp.cfg.n_columns;
+        debug_assert_eq!(inputs_flat.len(), t * input_bits);
+        let mut cols = vec![0u8; t * n_cols];
+        let mut anom = vec![0f32; t];
+        for ti in 0..t {
+            let off = ti * input_bits;
+            let input = &inputs_flat[off..off + input_bits];
+            let active_cols = self.sp.compute(input, learn);
+            let co = ti * n_cols;
+            for &c in &active_cols {
+                cols[co + c as usize] = 1;
+            }
+            anom[ti] = self.tm.compute(&active_cols, learn);
+        }
+        (cols, anom)
+    }
+}

overlay/htm_rust/src/sp.rs

@@ -1,302 +1,302 @@
-//! Numenta BAMI-spec Spatial Pooler.
-//!
-//! Implements:
-//!   - 2048 (configurable) mini-columns with proximal dendrites
-//!   - `potential_synapses` (default 40) synapses per column sampled from
-//!     `potential_radius` (default 1024) random input bits
-//!   - Permanence in [0.0, 1.0] (f32), connected_threshold = 0.5
-//!   - syn_perm_active_inc = +0.04, syn_perm_inactive_dec = -0.008
-//!   - Global k-WTA inhibition (top `sparsity` fraction of columns)
-//!   - Boost factor with exponential duty-cycle tracking (Numenta formula)
-//!
-//! Reference: BAMI "Spatial Pooling Algorithm Details" (Numenta, 2017).
-
-use rand::Rng;
-use rand::SeedableRng;
-use rand::seq::SliceRandom;
-use rand_xoshiro::Xoshiro256PlusPlus;
-
-/// A single proximal dendrite: a sparse set of potential synapses onto
-/// specific input bit indices, with per-synapse permanence values.
-#[derive(Clone)]
-pub struct ProximalDendrite {
-    /// Indices into the input SDR.  Length == potential_synapses.
-    pub inputs: Vec<u32>,
-    /// Permanence for each potential synapse (same length as `inputs`).
-    pub perms: Vec<f32>,
-}
-
-pub struct SpatialPoolerConfig {
-    pub input_bits: usize,
-    pub n_columns: usize,
-    /// Size of the random input sample per column.
-    pub potential_radius: usize,
-    /// Number of potential synapses per column's proximal dendrite.
-    pub potential_synapses: usize,
-    pub connected_threshold: f32,
-    pub syn_perm_active_inc: f32,
-    pub syn_perm_inactive_dec: f32,
-    /// Target fraction of columns active per step (e.g. 0.02 for 2%).
-    pub sparsity: f32,
-    /// Duty cycle EMA period.
-    pub duty_cycle_period: f32,
-    /// Boost strength. Set to 0.0 to disable boosting.
-    pub boost_strength: f32,
-    /// Initial permanence span around the connected threshold.
-    pub init_perm_span: f32,
-}
-
-impl Default for SpatialPoolerConfig {
-    fn default() -> Self {
-        Self {
-            input_bits: 16384,
-            n_columns: 2048,
-            potential_radius: 1024,
-            potential_synapses: 40,
-            connected_threshold: 0.5,
-            syn_perm_active_inc: 0.04,
-            syn_perm_inactive_dec: 0.008,
-            sparsity: 0.02,
-            duty_cycle_period: 1000.0,
-            boost_strength: 1.0,
-            init_perm_span: 0.1,
-        }
-    }
-}
-
-pub struct SpatialPooler {
-    pub cfg: SpatialPoolerConfig,
-    pub columns: Vec<ProximalDendrite>,
-    /// Exponential moving average of "column was active" per step.
-    pub active_duty_cycle: Vec<f32>,
-    /// Exponential moving average of "overlap exceeded threshold" per step.
-    pub overlap_duty_cycle: Vec<f32>,
-    /// Boost factor per column.
-    pub boost: Vec<f32>,
-    rng: Xoshiro256PlusPlus,
-    iter_count: u64,
-}
-
-impl SpatialPooler {
-    pub fn new(cfg: SpatialPoolerConfig, seed: u64) -> Self {
-        assert!(cfg.input_bits >= cfg.potential_radius,
-            "input_bits ({}) must be >= potential_radius ({})",
-            cfg.input_bits, cfg.potential_radius);
-        assert!(cfg.potential_radius >= cfg.potential_synapses,
-            "potential_radius ({}) must be >= potential_synapses ({})",
-            cfg.potential_radius, cfg.potential_synapses);
-
-        let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed);
-
-        let mut columns = Vec::with_capacity(cfg.n_columns);
-        for _ in 0..cfg.n_columns {
-            // Sample `potential_radius` distinct input indices, then from those
-            // pick `potential_synapses` as the actual proximal synapses.
-            // Using partial Fisher-Yates via shuffle on a pool index range.
-            let mut pool: Vec<u32> = (0..cfg.input_bits as u32).collect();
-            // Efficient partial shuffle: swap the first `potential_radius`
-            // items with random items from the rest (Durstenfeld step).
-            for i in 0..cfg.potential_radius.min(pool.len()) {
-                let j = rng.gen_range(i..pool.len());
-                pool.swap(i, j);
-            }
-            let window = &mut pool[..cfg.potential_radius];
-            window.shuffle(&mut rng);
-            let mut inputs: Vec<u32> = window[..cfg.potential_synapses].to_vec();
-            inputs.sort_unstable();
-
-            let perms: Vec<f32> = (0..cfg.potential_synapses)
-                .map(|_| {
-                    let delta: f32 = rng.gen_range(-cfg.init_perm_span..cfg.init_perm_span);
-                    (cfg.connected_threshold + delta).clamp(0.0, 1.0)
-                })
-                .collect();
-
-            columns.push(ProximalDendrite { inputs, perms });
-        }
-
-        let n = cfg.n_columns;
-        Self {
-            cfg,
-            columns,
-            active_duty_cycle: vec![0.0; n],
-            overlap_duty_cycle: vec![0.0; n],
-            boost: vec![1.0; n],
-            rng,
-            iter_count: 0,
-        }
-    }
-
-    /// Process one step: compute overlaps, inhibit, learn (if `learn`), update
-    /// duty cycles and boosts. Returns the set of active column indices.
-    pub fn compute(&mut self, input: &[bool], learn: bool) -> Vec<u32> {
-        assert_eq!(input.len(), self.cfg.input_bits);
-
-        // 1) Overlap score per column (sum of CONNECTED synapses onto active inputs).
-        //    Also track raw overlap for the overlap-duty-cycle.
-        let n = self.cfg.n_columns;
-        let mut overlaps: Vec<f32> = vec![0.0; n];
-        let mut raw_overlaps: Vec<u32> = vec![0; n];
-
-        for (ci, col) in self.columns.iter().enumerate() {
-            let mut s: u32 = 0;
-            for (syn_i, &inp) in col.inputs.iter().enumerate() {
-                if input[inp as usize] && col.perms[syn_i] >= self.cfg.connected_threshold {
-                    s += 1;
-                }
-            }
-            raw_overlaps[ci] = s;
-            overlaps[ci] = (s as f32) * self.boost[ci];
-        }
-
-        // 2) Global k-WTA inhibition. Select top-k columns by boosted overlap.
-        let k = ((self.cfg.sparsity * n as f32).round() as usize).max(1);
-        let active: Vec<u32> = top_k(&overlaps, k);
-
-        // 3) Hebbian learning on active columns.
-        if learn {
-            for &ci in &active {
-                let col = &mut self.columns[ci as usize];
-                for (syn_i, &inp) in col.inputs.iter().enumerate() {
-                    if input[inp as usize] {
-                        col.perms[syn_i] =
-                            (col.perms[syn_i] + self.cfg.syn_perm_active_inc).min(1.0);
-                    } else {
-                        col.perms[syn_i] =
-                            (col.perms[syn_i] - self.cfg.syn_perm_inactive_dec).max(0.0);
-                    }
-                }
-            }
-        }
-
-        // 4) Update duty cycles (EMA with period T -> alpha = 1/T).
-        let period = self.cfg.duty_cycle_period.max(1.0);
-        let alpha = 1.0 / period;
-        // Column is "overlapping enough" if raw overlap >= stimulus_threshold.
-        // Numenta uses min_overlap; we use 1 as a conservative floor.
-        let stimulus_threshold = 1.0_f32;
-
-        // Mark active columns.
-        let mut active_mask = vec![false; n];
-        for &ci in &active {
-            active_mask[ci as usize] = true;
-        }
-
-        for i in 0..n {
-            let active_sample = if active_mask[i] { 1.0 } else { 0.0 };
-            let overlap_sample = if (raw_overlaps[i] as f32) >= stimulus_threshold {
-                1.0
-            } else {
-                0.0
-            };
-            self.active_duty_cycle[i] =
-                (1.0 - alpha) * self.active_duty_cycle[i] + alpha * active_sample;
-            self.overlap_duty_cycle[i] =
-                (1.0 - alpha) * self.overlap_duty_cycle[i] + alpha * overlap_sample;
-        }
-
-        // 5) Boost factor: b_i = exp(-boost_strength * (duty_i - mean_duty)).
-        //    Under-used columns (duty < mean) get boost > 1.
-        if learn && self.cfg.boost_strength > 0.0 {
-            let mean_duty: f32 =
-                self.active_duty_cycle.iter().sum::<f32>() / (n as f32);
-            for i in 0..n {
-                self.boost[i] =
-                    (-self.cfg.boost_strength * (self.active_duty_cycle[i] - mean_duty)).exp();
-            }
-
-            // 6) Permanence bump for chronically under-stimulated columns.
-            //    If overlap_duty_cycle[i] < min_pct_overlap * max_duty_in_neighborhood,
-            //    bump all permanences by syn_perm_active_inc * 0.1.
-            //    With global inhibition, "neighborhood" = all columns.
-            let max_overlap_duty = self
-                .overlap_duty_cycle
-                .iter()
-                .cloned()
-                .fold(0.0_f32, f32::max);
-            let min_pct_overlap_duty = 0.001_f32 * max_overlap_duty;
-            if max_overlap_duty > 0.0 {
-                for i in 0..n {
-                    if self.overlap_duty_cycle[i] < min_pct_overlap_duty {
-                        for p in &mut self.columns[i].perms {
-                            *p = (*p + self.cfg.syn_perm_active_inc * 0.1).min(1.0);
-                        }
-                    }
-                }
-            }
-        }
-
-        self.iter_count = self.iter_count.wrapping_add(1);
-        let _ = &mut self.rng; // suppress unused-mut when learn=false
-        active
-    }
-}
-
-/// Return the indices of the top-k values in `scores`.
-/// Ties broken by index order. Output is sorted ascending.
-fn top_k(scores: &[f32], k: usize) -> Vec<u32> {
-    if k == 0 {
-        return Vec::new();
-    }
-    let mut idx: Vec<u32> = (0..scores.len() as u32).collect();
-    // Partial sort: put top-k at the front by descending score.
-    // Use select_nth_unstable_by on (desc score, asc index).
-    idx.select_nth_unstable_by(k - 1, |&a, &b| {
-        let sa = scores[a as usize];
-        let sb = scores[b as usize];
-        // Reverse for descending.
-        match sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal) {
-            std::cmp::Ordering::Equal => a.cmp(&b),
-            ord => ord,
-        }
-    });
-    let mut winners: Vec<u32> = idx[..k].to_vec();
-    winners.sort_unstable();
-    winners
-}
-
-// ---------------------------------------------------------------------------
-// Tests
-// ---------------------------------------------------------------------------
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use rand::Rng;
-    use rand::SeedableRng;
-    use rand_xoshiro::Xoshiro256PlusPlus;
-
-    #[test]
-    fn sp_sparsity_exact_2pct() {
-        // BAMI says "top ~2%"; with 2048 columns that's round(0.02*2048) = 41.
-        // The SP must produce *exactly* that count, no more, no less, and with
-        // no duplicate indices.
-        let cfg = SpatialPoolerConfig::default();
-        let expected_k = (cfg.sparsity * cfg.n_columns as f32).round() as usize;
-        assert!(expected_k > 0);
-
-        let input_bits = cfg.input_bits;
-        let mut sp = SpatialPooler::new(cfg, 42);
-        let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
-
-        for _ in 0..100 {
-            // 2% sparse random input SDR.
-            let on_bits = (0.02 * input_bits as f32) as usize;
-            let mut sdr = vec![false; input_bits];
-            for _ in 0..on_bits {
-                let i = rng.gen_range(0..input_bits);
-                sdr[i] = true;
-            }
-            let active = sp.compute(&sdr, true);
-            assert_eq!(
-                active.len(),
-                expected_k,
-                "SP must emit exactly {expected_k} active columns"
-            );
-            let mut a = active.clone();
-            a.sort_unstable();
-            a.dedup();
-            assert_eq!(a.len(), expected_k);
-        }
-    }
-}
+//! Numenta BAMI-spec Spatial Pooler.
+//!
+//! Implements:
+//!   - 2048 (configurable) mini-columns with proximal dendrites
+//!   - `potential_synapses` (default 40) synapses per column sampled from
+//!     `potential_radius` (default 1024) random input bits
+//!   - Permanence in [0.0, 1.0] (f32), connected_threshold = 0.5
+//!   - syn_perm_active_inc = +0.04, syn_perm_inactive_dec = -0.008
+//!   - Global k-WTA inhibition (top `sparsity` fraction of columns)
+//!   - Boost factor with exponential duty-cycle tracking (Numenta formula)
+//!
+//! Reference: BAMI "Spatial Pooling Algorithm Details" (Numenta, 2017).
+
+use rand::Rng;
+use rand::SeedableRng;
+use rand::seq::SliceRandom;
+use rand_xoshiro::Xoshiro256PlusPlus;
+
+/// A single proximal dendrite: a sparse set of potential synapses onto
+/// specific input bit indices, with per-synapse permanence values.
+#[derive(Clone)]
+pub struct ProximalDendrite {
+    /// Indices into the input SDR.  Length == potential_synapses.
+    pub inputs: Vec<u32>,
+    /// Permanence for each potential synapse (same length as `inputs`).
+    pub perms: Vec<f32>,
+}
+
+pub struct SpatialPoolerConfig {
+    pub input_bits: usize,
+    pub n_columns: usize,
+    /// Size of the random input sample per column.
+    pub potential_radius: usize,
+    /// Number of potential synapses per column's proximal dendrite.
+    pub potential_synapses: usize,
+    pub connected_threshold: f32,
+    pub syn_perm_active_inc: f32,
+    pub syn_perm_inactive_dec: f32,
+    /// Target fraction of columns active per step (e.g. 0.02 for 2%).
+    pub sparsity: f32,
+    /// Duty cycle EMA period.
+    pub duty_cycle_period: f32,
+    /// Boost strength. Set to 0.0 to disable boosting.
+    pub boost_strength: f32,
+    /// Initial permanence span around the connected threshold.
+    pub init_perm_span: f32,
+}
+
+impl Default for SpatialPoolerConfig {
+    fn default() -> Self {
+        Self {
+            input_bits: 16384,
+            n_columns: 2048,
+            potential_radius: 1024,
+            potential_synapses: 40,
+            connected_threshold: 0.5,
+            syn_perm_active_inc: 0.04,
+            syn_perm_inactive_dec: 0.008,
+            sparsity: 0.02,
+            duty_cycle_period: 1000.0,
+            boost_strength: 1.0,
+            init_perm_span: 0.1,
+        }
+    }
+}
+
+pub struct SpatialPooler {
+    pub cfg: SpatialPoolerConfig,
+    pub columns: Vec<ProximalDendrite>,
+    /// Exponential moving average of "column was active" per step.
+    pub active_duty_cycle: Vec<f32>,
+    /// Exponential moving average of "overlap exceeded threshold" per step.
+    pub overlap_duty_cycle: Vec<f32>,
+    /// Boost factor per column.
+    pub boost: Vec<f32>,
+    rng: Xoshiro256PlusPlus,
+    iter_count: u64,
+}
+
+impl SpatialPooler {
+    pub fn new(cfg: SpatialPoolerConfig, seed: u64) -> Self {
+        assert!(cfg.input_bits >= cfg.potential_radius,
+            "input_bits ({}) must be >= potential_radius ({})",
+            cfg.input_bits, cfg.potential_radius);
+        assert!(cfg.potential_radius >= cfg.potential_synapses,
+            "potential_radius ({}) must be >= potential_synapses ({})",
+            cfg.potential_radius, cfg.potential_synapses);
+
+        let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed);
+
+        let mut columns = Vec::with_capacity(cfg.n_columns);
+        for _ in 0..cfg.n_columns {
+            // Sample `potential_radius` distinct input indices, then from those
+            // pick `potential_synapses` as the actual proximal synapses.
+            // Using partial Fisher-Yates via shuffle on a pool index range.
+            let mut pool: Vec<u32> = (0..cfg.input_bits as u32).collect();
+            // Efficient partial shuffle: swap the first `potential_radius`
+            // items with random items from the rest (Durstenfeld step).
+            for i in 0..cfg.potential_radius.min(pool.len()) {
+                let j = rng.gen_range(i..pool.len());
+                pool.swap(i, j);
+            }
+            let window = &mut pool[..cfg.potential_radius];
+            window.shuffle(&mut rng);
+            let mut inputs: Vec<u32> = window[..cfg.potential_synapses].to_vec();
+            inputs.sort_unstable();
+
+            let perms: Vec<f32> = (0..cfg.potential_synapses)
+                .map(|_| {
+                    let delta: f32 = rng.gen_range(-cfg.init_perm_span..cfg.init_perm_span);
+                    (cfg.connected_threshold + delta).clamp(0.0, 1.0)
+                })
+                .collect();
+
+            columns.push(ProximalDendrite { inputs, perms });
+        }
+
+        let n = cfg.n_columns;
+        Self {
+            cfg,
+            columns,
+            active_duty_cycle: vec![0.0; n],
+            overlap_duty_cycle: vec![0.0; n],
+            boost: vec![1.0; n],
+            rng,
+            iter_count: 0,
+        }
+    }
+
+    /// Process one step: compute overlaps, inhibit, learn (if `learn`), update
+    /// duty cycles and boosts. Returns the set of active column indices.
+    pub fn compute(&mut self, input: &[bool], learn: bool) -> Vec<u32> {
+        assert_eq!(input.len(), self.cfg.input_bits);
+
+        // 1) Overlap score per column (sum of CONNECTED synapses onto active inputs).
+        //    Also track raw overlap for the overlap-duty-cycle.
+        let n = self.cfg.n_columns;
+        let mut overlaps: Vec<f32> = vec![0.0; n];
+        let mut raw_overlaps: Vec<u32> = vec![0; n];
+
+        for (ci, col) in self.columns.iter().enumerate() {
+            let mut s: u32 = 0;
+            for (syn_i, &inp) in col.inputs.iter().enumerate() {
+                if input[inp as usize] && col.perms[syn_i] >= self.cfg.connected_threshold {
+                    s += 1;
+                }
+            }
+            raw_overlaps[ci] = s;
+            overlaps[ci] = (s as f32) * self.boost[ci];
+        }
+
+        // 2) Global k-WTA inhibition. Select top-k columns by boosted overlap.
+        let k = ((self.cfg.sparsity * n as f32).round() as usize).max(1);
+        let active: Vec<u32> = top_k(&overlaps, k);
+
+        // 3) Hebbian learning on active columns.
+        if learn {
+            for &ci in &active {
+                let col = &mut self.columns[ci as usize];
+                for (syn_i, &inp) in col.inputs.iter().enumerate() {
+                    if input[inp as usize] {
+                        col.perms[syn_i] =
+                            (col.perms[syn_i] + self.cfg.syn_perm_active_inc).min(1.0);
+                    } else {
+                        col.perms[syn_i] =
+                            (col.perms[syn_i] - self.cfg.syn_perm_inactive_dec).max(0.0);
+                    }
+                }
+            }
+        }
+
+        // 4) Update duty cycles (EMA with period T -> alpha = 1/T).
+        let period = self.cfg.duty_cycle_period.max(1.0);
+        let alpha = 1.0 / period;
+        // Column is "overlapping enough" if raw overlap >= stimulus_threshold.
+        // Numenta uses min_overlap; we use 1 as a conservative floor.
+        let stimulus_threshold = 1.0_f32;
+
+        // Mark active columns.
+        let mut active_mask = vec![false; n];
+        for &ci in &active {
+            active_mask[ci as usize] = true;
+        }
+
+        for i in 0..n {
+            let active_sample = if active_mask[i] { 1.0 } else { 0.0 };
+            let overlap_sample = if (raw_overlaps[i] as f32) >= stimulus_threshold {
+                1.0
+            } else {
+                0.0
+            };
+            self.active_duty_cycle[i] =
+                (1.0 - alpha) * self.active_duty_cycle[i] + alpha * active_sample;
+            self.overlap_duty_cycle[i] =
+                (1.0 - alpha) * self.overlap_duty_cycle[i] + alpha * overlap_sample;
+        }
+
+        // 5) Boost factor: b_i = exp(-boost_strength * (duty_i - mean_duty)).
+        //    Under-used columns (duty < mean) get boost > 1.
+        if learn && self.cfg.boost_strength > 0.0 {
+            let mean_duty: f32 =
+                self.active_duty_cycle.iter().sum::<f32>() / (n as f32);
+            for i in 0..n {
+                self.boost[i] =
+                    (-self.cfg.boost_strength * (self.active_duty_cycle[i] - mean_duty)).exp();
+            }
+
+            // 6) Permanence bump for chronically under-stimulated columns.
+            //    If overlap_duty_cycle[i] < min_pct_overlap * max_duty_in_neighborhood,
+            //    bump all permanences by syn_perm_active_inc * 0.1.
+            //    With global inhibition, "neighborhood" = all columns.
+            let max_overlap_duty = self
+                .overlap_duty_cycle
+                .iter()
+                .cloned()
+                .fold(0.0_f32, f32::max);
+            let min_pct_overlap_duty = 0.001_f32 * max_overlap_duty;
+            if max_overlap_duty > 0.0 {
+                for i in 0..n {
+                    if self.overlap_duty_cycle[i] < min_pct_overlap_duty {
+                        for p in &mut self.columns[i].perms {
+                            *p = (*p + self.cfg.syn_perm_active_inc * 0.1).min(1.0);
+                        }
+                    }
+                }
+            }
+        }
+
+        self.iter_count = self.iter_count.wrapping_add(1);
+        let _ = &mut self.rng; // suppress unused-mut when learn=false
+        active
+    }
+}
+
+/// Return the indices of the top-k values in `scores`.
+/// Ties broken by index order. Output is sorted ascending.
+fn top_k(scores: &[f32], k: usize) -> Vec<u32> {
+    if k == 0 {
+        return Vec::new();
+    }
+    let mut idx: Vec<u32> = (0..scores.len() as u32).collect();
+    // Partial sort: put top-k at the front by descending score.
+    // Use select_nth_unstable_by on (desc score, asc index).
+    idx.select_nth_unstable_by(k - 1, |&a, &b| {
+        let sa = scores[a as usize];
+        let sb = scores[b as usize];
+        // Reverse for descending.
+        match sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal) {
+            std::cmp::Ordering::Equal => a.cmp(&b),
+            ord => ord,
+        }
+    });
+    let mut winners: Vec<u32> = idx[..k].to_vec();
+    winners.sort_unstable();
+    winners
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::Rng;
+    use rand::SeedableRng;
+    use rand_xoshiro::Xoshiro256PlusPlus;
+
+    #[test]
+    fn sp_sparsity_exact_2pct() {
+        // BAMI says "top ~2%"; with 2048 columns that's round(0.02*2048) = 41.
+        // The SP must produce *exactly* that count, no more, no less, and with
+        // no duplicate indices.
+        let cfg = SpatialPoolerConfig::default();
+        let expected_k = (cfg.sparsity * cfg.n_columns as f32).round() as usize;
+        assert!(expected_k > 0);
+
+        let input_bits = cfg.input_bits;
+        let mut sp = SpatialPooler::new(cfg, 42);
+        let mut rng = Xoshiro256PlusPlus::seed_from_u64(7);
+
+        for _ in 0..100 {
+            // 2% sparse random input SDR.
+            let on_bits = (0.02 * input_bits as f32) as usize;
+            let mut sdr = vec![false; input_bits];
+            for _ in 0..on_bits {
+                let i = rng.gen_range(0..input_bits);
+                sdr[i] = true;
+            }
+            let active = sp.compute(&sdr, true);
+            assert_eq!(
+                active.len(),
+                expected_k,
+                "SP must emit exactly {expected_k} active columns"
+            );
+            let mut a = active.clone();
+            a.sort_unstable();
+            a.dedup();
+            assert_eq!(a.len(), expected_k);
+        }
+    }
+}

overlay/htm_rust/src/tm.rs

@@ -1,545 +1,545 @@
-//! Numenta BAMI-spec Temporal Memory.
-//!
-//! Key parameters (Numenta defaults):
-//!   - cells_per_column = 32
-//!   - max_segments_per_cell = 255
-//!   - max_synapses_per_segment = 32
-//!   - activation_threshold = 15  (CONNECTED synapses onto active cells)
-//!   - learning_threshold     = 13  (POTENTIAL synapses onto active cells)
-//!     (often called `minThreshold` / match threshold in BAMI)
-//!   - initial_permanence = 0.21
-//!   - connected_permanence = 0.50
-//!   - permanence_increment = 0.10
-//!   - permanence_decrement = 0.10
-//!   - predicted_segment_decrement = 0.10  (decay for segments that predicted
-//!     inactive columns; called `predictedSegmentDecrement` in BAMI)
-//!   - max_new_synapse_count = 20  (max synapses to grow on a new/reinforced seg)
-//!
-//! Algorithm (one step):
-//!   Given `active_columns` from the Spatial Pooler, and segment activity
-//!   caches `active_segments` and `matching_segments` computed *at the end of
-//!   the previous step*:
-//!
-//!   1. For each active column:
-//!        - If it contains any predicted cell (any cell with an active segment
-//!          from the previous depolarization), mark those cells active and
-//!          learn on the segment that predicted it.
-//!        - Else BURST the column: mark all cells in it active, and grow a new
-//!          segment on the best-matching cell in the column (or, if none,
-//!          on the cell with the fewest segments).
-//!   2. For every column that was predicted but did NOT become active
-//!      (matching segments on inactive columns), apply the
-//!      `predicted_segment_decrement` decay so spurious predictions fade.
-//!   3. Winner cells = active cells chosen for learning (1 per active column).
-//!   4. Compute segment activity for NEXT step:
-//!        - A segment's CONNECTED activity = #synapses with perm >= connected_perm
-//!          whose presynaptic cell is in `active_cells`. If >= activation_threshold
-//!          -> segment is "active" -> its cell is "predicted".
-//!        - A segment's POTENTIAL activity = #synapses whose presynaptic cell is
-//!          in `active_cells` (regardless of permanence). If >= learning_threshold
-//!          -> segment is "matching".
-//!
-//! Anomaly score = (active columns with no prior predicted cells)
-//!                  / (# active columns).
-
-use rand::Rng;
-use rand::SeedableRng;
-use rand_xoshiro::Xoshiro256PlusPlus;
-
-type CellIdx = u32;
-type SegmentIdx = u32;
-
-#[derive(Clone)]
-pub struct Synapse {
-    pub presynaptic_cell: CellIdx,
-    pub permanence: f32,
-}
-
-#[derive(Clone)]
-pub struct Segment {
-    pub cell: CellIdx,
-    pub synapses: Vec<Synapse>,
-    /// Cached counters; recomputed each step.
-    pub num_active_connected: u32,
-    pub num_active_potential: u32,
-    /// Simple "last iter touched" stat for least-used cell selection.
-    pub last_used_iteration: u64,
-}
-
-pub struct TemporalMemoryConfig {
-    pub n_columns: usize,
-    pub cells_per_column: usize,
-    pub activation_threshold: u32,
-    pub learning_threshold: u32,
-    pub initial_permanence: f32,
-    pub connected_permanence: f32,
-    pub permanence_increment: f32,
-    pub permanence_decrement: f32,
-    pub predicted_segment_decrement: f32,
-    pub max_segments_per_cell: usize,
-    pub max_synapses_per_segment: usize,
-    pub max_new_synapse_count: usize,
-}
-
-impl Default for TemporalMemoryConfig {
-    fn default() -> Self {
-        Self {
-            n_columns: 2048,
-            cells_per_column: 32,
-            activation_threshold: 15,
-            learning_threshold: 13,
-            initial_permanence: 0.21,
-            connected_permanence: 0.50,
-            permanence_increment: 0.10,
-            permanence_decrement: 0.10,
-            predicted_segment_decrement: 0.10,
-            max_segments_per_cell: 255,
-            max_synapses_per_segment: 32,
-            max_new_synapse_count: 20,
-        }
-    }
-}
-
-pub struct TemporalMemory {
-    pub cfg: TemporalMemoryConfig,
-    /// All segments in the region. Indexed by SegmentIdx.
-    pub segments: Vec<Segment>,
-    /// For each cell, the list of segments that belong to it.
-    pub cell_segments: Vec<Vec<SegmentIdx>>,
-    /// Active cells in the current step.
-    pub active_cells: Vec<bool>,
-    /// Winner cells (subset of active_cells, 1 per active column) for learning.
-    pub winner_cells: Vec<bool>,
-    /// Predictive cells for the current step = cells whose segment became
-    /// active at the end of the previous step.
-    pub predictive_cells: Vec<bool>,
-    /// Cached list of segment indices that were "active" last compute().
-    active_segments_prev: Vec<SegmentIdx>,
-    /// Cached list of segment indices that were "matching" last compute().
-    matching_segments_prev: Vec<SegmentIdx>,
-    rng: Xoshiro256PlusPlus,
-    iter_count: u64,
-}
-
-impl TemporalMemory {
-    pub fn new(cfg: TemporalMemoryConfig, seed: u64) -> Self {
-        let total = cfg.n_columns * cfg.cells_per_column;
-        Self {
-            cell_segments: vec![Vec::new(); total],
-            active_cells: vec![false; total],
-            winner_cells: vec![false; total],
-            predictive_cells: vec![false; total],
-            cfg,
-            segments: Vec::new(),
-            active_segments_prev: Vec::new(),
-            matching_segments_prev: Vec::new(),
-            rng: Xoshiro256PlusPlus::seed_from_u64(seed),
-            iter_count: 0,
-        }
-    }
-
-    pub fn reset(&mut self) {
-        for v in self.active_cells.iter_mut() { *v = false; }
-        for v in self.winner_cells.iter_mut() { *v = false; }
-        for v in self.predictive_cells.iter_mut() { *v = false; }
-        self.active_segments_prev.clear();
-        self.matching_segments_prev.clear();
-    }
-
-    #[inline]
-    fn col_of(&self, cell: CellIdx) -> usize {
-        (cell as usize) / self.cfg.cells_per_column
-    }
-
-    #[inline]
-    fn cells_in_col(&self, col: usize) -> std::ops::Range<CellIdx> {
-        let base = (col * self.cfg.cells_per_column) as CellIdx;
-        base..(base + self.cfg.cells_per_column as CellIdx)
-    }
-
-    /// Process one step.
-    ///
-    /// `active_columns` is the set of column indices activated by the Spatial
-    /// Pooler this step.  Returns the anomaly score in [0, 1].
-    pub fn compute(&mut self, active_columns: &[u32], learn: bool) -> f32 {
-        self.iter_count = self.iter_count.wrapping_add(1);
-
-        // Snapshot previous-step cell activity (for learning on segments).
-        let prev_active_cells = self.active_cells.clone();
-        let prev_winner_cells = self.winner_cells.clone();
-
-        // Move current "predictive" (computed at the end of the last step)
-        // into local variables; we'll overwrite predictive_cells later.
-        let predictive_prev = self.predictive_cells.clone();
-
-        // Group active segments and matching segments by column of their
-        // owning cell, for the columns that are active this step.
-        let n_cols = self.cfg.n_columns;
-
-        // active_segs_by_col[col] = segment indices whose cell is in col and
-        // which were "active" in the previous depolarization.
-        // matching_segs_by_col[col] = similarly for "matching".
-        let mut active_segs_by_col: Vec<Vec<SegmentIdx>> = vec![Vec::new(); n_cols];
-        let mut matching_segs_by_col: Vec<Vec<SegmentIdx>> = vec![Vec::new(); n_cols];
-        for &seg in &self.active_segments_prev {
-            let col = self.col_of(self.segments[seg as usize].cell);
-            active_segs_by_col[col].push(seg);
-        }
-        for &seg in &self.matching_segments_prev {
-            let col = self.col_of(self.segments[seg as usize].cell);
-            matching_segs_by_col[col].push(seg);
-        }
-
-        // Columns that are active this step (for O(1) lookup).
-        let mut active_col_mask = vec![false; n_cols];
-        for &c in active_columns { active_col_mask[c as usize] = true; }
-
-        // Zero out current cell activations.
-        for v in self.active_cells.iter_mut() { *v = false; }
-        for v in self.winner_cells.iter_mut() { *v = false; }
-
-        // Track anomaly.
-        let mut unpredicted_cols = 0u32;
-
-        // We'll collect (segment, learn_mode) pairs for segment reinforcement
-        // so we can batch-apply permanence adjustments using prev_active_cells.
-        // learn_mode: "reinforce_correctly_predicted", "punish_incorrectly_matched"
-        enum LearnOp {
-            Reinforce(SegmentIdx),       // correctly predicted
-            Grow {                        // bursting column: grow on chosen segment
-                segment: SegmentIdx,
-                #[allow(dead_code)]
-                winner_cell: CellIdx,
-            },
-            Punish(SegmentIdx),           // matching segment on inactive column
-        }
-        let mut ops: Vec<LearnOp> = Vec::new();
-
-        // ---- 1) Process active columns ----
-        for &col in active_columns {
-            let col = col as usize;
-            let active_segs = &active_segs_by_col[col];
-            if !active_segs.is_empty() {
-                // "Activate predicted column": each cell with an active segment
-                // becomes active and is a winner; reinforce that segment.
-                let mut seen_cells: Vec<CellIdx> = Vec::new();
-                for &seg_i in active_segs {
-                    let seg = &self.segments[seg_i as usize];
-                    let cell = seg.cell;
-                    if !seen_cells.contains(&cell) {
-                        self.active_cells[cell as usize] = true;
-                        self.winner_cells[cell as usize] = true;
-                        seen_cells.push(cell);
-                    }
-                    if learn {
-                        ops.push(LearnOp::Reinforce(seg_i));
-                    }
-                }
-            } else {
-                // ----- BURST -----
-                unpredicted_cols += 1;
-                for c in self.cells_in_col(col) {
-                    self.active_cells[c as usize] = true;
-                }
-                // Pick a winner cell + segment for learning.
-                if learn {
-                    let matching = &matching_segs_by_col[col];
-                    let (winner_cell, target_segment) = if !matching.is_empty() {
-                        // Best-matching segment = highest num_active_potential.
-                        let mut best = matching[0];
-                        let mut best_score = self.segments[best as usize].num_active_potential;
-                        for &s in &matching[1..] {
-                            let score = self.segments[s as usize].num_active_potential;
-                            if score > best_score {
-                                best_score = score;
-                                best = s;
-                            }
-                        }
-                        let wc = self.segments[best as usize].cell;
-                        (wc, Some(best))
-                    } else {
-                        // Least-used cell in column, then grow a new segment.
-                        let winner = self.least_used_cell(col);
-                        (winner, None)
-                    };
-                    self.winner_cells[winner_cell as usize] = true;
-                    let segment_id = match target_segment {
-                        Some(s) => s,
-                        None => {
-                            // Create a fresh empty segment on winner cell.
-                            self.create_segment(winner_cell)
-                        }
-                    };
-                    ops.push(LearnOp::Grow { segment: segment_id, winner_cell });
-                } else {
-                    // No learning: still pick some winner cell (arbitrary)
-                    // so downstream code that inspects winner_cells isn't empty.
-                    let matching = &matching_segs_by_col[col];
-                    let winner_cell = if !matching.is_empty() {
-                        self.segments[matching[0] as usize].cell
-                    } else {
-                        self.least_used_cell(col)
-                    };
-                    self.winner_cells[winner_cell as usize] = true;
-                }
-            }
-        }
-
-        // ---- 2) Punish matching segments on INACTIVE columns ----
-        if learn && self.cfg.predicted_segment_decrement > 0.0 {
-            for &seg_i in &self.matching_segments_prev {
-                let col = self.col_of(self.segments[seg_i as usize].cell);
-                if !active_col_mask[col] {
-                    ops.push(LearnOp::Punish(seg_i));
-                }
-            }
-        }
-
-        // ---- 3) Apply learning ----
-        if learn {
-            for op in ops {
-                match op {
-                    LearnOp::Reinforce(seg_i) => {
-                        self.reinforce_segment(seg_i, &prev_active_cells);
-                        // Optionally grow up to N new synapses to winner cells
-                        // of the previous step.
-                        self.grow_synapses_on_segment(seg_i, &prev_winner_cells);
-                    }
-                    LearnOp::Grow { segment, winner_cell: _ } => {
-                        self.reinforce_segment(segment, &prev_active_cells);
-                        self.grow_synapses_on_segment(segment, &prev_winner_cells);
-                    }
-                    LearnOp::Punish(seg_i) => {
-                        let dec = self.cfg.predicted_segment_decrement;
-                        for syn in &mut self.segments[seg_i as usize].synapses {
-                            if prev_active_cells[syn.presynaptic_cell as usize] {
-                                syn.permanence = (syn.permanence - dec).max(0.0);
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        // ---- 4) Compute segment activity & predictive cells for NEXT step ----
-        // We have to use the *current* active_cells (just set above).
-        let mut next_active_segs: Vec<SegmentIdx> = Vec::new();
-        let mut next_matching_segs: Vec<SegmentIdx> = Vec::new();
-        for v in self.predictive_cells.iter_mut() { *v = false; }
-
-        let conn = self.cfg.connected_permanence;
-        let act_thr = self.cfg.activation_threshold;
-        let learn_thr = self.cfg.learning_threshold;
-
-        for (seg_i, seg) in self.segments.iter_mut().enumerate() {
-            let mut n_conn: u32 = 0;
-            let mut n_pot: u32 = 0;
-            for syn in &seg.synapses {
-                if self.active_cells[syn.presynaptic_cell as usize] {
-                    n_pot += 1;
-                    if syn.permanence >= conn { n_conn += 1; }
-                }
-            }
-            seg.num_active_connected = n_conn;
-            seg.num_active_potential = n_pot;
-            if n_conn >= act_thr {
-                next_active_segs.push(seg_i as SegmentIdx);
-                self.predictive_cells[seg.cell as usize] = true;
-            }
-            if n_pot >= learn_thr {
-                next_matching_segs.push(seg_i as SegmentIdx);
-            }
-        }
-        self.active_segments_prev = next_active_segs;
-        self.matching_segments_prev = next_matching_segs;
-
-        // Keep predictive_prev unused-guard; we no longer need it but
-        // retained to document intent.
-        let _ = predictive_prev;
-
-        // Anomaly.
-        if active_columns.is_empty() {
-            0.0
-        } else {
-            (unpredicted_cols as f32) / (active_columns.len() as f32)
-        }
-    }
-
-    /// Reinforce synapses on `seg`: +inc if presynaptic is active last step,
-    /// -dec otherwise.
-    fn reinforce_segment(&mut self, seg_i: SegmentIdx, prev_active_cells: &[bool]) {
-        let inc = self.cfg.permanence_increment;
-        let dec = self.cfg.permanence_decrement;
-        let seg = &mut self.segments[seg_i as usize];
-        seg.last_used_iteration = self.iter_count;
-        for syn in &mut seg.synapses {
-            if prev_active_cells[syn.presynaptic_cell as usize] {
-                syn.permanence = (syn.permanence + inc).min(1.0);
-            } else {
-                syn.permanence = (syn.permanence - dec).max(0.0);
-            }
-        }
-    }
-
-    /// Grow up to `max_new_synapse_count - current_potential` new synapses
-    /// from previous winner cells that are not already connected to this seg.
-    fn grow_synapses_on_segment(
-        &mut self,
-        seg_i: SegmentIdx,
-        prev_winner_cells: &[bool],
-    ) {
-        let initial_perm = self.cfg.initial_permanence;
-        let cap = self.cfg.max_synapses_per_segment;
-        let max_new = self.cfg.max_new_synapse_count;
-
-        // Gather candidate cells (prev winners not already presynaptic to this seg).
-        let already: Vec<CellIdx> = self.segments[seg_i as usize]
-            .synapses
-            .iter()
-            .map(|s| s.presynaptic_cell)
-            .collect();
-        let mut candidates: Vec<CellIdx> = Vec::new();
-        for (cell_i, &b) in prev_winner_cells.iter().enumerate() {
-            if b && !already.contains(&(cell_i as CellIdx)) {
-                candidates.push(cell_i as CellIdx);
-            }
-        }
-
-        // How many can we add?
-        let current_len = self.segments[seg_i as usize].synapses.len();
-        let room = cap.saturating_sub(current_len);
-        let mut to_add = max_new.min(candidates.len()).min(room);
-
-        // Random sample without replacement from candidates.
-        while to_add > 0 {
-            let idx = self.rng.gen_range(0..candidates.len());
-            let pre = candidates.swap_remove(idx);
-            self.segments[seg_i as usize].synapses.push(Synapse {
-                presynaptic_cell: pre,
-                permanence: initial_perm,
-            });
-            to_add -= 1;
-        }
-    }
-
-    fn create_segment(&mut self, cell: CellIdx) -> SegmentIdx {
-        // Enforce per-cell segment cap by evicting least-recently-used segment
-        // if necessary.
-        let cell_segs = &mut self.cell_segments[cell as usize];
-        if cell_segs.len() >= self.cfg.max_segments_per_cell {
-            // Find LRU segment.
-            let (lru_pos, &lru_id) = cell_segs
-                .iter()
-                .enumerate()
-                .min_by_key(|(_, &sid)| self.segments[sid as usize].last_used_iteration)
-                .expect("cell_segs non-empty");
-            // Clear that segment in place and reuse its index.
-            self.segments[lru_id as usize].synapses.clear();
-            self.segments[lru_id as usize].num_active_connected = 0;
-            self.segments[lru_id as usize].num_active_potential = 0;
-            self.segments[lru_id as usize].last_used_iteration = self.iter_count;
-            // Keep at same position in cell_segs.
-            let _ = lru_pos;
-            return lru_id;
-        }
-
-        let new_id = self.segments.len() as SegmentIdx;
-        self.segments.push(Segment {
-            cell,
-            synapses: Vec::with_capacity(self.cfg.max_new_synapse_count),
-            num_active_connected: 0,
-            num_active_potential: 0,
-            last_used_iteration: self.iter_count,
-        });
-        cell_segs.push(new_id);
-        new_id
-    }
-
-    fn least_used_cell(&mut self, col: usize) -> CellIdx {
-        // Cell with the fewest segments; break ties randomly.
-        let mut min_segs = usize::MAX;
-        let mut candidates: Vec<CellIdx> = Vec::new();
-        for c in self.cells_in_col(col) {
-            let n = self.cell_segments[c as usize].len();
-            if n < min_segs {
-                min_segs = n;
-                candidates.clear();
-                candidates.push(c);
-            } else if n == min_segs {
-                candidates.push(c);
-            }
-        }
-        let idx = self.rng.gen_range(0..candidates.len());
-        candidates[idx]
-    }
-}
-
-// ---------------------------------------------------------------------------
-// Tests
-// ---------------------------------------------------------------------------
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::sp::{SpatialPooler, SpatialPoolerConfig};
-    use rand::Rng;
-    use rand::SeedableRng;
-    use rand_xoshiro::Xoshiro256PlusPlus;
-
-    #[test]
-    fn tm_learns_repeating_sequence() {
-        // Sequence A -> B -> C -> A -> B -> C -> ... should drive anomaly down.
-        let cfg = SpatialPoolerConfig::default();
-        let mut sp = SpatialPooler::new(cfg, 123);
-        let mut tm = TemporalMemory::new(TemporalMemoryConfig::default(), 456);
-
-        // Build 3 fixed random SDRs of 2% sparsity.
-        let mut rng = Xoshiro256PlusPlus::seed_from_u64(99);
-        let input_bits = sp.cfg.input_bits;
-        let make_sdr = |rng: &mut Xoshiro256PlusPlus| {
-            let mut v = vec![false; input_bits];
-            let on = (0.02 * input_bits as f32) as usize;
-            let mut placed = 0;
-            while placed < on {
-                let i = rng.gen_range(0..input_bits);
-                if !v[i] {
-                    v[i] = true;
-                    placed += 1;
-                }
-            }
-            v
-        };
-        let seqs = [make_sdr(&mut rng), make_sdr(&mut rng), make_sdr(&mut rng)];
-
-        // Warm up SP first so that columns are reliable for each symbol.
-        for _ in 0..200 {
-            for s in &seqs {
-                sp.compute(s, true);
-            }
-        }
-
-        // Reset TM so prediction state is clean.
-        tm.reset();
-
-        // Record anomaly over a window early and late.
-        let mut early_anoms: Vec<f32> = Vec::new();
-        let mut late_anoms: Vec<f32> = Vec::new();
-        for iter in 0..250 {
-            for s in &seqs {
-                let active = sp.compute(s, false);
-                let anomaly = tm.compute(&active, true);
-                if iter == 10 { early_anoms.push(anomaly); }
-                if iter == 249 { late_anoms.push(anomaly); }
-            }
-        }
-
-        let mean = |v: &[f32]| v.iter().sum::<f32>() / (v.len() as f32);
-        let early = mean(&early_anoms);
-        let late = mean(&late_anoms);
-        println!("early_anomaly={early}, late_anomaly={late}");
-        assert!(
-            late < 0.5 * early + 1e-6,
-            "late anomaly ({late}) should be < 0.5 * early anomaly ({early})"
-        );
-    }
-}
+//! Numenta BAMI-spec Temporal Memory.
+//!
+//! Key parameters (Numenta defaults):
+//!   - cells_per_column = 32
+//!   - max_segments_per_cell = 255
+//!   - max_synapses_per_segment = 32
+//!   - activation_threshold = 15  (CONNECTED synapses onto active cells)
+//!   - learning_threshold     = 13  (POTENTIAL synapses onto active cells)
+//!     (often called `minThreshold` / match threshold in BAMI)
+//!   - initial_permanence = 0.21
+//!   - connected_permanence = 0.50
+//!   - permanence_increment = 0.10
+//!   - permanence_decrement = 0.10
+//!   - predicted_segment_decrement = 0.10  (decay for segments that predicted
+//!     inactive columns; called `predictedSegmentDecrement` in BAMI)
+//!   - max_new_synapse_count = 20  (max synapses to grow on a new/reinforced seg)
+//!
+//! Algorithm (one step):
+//!   Given `active_columns` from the Spatial Pooler, and segment activity
+//!   caches `active_segments` and `matching_segments` computed *at the end of
+//!   the previous step*:
+//!
+//!   1. For each active column:
+//!        - If it contains any predicted cell (any cell with an active segment
+//!          from the previous depolarization), mark those cells active and
+//!          learn on the segment that predicted it.
+//!        - Else BURST the column: mark all cells in it active, and grow a new
+//!          segment on the best-matching cell in the column (or, if none,
+//!          on the cell with the fewest segments).
+//!   2. For every column that was predicted but did NOT become active
+//!      (matching segments on inactive columns), apply the
+//!      `predicted_segment_decrement` decay so spurious predictions fade.
+//!   3. Winner cells = active cells chosen for learning (1 per active column).
+//!   4. Compute segment activity for NEXT step:
+//!        - A segment's CONNECTED activity = #synapses with perm >= connected_perm
+//!          whose presynaptic cell is in `active_cells`. If >= activation_threshold
+//!          -> segment is "active" -> its cell is "predicted".
+//!        - A segment's POTENTIAL activity = #synapses whose presynaptic cell is
+//!          in `active_cells` (regardless of permanence). If >= learning_threshold
+//!          -> segment is "matching".
+//!
+//! Anomaly score = (active columns with no prior predicted cells)
+//!                  / (# active columns).
+
+use rand::Rng;
+use rand::SeedableRng;
+use rand_xoshiro::Xoshiro256PlusPlus;
+
+type CellIdx = u32;
+type SegmentIdx = u32;
+
+#[derive(Clone)]
+pub struct Synapse {
+    pub presynaptic_cell: CellIdx,
+    pub permanence: f32,
+}
+
+#[derive(Clone)]
+pub struct Segment {
+    pub cell: CellIdx,
+    pub synapses: Vec<Synapse>,
+    /// Cached counters; recomputed each step.
+    pub num_active_connected: u32,
+    pub num_active_potential: u32,
+    /// Simple "last iter touched" stat for least-used cell selection.
+    pub last_used_iteration: u64,
+}
+
+pub struct TemporalMemoryConfig {
+    pub n_columns: usize,
+    pub cells_per_column: usize,
+    pub activation_threshold: u32,
+    pub learning_threshold: u32,
+    pub initial_permanence: f32,
+    pub connected_permanence: f32,
+    pub permanence_increment: f32,
+    pub permanence_decrement: f32,
+    pub predicted_segment_decrement: f32,
+    pub max_segments_per_cell: usize,
+    pub max_synapses_per_segment: usize,
+    pub max_new_synapse_count: usize,
+}
+
+impl Default for TemporalMemoryConfig {
+    fn default() -> Self {
+        Self {
+            n_columns: 2048,
+            cells_per_column: 32,
+            activation_threshold: 15,
+            learning_threshold: 13,
+            initial_permanence: 0.21,
+            connected_permanence: 0.50,
+            permanence_increment: 0.10,
+            permanence_decrement: 0.10,
+            predicted_segment_decrement: 0.10,
+            max_segments_per_cell: 255,
+            max_synapses_per_segment: 32,
+            max_new_synapse_count: 20,
+        }
+    }
+}
+
+pub struct TemporalMemory {
+    pub cfg: TemporalMemoryConfig,
+    /// All segments in the region. Indexed by SegmentIdx.
+    pub segments: Vec<Segment>,
+    /// For each cell, the list of segments that belong to it.
+    pub cell_segments: Vec<Vec<SegmentIdx>>,
+    /// Active cells in the current step.
+    pub active_cells: Vec<bool>,
+    /// Winner cells (subset of active_cells, 1 per active column) for learning.
+    pub winner_cells: Vec<bool>,
+    /// Predictive cells for the current step = cells whose segment became
+    /// active at the end of the previous step.
+    pub predictive_cells: Vec<bool>,
+    /// Cached list of segment indices that were "active" last compute().
+    active_segments_prev: Vec<SegmentIdx>,
+    /// Cached list of segment indices that were "matching" last compute().
+    matching_segments_prev: Vec<SegmentIdx>,
+    rng: Xoshiro256PlusPlus,
+    iter_count: u64,
+}
+
+impl TemporalMemory {
+    pub fn new(cfg: TemporalMemoryConfig, seed: u64) -> Self {
+        let total = cfg.n_columns * cfg.cells_per_column;
+        Self {
+            cell_segments: vec![Vec::new(); total],
+            active_cells: vec![false; total],
+            winner_cells: vec![false; total],
+            predictive_cells: vec![false; total],
+            cfg,
+            segments: Vec::new(),
+            active_segments_prev: Vec::new(),
+            matching_segments_prev: Vec::new(),
+            rng: Xoshiro256PlusPlus::seed_from_u64(seed),
+            iter_count: 0,
+        }
+    }
+
+    pub fn reset(&mut self) {
+        for v in self.active_cells.iter_mut() { *v = false; }
+        for v in self.winner_cells.iter_mut() { *v = false; }
+        for v in self.predictive_cells.iter_mut() { *v = false; }
+        self.active_segments_prev.clear();
+        self.matching_segments_prev.clear();
+    }
+
+    #[inline]
+    fn col_of(&self, cell: CellIdx) -> usize {
+        (cell as usize) / self.cfg.cells_per_column
+    }
+
+    #[inline]
+    fn cells_in_col(&self, col: usize) -> std::ops::Range<CellIdx> {
+        let base = (col * self.cfg.cells_per_column) as CellIdx;
+        base..(base + self.cfg.cells_per_column as CellIdx)
+    }
+
+    /// Process one step.
+    ///
+    /// `active_columns` is the set of column indices activated by the Spatial
+    /// Pooler this step.  Returns the anomaly score in [0, 1].
+    pub fn compute(&mut self, active_columns: &[u32], learn: bool) -> f32 {
+        self.iter_count = self.iter_count.wrapping_add(1);
+
+        // Snapshot previous-step cell activity (for learning on segments).
+        let prev_active_cells = self.active_cells.clone();
+        let prev_winner_cells = self.winner_cells.clone();
+
+        // Move current "predictive" (computed at the end of the last step)
+        // into local variables; we'll overwrite predictive_cells later.
+        let predictive_prev = self.predictive_cells.clone();
+
+        // Group active segments and matching segments by column of their
+        // owning cell, for the columns that are active this step.
+        let n_cols = self.cfg.n_columns;
+
+        // active_segs_by_col[col] = segment indices whose cell is in col and
+        // which were "active" in the previous depolarization.
+        // matching_segs_by_col[col] = similarly for "matching".
+        let mut active_segs_by_col: Vec<Vec<SegmentIdx>> = vec![Vec::new(); n_cols];
+        let mut matching_segs_by_col: Vec<Vec<SegmentIdx>> = vec![Vec::new(); n_cols];
+        for &seg in &self.active_segments_prev {
+            let col = self.col_of(self.segments[seg as usize].cell);
+            active_segs_by_col[col].push(seg);
+        }
+        for &seg in &self.matching_segments_prev {
+            let col = self.col_of(self.segments[seg as usize].cell);
+            matching_segs_by_col[col].push(seg);
+        }
+
+        // Columns that are active this step (for O(1) lookup).
+        let mut active_col_mask = vec![false; n_cols];
+        for &c in active_columns { active_col_mask[c as usize] = true; }
+
+        // Zero out current cell activations.
+        for v in self.active_cells.iter_mut() { *v = false; }
+        for v in self.winner_cells.iter_mut() { *v = false; }
+
+        // Track anomaly.
+        let mut unpredicted_cols = 0u32;
+
+        // We'll collect (segment, learn_mode) pairs for segment reinforcement
+        // so we can batch-apply permanence adjustments using prev_active_cells.
+        // learn_mode: "reinforce_correctly_predicted", "punish_incorrectly_matched"
+        enum LearnOp {
+            Reinforce(SegmentIdx),       // correctly predicted
+            Grow {                        // bursting column: grow on chosen segment
+                segment: SegmentIdx,
+                #[allow(dead_code)]
+                winner_cell: CellIdx,
+            },
+            Punish(SegmentIdx),           // matching segment on inactive column
+        }
+        let mut ops: Vec<LearnOp> = Vec::new();
+
+        // ---- 1) Process active columns ----
+        for &col in active_columns {
+            let col = col as usize;
+            let active_segs = &active_segs_by_col[col];
+            if !active_segs.is_empty() {
+                // "Activate predicted column": each cell with an active segment
+                // becomes active and is a winner; reinforce that segment.
+                let mut seen_cells: Vec<CellIdx> = Vec::new();
+                for &seg_i in active_segs {
+                    let seg = &self.segments[seg_i as usize];
+                    let cell = seg.cell;
+                    if !seen_cells.contains(&cell) {
+                        self.active_cells[cell as usize] = true;
+                        self.winner_cells[cell as usize] = true;
+                        seen_cells.push(cell);
+                    }
+                    if learn {
+                        ops.push(LearnOp::Reinforce(seg_i));
+                    }
+                }
+            } else {
+                // ----- BURST -----
+                unpredicted_cols += 1;
+                for c in self.cells_in_col(col) {
+                    self.active_cells[c as usize] = true;
+                }
+                // Pick a winner cell + segment for learning.
+                if learn {
+                    let matching = &matching_segs_by_col[col];
+                    let (winner_cell, target_segment) = if !matching.is_empty() {
+                        // Best-matching segment = highest num_active_potential.
+                        let mut best = matching[0];
+                        let mut best_score = self.segments[best as usize].num_active_potential;
+                        for &s in &matching[1..] {
+                            let score = self.segments[s as usize].num_active_potential;
+                            if score > best_score {
+                                best_score = score;
+                                best = s;
+                            }
+                        }
+                        let wc = self.segments[best as usize].cell;
+                        (wc, Some(best))
+                    } else {
+                        // Least-used cell in column, then grow a new segment.
+                        let winner = self.least_used_cell(col);
+                        (winner, None)
+                    };
+                    self.winner_cells[winner_cell as usize] = true;
+                    let segment_id = match target_segment {
+                        Some(s) => s,
+                        None => {
+                            // Create a fresh empty segment on winner cell.
+                            self.create_segment(winner_cell)
+                        }
+                    };
+                    ops.push(LearnOp::Grow { segment: segment_id, winner_cell });
+                } else {
+                    // No learning: still pick some winner cell (arbitrary)
+                    // so downstream code that inspects winner_cells isn't empty.
+                    let matching = &matching_segs_by_col[col];
+                    let winner_cell = if !matching.is_empty() {
+                        self.segments[matching[0] as usize].cell
+                    } else {
+                        self.least_used_cell(col)
+                    };
+                    self.winner_cells[winner_cell as usize] = true;
+                }
+            }
+        }
+
+        // ---- 2) Punish matching segments on INACTIVE columns ----
+        if learn && self.cfg.predicted_segment_decrement > 0.0 {
+            for &seg_i in &self.matching_segments_prev {
+                let col = self.col_of(self.segments[seg_i as usize].cell);
+                if !active_col_mask[col] {
+                    ops.push(LearnOp::Punish(seg_i));
+                }
+            }
+        }
+
+        // ---- 3) Apply learning ----
+        if learn {
+            for op in ops {
+                match op {
+                    LearnOp::Reinforce(seg_i) => {
+                        self.reinforce_segment(seg_i, &prev_active_cells);
+                        // Optionally grow up to N new synapses to winner cells
+                        // of the previous step.
+                        self.grow_synapses_on_segment(seg_i, &prev_winner_cells);
+                    }
+                    LearnOp::Grow { segment, winner_cell: _ } => {
+                        self.reinforce_segment(segment, &prev_active_cells);
+                        self.grow_synapses_on_segment(segment, &prev_winner_cells);
+                    }
+                    LearnOp::Punish(seg_i) => {
+                        let dec = self.cfg.predicted_segment_decrement;
+                        for syn in &mut self.segments[seg_i as usize].synapses {
+                            if prev_active_cells[syn.presynaptic_cell as usize] {
+                                syn.permanence = (syn.permanence - dec).max(0.0);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        // ---- 4) Compute segment activity & predictive cells for NEXT step ----
+        // We have to use the *current* active_cells (just set above).
+        let mut next_active_segs: Vec<SegmentIdx> = Vec::new();
+        let mut next_matching_segs: Vec<SegmentIdx> = Vec::new();
+        for v in self.predictive_cells.iter_mut() { *v = false; }
+
+        let conn = self.cfg.connected_permanence;
+        let act_thr = self.cfg.activation_threshold;
+        let learn_thr = self.cfg.learning_threshold;
+
+        for (seg_i, seg) in self.segments.iter_mut().enumerate() {
+            let mut n_conn: u32 = 0;
+            let mut n_pot: u32 = 0;
+            for syn in &seg.synapses {
+                if self.active_cells[syn.presynaptic_cell as usize] {
+                    n_pot += 1;
+                    if syn.permanence >= conn { n_conn += 1; }
+                }
+            }
+            seg.num_active_connected = n_conn;
+            seg.num_active_potential = n_pot;
+            if n_conn >= act_thr {
+                next_active_segs.push(seg_i as SegmentIdx);
+                self.predictive_cells[seg.cell as usize] = true;
+            }
+            if n_pot >= learn_thr {
+                next_matching_segs.push(seg_i as SegmentIdx);
+            }
+        }
+        self.active_segments_prev = next_active_segs;
+        self.matching_segments_prev = next_matching_segs;
+
+        // Keep predictive_prev unused-guard; we no longer need it but
+        // retained to document intent.
+        let _ = predictive_prev;
+
+        // Anomaly.
+        if active_columns.is_empty() {
+            0.0
+        } else {
+            (unpredicted_cols as f32) / (active_columns.len() as f32)
+        }
+    }
+
+    /// Reinforce synapses on `seg`: +inc if presynaptic is active last step,
+    /// -dec otherwise.
+    fn reinforce_segment(&mut self, seg_i: SegmentIdx, prev_active_cells: &[bool]) {
+        let inc = self.cfg.permanence_increment;
+        let dec = self.cfg.permanence_decrement;
+        let seg = &mut self.segments[seg_i as usize];
+        seg.last_used_iteration = self.iter_count;
+        for syn in &mut seg.synapses {
+            if prev_active_cells[syn.presynaptic_cell as usize] {
+                syn.permanence = (syn.permanence + inc).min(1.0);
+            } else {
+                syn.permanence = (syn.permanence - dec).max(0.0);
+            }
+        }
+    }
+
+    /// Grow up to `max_new_synapse_count - current_potential` new synapses
+    /// from previous winner cells that are not already connected to this seg.
+    fn grow_synapses_on_segment(
+        &mut self,
+        seg_i: SegmentIdx,
+        prev_winner_cells: &[bool],
+    ) {
+        let initial_perm = self.cfg.initial_permanence;
+        let cap = self.cfg.max_synapses_per_segment;
+        let max_new = self.cfg.max_new_synapse_count;
+
+        // Gather candidate cells (prev winners not already presynaptic to this seg).
+        let already: Vec<CellIdx> = self.segments[seg_i as usize]
+            .synapses
+            .iter()
+            .map(|s| s.presynaptic_cell)
+            .collect();
+        let mut candidates: Vec<CellIdx> = Vec::new();
+        for (cell_i, &b) in prev_winner_cells.iter().enumerate() {
+            if b && !already.contains(&(cell_i as CellIdx)) {
+                candidates.push(cell_i as CellIdx);
+            }
+        }
+
+        // How many can we add?
+        let current_len = self.segments[seg_i as usize].synapses.len();
+        let room = cap.saturating_sub(current_len);
+        let mut to_add = max_new.min(candidates.len()).min(room);
+
+        // Random sample without replacement from candidates.
+        while to_add > 0 {
+            let idx = self.rng.gen_range(0..candidates.len());
+            let pre = candidates.swap_remove(idx);
+            self.segments[seg_i as usize].synapses.push(Synapse {
+                presynaptic_cell: pre,
+                permanence: initial_perm,
+            });
+            to_add -= 1;
+        }
+    }
+
+    fn create_segment(&mut self, cell: CellIdx) -> SegmentIdx {
+        // Enforce per-cell segment cap by evicting least-recently-used segment
+        // if necessary.
+        let cell_segs = &mut self.cell_segments[cell as usize];
+        if cell_segs.len() >= self.cfg.max_segments_per_cell {
+            // Find LRU segment.
+            let (lru_pos, &lru_id) = cell_segs
+                .iter()
+                .enumerate()
+                .min_by_key(|(_, &sid)| self.segments[sid as usize].last_used_iteration)
+                .expect("cell_segs non-empty");
+            // Clear that segment in place and reuse its index.
+            self.segments[lru_id as usize].synapses.clear();
+            self.segments[lru_id as usize].num_active_connected = 0;
+            self.segments[lru_id as usize].num_active_potential = 0;
+            self.segments[lru_id as usize].last_used_iteration = self.iter_count;
+            // Keep at same position in cell_segs.
+            let _ = lru_pos;
+            return lru_id;
+        }
+
+        let new_id = self.segments.len() as SegmentIdx;
+        self.segments.push(Segment {
+            cell,
+            synapses: Vec::with_capacity(self.cfg.max_new_synapse_count),
+            num_active_connected: 0,
+            num_active_potential: 0,
+            last_used_iteration: self.iter_count,
+        });
+        cell_segs.push(new_id);
+        new_id
+    }
+
+    fn least_used_cell(&mut self, col: usize) -> CellIdx {
+        // Cell with the fewest segments; break ties randomly.
+        let mut min_segs = usize::MAX;
+        let mut candidates: Vec<CellIdx> = Vec::new();
+        for c in self.cells_in_col(col) {
+            let n = self.cell_segments[c as usize].len();
+            if n < min_segs {
+                min_segs = n;
+                candidates.clear();
+                candidates.push(c);
+            } else if n == min_segs {
+                candidates.push(c);
+            }
+        }
+        let idx = self.rng.gen_range(0..candidates.len());
+        candidates[idx]
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::sp::{SpatialPooler, SpatialPoolerConfig};
+    use rand::Rng;
+    use rand::SeedableRng;
+    use rand_xoshiro::Xoshiro256PlusPlus;
+
+    #[test]
+    fn tm_learns_repeating_sequence() {
+        // Sequence A -> B -> C -> A -> B -> C -> ... should drive anomaly down.
+        let cfg = SpatialPoolerConfig::default();
+        let mut sp = SpatialPooler::new(cfg, 123);
+        let mut tm = TemporalMemory::new(TemporalMemoryConfig::default(), 456);
+
+        // Build 3 fixed random SDRs of 2% sparsity.
+        let mut rng = Xoshiro256PlusPlus::seed_from_u64(99);
+        let input_bits = sp.cfg.input_bits;
+        let make_sdr = |rng: &mut Xoshiro256PlusPlus| {
+            let mut v = vec![false; input_bits];
+            let on = (0.02 * input_bits as f32) as usize;
+            let mut placed = 0;
+            while placed < on {
+                let i = rng.gen_range(0..input_bits);
+                if !v[i] {
+                    v[i] = true;
+                    placed += 1;
+                }
+            }
+            v
+        };
+        let seqs = [make_sdr(&mut rng), make_sdr(&mut rng), make_sdr(&mut rng)];
+
+        // Warm up SP first so that columns are reliable for each symbol.
+        for _ in 0..200 {
+            for s in &seqs {
+                sp.compute(s, true);
+            }
+        }
+
+        // Reset TM so prediction state is clean.
+        tm.reset();
+
+        // Record anomaly over a window early and late.
+        let mut early_anoms: Vec<f32> = Vec::new();
+        let mut late_anoms: Vec<f32> = Vec::new();
+        for iter in 0..250 {
+            for s in &seqs {
+                let active = sp.compute(s, false);
+                let anomaly = tm.compute(&active, true);
+                if iter == 10 { early_anoms.push(anomaly); }
+                if iter == 249 { late_anoms.push(anomaly); }
+            }
+        }
+
+        let mean = |v: &[f32]| v.iter().sum::<f32>() / (v.len() as f32);
+        let early = mean(&early_anoms);
+        let late = mean(&late_anoms);
+        println!("early_anomaly={early}, late_anomaly={late}");
+        assert!(
+            late < 0.5 * early + 1e-6,
+            "late anomaly ({late}) should be < 0.5 * early anomaly ({early})"
+        );
+    }
+}

overlay/htm_rust/uv.lock

@@ -1,8 +1,8 @@
-version = 1
-revision = 3
-requires-python = ">=3.11"
-
-[[package]]
-name = "htm-rust"
-version = "0.1.0"
-source = { editable = "." }
+version = 1
+revision = 3
+requires-python = ">=3.11"
+
+[[package]]
+name = "htm-rust"
+version = "0.1.0"
+source = { editable = "." }

overlay/hydra/__init__.py

@@ -1,31 +1,31 @@
-"""HYDRA training package.
-
-Thin facade re-exporting the public API used by train.py, the test suite,
-and external research scripts. Imports are lazy where possible to keep
-`import hydra` cheap (prepare.py and mamba-ssm are the heavy deps).
-"""
-
-from hydra.config import PostSemClawConfig
-from hydra.engram import GPUEngram
-from hydra.model import PostSemClawModel, norm
-from hydra.optimizer import MuonAdamW, adamw_step_fused, muon_step_fused
-
-# config_from_dict is imported lazily (via attribute access on hydra.training)
-# to keep `import hydra` cheap; re-export here for convenience.
-def __getattr__(name: str):
-    if name == "config_from_dict":
-        from hydra.training import config_from_dict as _cfd
-        return _cfd
-    raise AttributeError(name)
-
-
-__all__ = [
-    "PostSemClawConfig",
-    "GPUEngram",
-    "PostSemClawModel",
-    "norm",
-    "MuonAdamW",
-    "adamw_step_fused",
-    "muon_step_fused",
-    "config_from_dict",
-]
+"""HYDRA training package.
+
+Thin facade re-exporting the public API used by train.py, the test suite,
+and external research scripts. Imports are lazy where possible to keep
+`import hydra` cheap (prepare.py and mamba-ssm are the heavy deps).
+"""
+
+from hydra.config import PostSemClawConfig
+from hydra.engram import GPUEngram
+from hydra.model import PostSemClawModel, norm
+from hydra.optimizer import MuonAdamW, adamw_step_fused, muon_step_fused
+
+# config_from_dict is imported lazily (via attribute access on hydra.training)
+# to keep `import hydra` cheap; re-export here for convenience.
+def __getattr__(name: str):
+    if name == "config_from_dict":
+        from hydra.training import config_from_dict as _cfd
+        return _cfd
+    raise AttributeError(name)
+
+
+__all__ = [
+    "PostSemClawConfig",
+    "GPUEngram",
+    "PostSemClawModel",
+    "norm",
+    "MuonAdamW",
+    "adamw_step_fused",
+    "muon_step_fused",
+    "config_from_dict",
+]

overlay/hydra/config.py

@@ -1,225 +1,225 @@
-"""HYDRA training configuration — dataclass + env-var constants.
-
-Extracted from the monolithic train.py as part of W1 modularization. All
-env-var reads and the PostSemClawConfig dataclass live here. The training
-body imports these constants; zero behavior change from the extraction.
-"""
-
-from __future__ import annotations
-
-import os
-from dataclasses import dataclass, field
-
-
-def _parse_hyena_layers_env() -> tuple[int, ...]:
-    """Parse HYDRA_HYENA_LAYERS env var into a sorted tuple of layer indices.
-
-    Used as the default_factory for PostSemClawConfig.hyena_layers so a fresh
-    config construction reads the current env var, but once constructed the
-    value is first-class and travels with checkpoints (see asdict(config) in
-    save_ckpt). Ckpt-load sets the dataclass field explicitly, overriding the
-    env-var default.
-
-    Returns empty tuple when env var is unset/empty (byte-identical to
-    pre-port behavior: no Hyena layers).
-    """
-    raw = os.environ.get("HYDRA_HYENA_LAYERS", "")
-    if not raw:
-        return ()
-    return tuple(sorted({int(s.strip()) for s in raw.split(",") if s.strip()}))
-
-
-def _parse_gdn_layers_env() -> tuple[int, ...]:
-    """Parse HYDRA_GDN_LAYERS env var into a sorted tuple of layer indices.
-
-    Same contract as _parse_hyena_layers_env: layers whose index is listed
-    here use GatedDeltaNet (fla.layers.GatedDeltaNet) as a drop-in
-    replacement for Mamba3. Empty tuple = no GDN layers (byte-identical
-    to baseline).
-    """
-    raw = os.environ.get("HYDRA_GDN_LAYERS", "")
-    if not raw:
-        return ()
-    return tuple(sorted({int(s.strip()) for s in raw.split(",") if s.strip()}))
-
-# ---------------------------------------------------------------------------
-# CUDA env — set before importing torch in entry point. Kept here so any
-# module that `from hydra.config import ...` also benefits (import order is
-# top-down in Python, and train.py used to set these at module top).
-# ---------------------------------------------------------------------------
-os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
-if "/usr/local/cuda/bin" not in os.environ.get("PATH", ""):
-    os.environ["PATH"] = "/usr/local/cuda/bin:" + os.environ.get("PATH", "")
-os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
-
-
-# ---------------------------------------------------------------------------
-# Model Configuration
-# ---------------------------------------------------------------------------
-
-@dataclass
-class PostSemClawConfig:
-    """Full-architecture model config. Defaults reflect Phase-1 baseline;
-    the training entry overrides d_model/n_layer/etc. from env vars."""
-    # Sequence
-    sequence_len: int = 2048
-    vocab_size: int = 8192  # Must match prepare.py VOCAB_SIZE
-
-    # Mamba-3 SSM
-    n_layer: int = 6
-    d_model: int = 384
-    d_state: int = 64       # SSM state dimension
-    headdim: int = 48       # head dimension for SSM
-    n_heads: int = 8        # d_model // headdim
-    expand: int = 2         # inner_dim = expand * d_model
-
-    # Engram (conditional memory with Hebbian writes)
-    engram_n_columns: int = 4096
-    engram_key_dim: int = 64
-    engram_layer_idx: int = 1  # which layer gets engram (0-indexed, mid-layer)
-
-    # SemanticFoldingSDR (offline retina with STE; no-bypass, runs every step)
-    sdr_n_bits: int = 16384          # retina width
-    # Default 327 = 2% sparsity (Webber/Numenta canonical). Override with
-    # HYDRA_SDR_TARGET_ACTIVE env var; value MUST match subsystems/sdr_retina.py
-    # TARGET_ACTIVE (same env var is read there, so just setting it once works).
-    sdr_target_active: int = int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327"))
-    sdr_delta_rank: int = 32         # low-rank STE delta rank
-    sdr_som_warmup: int = 500
-    sdr_som_interval: int = 100
-
-    # HTMLayer (Rust-backed, Hebbian; no-bypass, runs every step)
-    htm_n_columns: int = 2048
-    htm_cells_per_column: int = 32
-
-    # Hyena supplement layer indices (sorted tuple). Defaults to the
-    # HYDRA_HYENA_LAYERS env var at config-construction time, but once
-    # persisted in a checkpoint the value is first-class and survives even
-    # when the env var is unset at resume time. This fixes the ckpt-reload
-    # crash path where a model trained with `HYDRA_HYENA_LAYERS=3,7` saves
-    # HyenaBlock params but a fresh process without the env var would try
-    # to build a pure-Mamba3 architecture and reject the state_dict as
-    # `Missing/Unexpected key(s)`.
-    hyena_layers: tuple[int, ...] = field(default_factory=_parse_hyena_layers_env)
-
-    # GatedDeltaNet supplement layer indices (sorted tuple). Same semantics
-    # as hyena_layers — a layer index listed here uses GDNBlock (fla-backed
-    # Gated DeltaNet) instead of Mamba3. Selections are mutually exclusive
-    # with hyena_layers at construction time (hyena wins on overlap; the
-    # model loop checks hyena first).
-    gdn_layers: tuple[int, ...] = field(default_factory=_parse_gdn_layers_env)
-
-    # Label smoothing + Z-loss
-    label_smoothing: float = 0.0   # disabled: any smoothing hurts in 5-min budget
-    z_loss_weight: float = 1e-4
-
-
-# ---------------------------------------------------------------------------
-# Hyperparameters (autoresearch agent modifies these via env vars)
-# ---------------------------------------------------------------------------
-
-# Model architecture
-D_MODEL = int(os.environ.get("HYDRA_D_MODEL", "256"))
-N_LAYER = int(os.environ.get("HYDRA_N_LAYER", "4"))
-D_STATE = int(os.environ.get("HYDRA_D_STATE", "64"))
-HEADDIM = int(os.environ.get("HYDRA_HEADDIM", "32"))
-N_HEADS = D_MODEL // HEADDIM
-EXPAND = int(os.environ.get("HYDRA_EXPAND", "2"))
-
-# Engram
-ENGRAM_N_COLUMNS = int(os.environ.get("HYDRA_ENGRAM_N_COLUMNS", "1024"))
-ENGRAM_KEY_DIM = 64
-ENGRAM_LAYER_IDX = int(os.environ.get("HYDRA_ENGRAM_LAYER_IDX", "1"))
-
-# Optimization
-DEVICE_BATCH_SIZE = int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
-TOTAL_BATCH_SIZE = int(os.environ.get("HYDRA_TOTAL_BATCH", "32768"))
-MATRIX_LR = float(os.environ.get("HYDRA_MATRIX_LR", "0.12"))
-EMBEDDING_LR = float(os.environ.get("HYDRA_EMBED_LR", "1.0"))
-UNEMBEDDING_LR = float(os.environ.get("HYDRA_UNEMBED_LR", "0.005"))
-# Scalar/vector params include Hyena implicit-filter vectors, norms, gate/bias
-# terms, and SDR delta_u/delta_v.  They are AdamW-scaled by d_model and can be
-# the hidden instability path when the high-throughput HF recipe pushes a large
-# device batch for hours.  Keep the historical default, but make it controllable
-# from launch scripts so cloud jobs can cool scalars without editing code.
-SCALAR_LR = float(os.environ.get("HYDRA_SCALAR_LR", "0.5"))
-WEIGHT_DECAY = float(os.environ.get("HYDRA_WEIGHT_DECAY", "0.01"))
-ADAM_BETAS = (0.9, 0.95)
-WARMUP_RATIO = float(os.environ.get("HYDRA_WARMUP_RATIO", "0.0"))
-WARMDOWN_RATIO = 0.5
-FINAL_LR_FRAC = float(os.environ.get("HYDRA_LR_MIN_MULT", "0.0"))
-
-# Runtime
-SEED = int(os.environ.get("HYDRA_SEED", "42"))
-# BF16 TFLOPS peak (RTX 3060=25.5, A100 SXM4=312, H100 SXM5=989)
-GPU_BF16_PEAK_FLOPS = float(os.environ.get("HYDRA_GPU_BF16_TFLOPS", "25.5")) * 1e12
-
-# Loss / inference knobs read by the model
-CE_CHUNK = int(os.environ.get("HYDRA_CE_CHUNK", "1024"))
-DROPOUT = float(os.environ.get("HYDRA_DROPOUT", "0.2"))
-FUSED_ADAMW = os.environ.get("HYDRA_FUSED_ADAMW", "1") == "1"
-
-# ---------------------------------------------------------------------------
-# Learnability knobs (all OFF by default — zero behavior change unless set)
-# ---------------------------------------------------------------------------
-# 1) Multi-Token Prediction (Llama-3 style). K=1 disables (next-1 only). K=4
-#    adds 3 extra weight-tied heads; loss = mean of K position-shifted CEs.
-MTP_K = int(os.environ.get("HYDRA_MTP_K", "1"))
-# 2) Exponential Moving Average of model weights (decay=0.999). Saves an
-#    additional latest_ema.pt at the end of training.
-USE_EMA = os.environ.get("HYDRA_USE_EMA", "0") == "1"
-EMA_DECAY = float(os.environ.get("HYDRA_EMA_DECAY", "0.999"))
-# 3) Gradient checkpointing on Mamba3 block forward. Trades ~30% compute for
-#    ~40% activation memory savings — lets you push B upward on a 3060.
-GRAD_CKPT = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
-# 4) Doc-separator masking in packed sequences: at every packed-BOS position
-#    in the targets tensor, mask the loss (ignore_index=-1) so the model is
-#    not forced to predict doc B from doc A's context.
-DOC_SEP_MASK = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
-# 5) Stop-gradient on HTM state (belt-and-braces: htm_rust already runs under
-#    torch.no_grad() so the tensor returned has requires_grad=False; this
-#    simply detaches explicitly to harden graph hygiene against future refactors).
-HTM_STOP_GRAD = os.environ.get("HYDRA_HTM_STOP_GRAD", "0") == "1"
-# 6) Output entropy penalty: loss += -lambda * H(softmax(logits)). Negative
-#    entropy penalizes peaked distributions and breaks repetition loops.
-ENTROPY_PENALTY = float(os.environ.get("HYDRA_ENTROPY_PENALTY", "0.0"))
-# 7) Curriculum: first N optimizer steps use short seq_len, then switch to
-#    full. 0 disables (no curriculum).
-CURRICULUM_SHORT_STEPS = int(os.environ.get("HYDRA_CURRICULUM_SHORT_STEPS", "0"))
-CURRICULUM_SHORT_SEQ_LEN = int(os.environ.get("HYDRA_CURRICULUM_SHORT_SEQ_LEN", "256"))
-
-# ---------------------------------------------------------------------------
-# Hyena supplement (additional block type for selected layer indices).
-# Hyena replaces Mamba3 at the specified layer indices while all other layers
-# remain Mamba3. Empty string (default) → no Hyena layers, byte-identical to
-# pre-port behavior.
-#   HYDRA_HYENA_LAYERS       "3,7"  — comma-separated 0-indexed layer ids
-#   HYDRA_HYENA_ORDER         2     — Hyena recurrence order (>= 2)
-#   HYDRA_HYENA_FILTER_DIM    64    — implicit-filter MLP hidden width
-# Hyena reference: https://arxiv.org/pdf/2302.10866.pdf (HazyResearch/safari).
-# ---------------------------------------------------------------------------
-HYENA_LAYERS = os.environ.get("HYDRA_HYENA_LAYERS", "")
-HYENA_ORDER = int(os.environ.get("HYDRA_HYENA_ORDER", "2"))
-HYENA_FILTER_DIM = int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64"))
-# Filter-rfft cache modes (see subsystems/hyena_pure.py):
-#   HYDRA_HYENA_FILTER_CACHE=1 — eval-only cache. Safe under torch.no_grad()
-#       where PyTorch never saves intermediate tensors. Off by default.
-#   HYDRA_HYENA_TRAIN_CACHE=1  — training-safe cache using a deferred
-#       gradient pattern. Cuts the implicit filter MLP forward to ONCE per
-#       optimizer step regardless of grad-accumulation factor. Requires the
-#       training loop (see hydra/lightning_module.py::optimizer_step) to
-#       call `model.flush_hyena_pending_grads()` before optimizer.step().
-#       Off by default.
-HYENA_FILTER_CACHE = os.environ.get("HYDRA_HYENA_FILTER_CACHE", "0") == "1"
-HYENA_TRAIN_CACHE = os.environ.get("HYDRA_HYENA_TRAIN_CACHE", "0") == "1"
-
-# Factual eval knobs
-FACTUAL_SAMPLES = int(os.environ.get("HYDRA_FACTUAL_SAMPLES", "3"))
-FACTUAL_BATCH = int(os.environ.get("HYDRA_FACTUAL_BATCH", "32"))
-# F6 (partial): Full incremental SSM decode integration deferred — would require
-# threading mamba_ssm InferenceParams through PostSemClawModel.forward and all
-# auxiliary subsystems (HTM, SDR, Engram) which currently run full-sequence each
-# call. As a stopgap we reduce default from 16 -> 4 so the per-prompt cost is
-# quartered (each gen-tok does a full re-encode of ctx+k tokens). Override with
-# HYDRA_FACTUAL_GEN_TOKENS to restore prior behavior. See docs/OPTIMIZATION_PLAN.md.
-FACTUAL_GEN_TOKENS = int(os.environ.get("HYDRA_FACTUAL_GEN_TOKENS", "2"))
+"""HYDRA training configuration — dataclass + env-var constants.
+
+Extracted from the monolithic train.py as part of W1 modularization. All
+env-var reads and the PostSemClawConfig dataclass live here. The training
+body imports these constants; zero behavior change from the extraction.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+
+
+def _parse_hyena_layers_env() -> tuple[int, ...]:
+    """Parse HYDRA_HYENA_LAYERS env var into a sorted tuple of layer indices.
+
+    Used as the default_factory for PostSemClawConfig.hyena_layers so a fresh
+    config construction reads the current env var, but once constructed the
+    value is first-class and travels with checkpoints (see asdict(config) in
+    save_ckpt). Ckpt-load sets the dataclass field explicitly, overriding the
+    env-var default.
+
+    Returns empty tuple when env var is unset/empty (byte-identical to
+    pre-port behavior: no Hyena layers).
+    """
+    raw = os.environ.get("HYDRA_HYENA_LAYERS", "")
+    if not raw:
+        return ()
+    return tuple(sorted({int(s.strip()) for s in raw.split(",") if s.strip()}))
+
+
+def _parse_gdn_layers_env() -> tuple[int, ...]:
+    """Parse HYDRA_GDN_LAYERS env var into a sorted tuple of layer indices.
+
+    Same contract as _parse_hyena_layers_env: layers whose index is listed
+    here use GatedDeltaNet (fla.layers.GatedDeltaNet) as a drop-in
+    replacement for Mamba3. Empty tuple = no GDN layers (byte-identical
+    to baseline).
+    """
+    raw = os.environ.get("HYDRA_GDN_LAYERS", "")
+    if not raw:
+        return ()
+    return tuple(sorted({int(s.strip()) for s in raw.split(",") if s.strip()}))
+
+# ---------------------------------------------------------------------------
+# CUDA env — set before importing torch in entry point. Kept here so any
+# module that `from hydra.config import ...` also benefits (import order is
+# top-down in Python, and train.py used to set these at module top).
+# ---------------------------------------------------------------------------
+os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
+if "/usr/local/cuda/bin" not in os.environ.get("PATH", ""):
+    os.environ["PATH"] = "/usr/local/cuda/bin:" + os.environ.get("PATH", "")
+os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
+
+
+# ---------------------------------------------------------------------------
+# Model Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class PostSemClawConfig:
+    """Full-architecture model config. Defaults reflect Phase-1 baseline;
+    the training entry overrides d_model/n_layer/etc. from env vars."""
+    # Sequence
+    sequence_len: int = 2048
+    vocab_size: int = 8192  # Must match prepare.py VOCAB_SIZE
+
+    # Mamba-3 SSM
+    n_layer: int = 6
+    d_model: int = 384
+    d_state: int = 64       # SSM state dimension
+    headdim: int = 48       # head dimension for SSM
+    n_heads: int = 8        # d_model // headdim
+    expand: int = 2         # inner_dim = expand * d_model
+
+    # Engram (conditional memory with Hebbian writes)
+    engram_n_columns: int = 4096
+    engram_key_dim: int = 64
+    engram_layer_idx: int = 1  # which layer gets engram (0-indexed, mid-layer)
+
+    # SemanticFoldingSDR (offline retina with STE; no-bypass, runs every step)
+    sdr_n_bits: int = 16384          # retina width
+    # Default 327 = 2% sparsity (Webber/Numenta canonical). Override with
+    # HYDRA_SDR_TARGET_ACTIVE env var; value MUST match subsystems/sdr_retina.py
+    # TARGET_ACTIVE (same env var is read there, so just setting it once works).
+    sdr_target_active: int = int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327"))
+    sdr_delta_rank: int = 32         # low-rank STE delta rank
+    sdr_som_warmup: int = 500
+    sdr_som_interval: int = 100
+
+    # HTMLayer (Rust-backed, Hebbian; no-bypass, runs every step)
+    htm_n_columns: int = 2048
+    htm_cells_per_column: int = 32
+
+    # Hyena supplement layer indices (sorted tuple). Defaults to the
+    # HYDRA_HYENA_LAYERS env var at config-construction time, but once
+    # persisted in a checkpoint the value is first-class and survives even
+    # when the env var is unset at resume time. This fixes the ckpt-reload
+    # crash path where a model trained with `HYDRA_HYENA_LAYERS=3,7` saves
+    # HyenaBlock params but a fresh process without the env var would try
+    # to build a pure-Mamba3 architecture and reject the state_dict as
+    # `Missing/Unexpected key(s)`.
+    hyena_layers: tuple[int, ...] = field(default_factory=_parse_hyena_layers_env)
+
+    # GatedDeltaNet supplement layer indices (sorted tuple). Same semantics
+    # as hyena_layers — a layer index listed here uses GDNBlock (fla-backed
+    # Gated DeltaNet) instead of Mamba3. Selections are mutually exclusive
+    # with hyena_layers at construction time (hyena wins on overlap; the
+    # model loop checks hyena first).
+    gdn_layers: tuple[int, ...] = field(default_factory=_parse_gdn_layers_env)
+
+    # Label smoothing + Z-loss
+    label_smoothing: float = 0.0   # disabled: any smoothing hurts in 5-min budget
+    z_loss_weight: float = 1e-4
+
+
+# ---------------------------------------------------------------------------
+# Hyperparameters (autoresearch agent modifies these via env vars)
+# ---------------------------------------------------------------------------
+
+# Model architecture
+D_MODEL = int(os.environ.get("HYDRA_D_MODEL", "256"))
+N_LAYER = int(os.environ.get("HYDRA_N_LAYER", "4"))
+D_STATE = int(os.environ.get("HYDRA_D_STATE", "64"))
+HEADDIM = int(os.environ.get("HYDRA_HEADDIM", "32"))
+N_HEADS = D_MODEL // HEADDIM
+EXPAND = int(os.environ.get("HYDRA_EXPAND", "2"))
+
+# Engram
+ENGRAM_N_COLUMNS = int(os.environ.get("HYDRA_ENGRAM_N_COLUMNS", "1024"))
+ENGRAM_KEY_DIM = 64
+ENGRAM_LAYER_IDX = int(os.environ.get("HYDRA_ENGRAM_LAYER_IDX", "1"))
+
+# Optimization
+DEVICE_BATCH_SIZE = int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
+TOTAL_BATCH_SIZE = int(os.environ.get("HYDRA_TOTAL_BATCH", "32768"))
+MATRIX_LR = float(os.environ.get("HYDRA_MATRIX_LR", "0.12"))
+EMBEDDING_LR = float(os.environ.get("HYDRA_EMBED_LR", "1.0"))
+UNEMBEDDING_LR = float(os.environ.get("HYDRA_UNEMBED_LR", "0.005"))
+# Scalar/vector params include Hyena implicit-filter vectors, norms, gate/bias
+# terms, and SDR delta_u/delta_v.  They are AdamW-scaled by d_model and can be
+# the hidden instability path when the high-throughput HF recipe pushes a large
+# device batch for hours.  Keep the historical default, but make it controllable
+# from launch scripts so cloud jobs can cool scalars without editing code.
+SCALAR_LR = float(os.environ.get("HYDRA_SCALAR_LR", "0.5"))
+WEIGHT_DECAY = float(os.environ.get("HYDRA_WEIGHT_DECAY", "0.01"))
+ADAM_BETAS = (0.9, 0.95)
+WARMUP_RATIO = float(os.environ.get("HYDRA_WARMUP_RATIO", "0.0"))
+WARMDOWN_RATIO = 0.5
+FINAL_LR_FRAC = float(os.environ.get("HYDRA_LR_MIN_MULT", "0.0"))
+
+# Runtime
+SEED = int(os.environ.get("HYDRA_SEED", "42"))
+# BF16 TFLOPS peak (RTX 3060=25.5, A100 SXM4=312, H100 SXM5=989)
+GPU_BF16_PEAK_FLOPS = float(os.environ.get("HYDRA_GPU_BF16_TFLOPS", "25.5")) * 1e12
+
+# Loss / inference knobs read by the model
+CE_CHUNK = int(os.environ.get("HYDRA_CE_CHUNK", "1024"))
+DROPOUT = float(os.environ.get("HYDRA_DROPOUT", "0.2"))
+FUSED_ADAMW = os.environ.get("HYDRA_FUSED_ADAMW", "1") == "1"
+
+# ---------------------------------------------------------------------------
+# Learnability knobs (all OFF by default — zero behavior change unless set)
+# ---------------------------------------------------------------------------
+# 1) Multi-Token Prediction (Llama-3 style). K=1 disables (next-1 only). K=4
+#    adds 3 extra weight-tied heads; loss = mean of K position-shifted CEs.
+MTP_K = int(os.environ.get("HYDRA_MTP_K", "1"))
+# 2) Exponential Moving Average of model weights (decay=0.999). Saves an
+#    additional latest_ema.pt at the end of training.
+USE_EMA = os.environ.get("HYDRA_USE_EMA", "0") == "1"
+EMA_DECAY = float(os.environ.get("HYDRA_EMA_DECAY", "0.999"))
+# 3) Gradient checkpointing on Mamba3 block forward. Trades ~30% compute for
+#    ~40% activation memory savings — lets you push B upward on a 3060.
+GRAD_CKPT = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
+# 4) Doc-separator masking in packed sequences: at every packed-BOS position
+#    in the targets tensor, mask the loss (ignore_index=-1) so the model is
+#    not forced to predict doc B from doc A's context.
+DOC_SEP_MASK = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
+# 5) Stop-gradient on HTM state (belt-and-braces: htm_rust already runs under
+#    torch.no_grad() so the tensor returned has requires_grad=False; this
+#    simply detaches explicitly to harden graph hygiene against future refactors).
+HTM_STOP_GRAD = os.environ.get("HYDRA_HTM_STOP_GRAD", "0") == "1"
+# 6) Output entropy penalty: loss += -lambda * H(softmax(logits)). Negative
+#    entropy penalizes peaked distributions and breaks repetition loops.
+ENTROPY_PENALTY = float(os.environ.get("HYDRA_ENTROPY_PENALTY", "0.0"))
+# 7) Curriculum: first N optimizer steps use short seq_len, then switch to
+#    full. 0 disables (no curriculum).
+CURRICULUM_SHORT_STEPS = int(os.environ.get("HYDRA_CURRICULUM_SHORT_STEPS", "0"))
+CURRICULUM_SHORT_SEQ_LEN = int(os.environ.get("HYDRA_CURRICULUM_SHORT_SEQ_LEN", "256"))
+
+# ---------------------------------------------------------------------------
+# Hyena supplement (additional block type for selected layer indices).
+# Hyena replaces Mamba3 at the specified layer indices while all other layers
+# remain Mamba3. Empty string (default) → no Hyena layers, byte-identical to
+# pre-port behavior.
+#   HYDRA_HYENA_LAYERS       "3,7"  — comma-separated 0-indexed layer ids
+#   HYDRA_HYENA_ORDER         2     — Hyena recurrence order (>= 2)
+#   HYDRA_HYENA_FILTER_DIM    64    — implicit-filter MLP hidden width
+# Hyena reference: https://arxiv.org/pdf/2302.10866.pdf (HazyResearch/safari).
+# ---------------------------------------------------------------------------
+HYENA_LAYERS = os.environ.get("HYDRA_HYENA_LAYERS", "")
+HYENA_ORDER = int(os.environ.get("HYDRA_HYENA_ORDER", "2"))
+HYENA_FILTER_DIM = int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64"))
+# Filter-rfft cache modes (see subsystems/hyena_pure.py):
+#   HYDRA_HYENA_FILTER_CACHE=1 — eval-only cache. Safe under torch.no_grad()
+#       where PyTorch never saves intermediate tensors. Off by default.
+#   HYDRA_HYENA_TRAIN_CACHE=1  — training-safe cache using a deferred
+#       gradient pattern. Cuts the implicit filter MLP forward to ONCE per
+#       optimizer step regardless of grad-accumulation factor. Requires the
+#       training loop (see hydra/lightning_module.py::optimizer_step) to
+#       call `model.flush_hyena_pending_grads()` before optimizer.step().
+#       Off by default.
+HYENA_FILTER_CACHE = os.environ.get("HYDRA_HYENA_FILTER_CACHE", "0") == "1"
+HYENA_TRAIN_CACHE = os.environ.get("HYDRA_HYENA_TRAIN_CACHE", "0") == "1"
+
+# Factual eval knobs
+FACTUAL_SAMPLES = int(os.environ.get("HYDRA_FACTUAL_SAMPLES", "3"))
+FACTUAL_BATCH = int(os.environ.get("HYDRA_FACTUAL_BATCH", "32"))
+# F6 (partial): Full incremental SSM decode integration deferred — would require
+# threading mamba_ssm InferenceParams through PostSemClawModel.forward and all
+# auxiliary subsystems (HTM, SDR, Engram) which currently run full-sequence each
+# call. As a stopgap we reduce default from 16 -> 4 so the per-prompt cost is
+# quartered (each gen-tok does a full re-encode of ctx+k tokens). Override with
+# HYDRA_FACTUAL_GEN_TOKENS to restore prior behavior. See docs/OPTIMIZATION_PLAN.md.
+FACTUAL_GEN_TOKENS = int(os.environ.get("HYDRA_FACTUAL_GEN_TOKENS", "2"))

overlay/hydra/data_module.py

@@ -1,288 +1,288 @@
-"""Lightning DataModule + IterableDataset for HYDRA pretraining.
-
-Replaces the custom threading/queue pipeline in prepare_nemotron.make_dataloader
-with a standard multiprocessing DataLoader approach.
-
-Design:
-  • IterableStreamDataset: each worker opens its own HF streams for the 7-way
-    blend, tokenizes with rustbpe, packs into (T+1,) rows via best-fit, and
-    yields one row per __next__.
-  • HydraDataModule: wraps the dataset with a standard DataLoader using
-    num_workers>=1, prefetch_factor=4, pin_memory=True. Lightning handles
-    device transfer.
-  • Val stream: deterministic seed 12345, weights match training blend.
-
-The worker RNG is seeded per-worker so the weighted-sampling schedule is
-independent across workers (else all workers request the same config at
-the same step and prefetching serializes).
-
-Env vars (all preserved from prepare_nemotron):
-  HYDRA_SEQ_LEN                  — sequence length T (default 512)
-  HYDRA_BATCH_SIZE               — batch size B (default 1) — passed through
-                                    to DataLoader
-  HYDRA_STREAM_SHUFFLE_BUFFER    — HF shuffle buffer (default 2048)
-  HYDRA_USE_FULL_BLEND           — 7-way blend vs 5-way Nemotron phase
-  HYDRA_USE_NEMOTRON             — enables streaming path (else shard path)
-  HYDRA_FACTUAL_INJECT_RATE      — factual doc injection cadence
-  HYDRA_NEMOTRON_PHASE           — phase1|phase2 (when not full blend)
-  HYDRA_DATA_NUM_WORKERS         — DataLoader num_workers (default 2)
-  HYDRA_DATA_PREFETCH            — DataLoader prefetch_factor (default 4)
-  HYDRA_DATA_BUFFER              — doc_buffer size for best-fit packing
-                                    (default 1000)
-"""
-from __future__ import annotations
-
-import os
-import random
-from typing import Iterator
-
-import numpy as np
-import torch
-import lightning as L
-from torch.utils.data import DataLoader, IterableDataset, get_worker_info
-
-import prepare as _prepare
-import prepare_nemotron as _p_nemo
-from prepare_nemotron import (
-    FULL_BLEND_WEIGHTS,
-    PHASE1_WEIGHTS,
-    PHASE2_WEIGHTS,
-    _BLEND_REGISTRY,
-    _extract_text,
-    _open_stream,
-)
-
-
-# ---------------------------------------------------------------------------
-# Worker-local weighted stream. A stripped version of prepare_nemotron's
-# _WeightedStream that is constructed inside each worker. Adds worker sharding:
-# when num_workers > 1 the RNG is seeded per-worker, so different workers
-# sample different config sequences and pull disjoint shard assignments from
-# HF's shuffle buffer.
-# ---------------------------------------------------------------------------
-
-
-class _WorkerWeightedStream:
-    def __init__(self, weights: dict[str, float], base_seed: int, worker_id: int):
-        self.configs = list(weights.keys())
-        self.weights = [weights[c] for c in self.configs]
-        self.base_seed = base_seed
-        self.worker_id = worker_id
-        # Each worker opens its own HF streams. _open_stream returns an iter()
-        # over a streaming dataset, with an internal shuffle buffer.
-        self.streams = {c: _open_stream(c, "train") for c in self.configs}
-        # Per-worker RNG so the config-choice trajectory is independent.
-        self.rng = random.Random(base_seed + worker_id * 7919)
-        self.epoch = 1
-
-        # Lazy-init factual docs (once per worker). The main-process version
-        # in prepare_nemotron._WeightedStream reads these on first __next__.
-        self._factual_docs: list[str] | None = None
-        self._factual_idx = 0
-        self._inject_counter = 0
-        inject_rate = int(os.environ.get("HYDRA_FACTUAL_INJECT_RATE", "50"))
-        self._inject_rate = inject_rate
-        if inject_rate > 0:
-            factual_path = os.path.join(
-                os.path.dirname(os.path.abspath(_p_nemo.__file__)),
-                "data", "factual", "facts.txt",
-            )
-            if os.path.exists(factual_path):
-                with open(factual_path) as fh:
-                    self._factual_docs = fh.read().strip().split("\n")
-
-    def _reopen(self, config: str) -> None:
-        self.streams[config] = _open_stream(config, "train")
-        self.epoch += 1
-
-    def __iter__(self):
-        return self
-
-    def __next__(self) -> tuple[str, int]:
-        # Factual injection (preserves prepare_nemotron cadence).
-        if self._inject_rate > 0 and self._factual_docs:
-            self._inject_counter += 1
-            if self._inject_counter >= self._inject_rate:
-                self._inject_counter = 0
-                doc = self._factual_docs[self._factual_idx % len(self._factual_docs)]
-                self._factual_idx += 1
-                return doc, self.epoch
-
-        config = self.rng.choices(self.configs, weights=self.weights, k=1)[0]
-        try:
-            row = next(self.streams[config])
-        except StopIteration:
-            self._reopen(config)
-            row = next(self.streams[config])
-        return _extract_text(row), self.epoch
-
-
-# ---------------------------------------------------------------------------
-# IterableStreamDataset — yields (T+1,) packed rows. No threads. No queues.
-# Lives inside each DataLoader worker. DataLoader's own multiprocessing stacks
-# rows into batches of shape (B, T+1) and sends them to the main process.
-# ---------------------------------------------------------------------------
-
-
-class IterableStreamDataset(IterableDataset):
-    """Streams docs, tokenizes, packs into (T+1,) rows via best-fit.
-
-    Each worker gets its own instance (via fork/spawn) and initializes its
-    own HF streams + rustbpe tokenizer + factual injector. The tokenizer
-    pickled blob is small (~1 MB) and thread-safe per tiktoken docs.
-    """
-
-    def __init__(
-        self,
-        split: str,
-        seq_len: int,
-        *,
-        base_seed: int = 0,
-        doc_buffer_size: int = 1000,
-        tokenizer_batch: int = 128,
-    ):
-        super().__init__()
-        assert split in ("train", "val"), split
-        self.split = split
-        self.seq_len = seq_len
-        self.row_capacity = seq_len + 1
-        self.base_seed = base_seed
-        self.doc_buffer_size = doc_buffer_size
-        self.tokenizer_batch = tokenizer_batch
-
-    def _pick_weights(self) -> dict[str, float]:
-        if self.split == "val":
-            if os.environ.get("HYDRA_USE_FULL_BLEND", "0") == "1":
-                return FULL_BLEND_WEIGHTS
-            return {"Nemotron-Pretraining-Multiple-Choice": 1.0}
-        if os.environ.get("HYDRA_USE_FULL_BLEND", "0") == "1":
-            return FULL_BLEND_WEIGHTS
-        phase = os.environ.get("HYDRA_NEMOTRON_PHASE", "phase1").strip().lower()
-        return PHASE2_WEIGHTS if phase == "phase2" else PHASE1_WEIGHTS
-
-    def __iter__(self) -> Iterator[torch.Tensor]:
-        info = get_worker_info()
-        worker_id = 0 if info is None else info.id
-
-        # Each worker builds its own tokenizer instance. tiktoken's Encoding
-        # object is pickleable and the underlying C++ BPE is thread-safe;
-        # per-worker instantiation avoids cross-process sharing headaches.
-        tokenizer = _prepare.Tokenizer.from_directory()
-        bos = tokenizer.get_bos_token_id()
-
-        # Each worker gets its own weighted HF stream. Seed offset ensures
-        # disjoint config-choice trajectories; HF's own shuffle buffer handles
-        # shard randomization.
-        val_seed = 12345  # deterministic val
-        seed = val_seed if self.split == "val" else self.base_seed
-        stream = _WorkerWeightedStream(
-            self._pick_weights(), base_seed=seed, worker_id=worker_id,
-        )
-
-        row_capacity = self.row_capacity
-        doc_buffer: list[list[int]] = []
-        doc_batch_size = self.tokenizer_batch
-
-        def refill_buffer() -> None:
-            # Collect doc_batch_size text strings, then batch-tokenize.
-            texts: list[str] = []
-            for _ in range(doc_batch_size):
-                text, _epoch = next(stream)
-                if text:
-                    texts.append(text)
-            if texts:
-                token_lists = tokenizer.encode(texts, prepend=bos)
-                doc_buffer.extend(token_lists)
-
-        while True:
-            pos = 0
-            row = torch.empty(row_capacity, dtype=torch.long)
-            while pos < row_capacity:
-                while len(doc_buffer) < self.doc_buffer_size:
-                    refill_buffer()
-
-                remaining = row_capacity - pos
-
-                # Best-fit packing: largest doc that fully fits.
-                best_idx = -1
-                best_len = 0
-                for i, doc in enumerate(doc_buffer):
-                    dlen = len(doc)
-                    if dlen <= remaining and dlen > best_len:
-                        best_idx = i
-                        best_len = dlen
-
-                if best_idx >= 0:
-                    doc = doc_buffer.pop(best_idx)
-                    row[pos : pos + len(doc)] = torch.tensor(doc, dtype=torch.long)
-                    pos += len(doc)
-                else:
-                    # No doc fits remaining space — crop shortest to fill.
-                    shortest_idx = min(
-                        range(len(doc_buffer)),
-                        key=lambda i: len(doc_buffer[i]),
-                    )
-                    doc = doc_buffer.pop(shortest_idx)
-                    row[pos : pos + remaining] = torch.tensor(
-                        doc[:remaining], dtype=torch.long,
-                    )
-                    pos += remaining
-
-            yield row
-
-
-# ---------------------------------------------------------------------------
-# LightningDataModule
-# ---------------------------------------------------------------------------
-
-
-class HydraDataModule(L.LightningDataModule):
-    def __init__(
-        self,
-        batch_size: int | None = None,
-        seq_len: int | None = None,
-        num_workers: int | None = None,
-        prefetch_factor: int | None = None,
-    ):
-        super().__init__()
-        self.batch_size = batch_size or int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
-        self.seq_len = seq_len or int(os.environ.get("HYDRA_SEQ_LEN", "512"))
-        self.num_workers = (
-            num_workers
-            if num_workers is not None
-            else int(os.environ.get("HYDRA_DATA_NUM_WORKERS", "2"))
-        )
-        self.prefetch_factor = (
-            prefetch_factor
-            if prefetch_factor is not None
-            else int(os.environ.get("HYDRA_DATA_PREFETCH", "4"))
-        )
-        self.doc_buffer = int(os.environ.get("HYDRA_DATA_BUFFER", "1000"))
-
-    def _make_loader(self, split: str, seed: int) -> DataLoader:
-        dataset = IterableStreamDataset(
-            split=split,
-            seq_len=self.seq_len,
-            base_seed=seed,
-            doc_buffer_size=self.doc_buffer,
-        )
-        # num_workers=0 → main-process iteration (useful for debugging). With
-        # IterableDataset the DataLoader batches the rows into (B, T+1) via
-        # default torch.stack-collate.
-        kw: dict = dict(
-            dataset=dataset,
-            batch_size=self.batch_size,
-            num_workers=self.num_workers,
-            pin_memory=True,
-            drop_last=True,
-        )
-        if self.num_workers > 0:
-            kw["prefetch_factor"] = self.prefetch_factor
-            kw["persistent_workers"] = True
-        return DataLoader(**kw)
-
-    def train_dataloader(self) -> DataLoader:
-        return self._make_loader("train", seed=0)
-
-    def val_dataloader(self) -> DataLoader:
-        return self._make_loader("val", seed=12345)
+"""Lightning DataModule + IterableDataset for HYDRA pretraining.
+
+Replaces the custom threading/queue pipeline in prepare_nemotron.make_dataloader
+with a standard multiprocessing DataLoader approach.
+
+Design:
+  • IterableStreamDataset: each worker opens its own HF streams for the 7-way
+    blend, tokenizes with rustbpe, packs into (T+1,) rows via best-fit, and
+    yields one row per __next__.
+  • HydraDataModule: wraps the dataset with a standard DataLoader using
+    num_workers>=1, prefetch_factor=4, pin_memory=True. Lightning handles
+    device transfer.
+  • Val stream: deterministic seed 12345, weights match training blend.
+
+The worker RNG is seeded per-worker so the weighted-sampling schedule is
+independent across workers (else all workers request the same config at
+the same step and prefetching serializes).
+
+Env vars (all preserved from prepare_nemotron):
+  HYDRA_SEQ_LEN                  — sequence length T (default 512)
+  HYDRA_BATCH_SIZE               — batch size B (default 1) — passed through
+                                    to DataLoader
+  HYDRA_STREAM_SHUFFLE_BUFFER    — HF shuffle buffer (default 2048)
+  HYDRA_USE_FULL_BLEND           — 7-way blend vs 5-way Nemotron phase
+  HYDRA_USE_NEMOTRON             — enables streaming path (else shard path)
+  HYDRA_FACTUAL_INJECT_RATE      — factual doc injection cadence
+  HYDRA_NEMOTRON_PHASE           — phase1|phase2 (when not full blend)
+  HYDRA_DATA_NUM_WORKERS         — DataLoader num_workers (default 2)
+  HYDRA_DATA_PREFETCH            — DataLoader prefetch_factor (default 4)
+  HYDRA_DATA_BUFFER              — doc_buffer size for best-fit packing
+                                    (default 1000)
+"""
+from __future__ import annotations
+
+import os
+import random
+from typing import Iterator
+
+import numpy as np
+import torch
+import lightning as L
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
+
+import prepare as _prepare
+import prepare_nemotron as _p_nemo
+from prepare_nemotron import (
+    FULL_BLEND_WEIGHTS,
+    PHASE1_WEIGHTS,
+    PHASE2_WEIGHTS,
+    _BLEND_REGISTRY,
+    _extract_text,
+    _open_stream,
+)
+
+
+# ---------------------------------------------------------------------------
+# Worker-local weighted stream. A stripped version of prepare_nemotron's
+# _WeightedStream that is constructed inside each worker. Adds worker sharding:
+# when num_workers > 1 the RNG is seeded per-worker, so different workers
+# sample different config sequences and pull disjoint shard assignments from
+# HF's shuffle buffer.
+# ---------------------------------------------------------------------------
+
+
+class _WorkerWeightedStream:
+    def __init__(self, weights: dict[str, float], base_seed: int, worker_id: int):
+        self.configs = list(weights.keys())
+        self.weights = [weights[c] for c in self.configs]
+        self.base_seed = base_seed
+        self.worker_id = worker_id
+        # Each worker opens its own HF streams. _open_stream returns an iter()
+        # over a streaming dataset, with an internal shuffle buffer.
+        self.streams = {c: _open_stream(c, "train") for c in self.configs}
+        # Per-worker RNG so the config-choice trajectory is independent.
+        self.rng = random.Random(base_seed + worker_id * 7919)
+        self.epoch = 1
+
+        # Lazy-init factual docs (once per worker). The main-process version
+        # in prepare_nemotron._WeightedStream reads these on first __next__.
+        self._factual_docs: list[str] | None = None
+        self._factual_idx = 0
+        self._inject_counter = 0
+        inject_rate = int(os.environ.get("HYDRA_FACTUAL_INJECT_RATE", "50"))
+        self._inject_rate = inject_rate
+        if inject_rate > 0:
+            factual_path = os.path.join(
+                os.path.dirname(os.path.abspath(_p_nemo.__file__)),
+                "data", "factual", "facts.txt",
+            )
+            if os.path.exists(factual_path):
+                with open(factual_path) as fh:
+                    self._factual_docs = fh.read().strip().split("\n")
+
+    def _reopen(self, config: str) -> None:
+        self.streams[config] = _open_stream(config, "train")
+        self.epoch += 1
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> tuple[str, int]:
+        # Factual injection (preserves prepare_nemotron cadence).
+        if self._inject_rate > 0 and self._factual_docs:
+            self._inject_counter += 1
+            if self._inject_counter >= self._inject_rate:
+                self._inject_counter = 0
+                doc = self._factual_docs[self._factual_idx % len(self._factual_docs)]
+                self._factual_idx += 1
+                return doc, self.epoch
+
+        config = self.rng.choices(self.configs, weights=self.weights, k=1)[0]
+        try:
+            row = next(self.streams[config])
+        except StopIteration:
+            self._reopen(config)
+            row = next(self.streams[config])
+        return _extract_text(row), self.epoch
+
+
+# ---------------------------------------------------------------------------
+# IterableStreamDataset — yields (T+1,) packed rows. No threads. No queues.
+# Lives inside each DataLoader worker. DataLoader's own multiprocessing stacks
+# rows into batches of shape (B, T+1) and sends them to the main process.
+# ---------------------------------------------------------------------------
+
+
+class IterableStreamDataset(IterableDataset):
+    """Streams docs, tokenizes, packs into (T+1,) rows via best-fit.
+
+    Each worker gets its own instance (via fork/spawn) and initializes its
+    own HF streams + rustbpe tokenizer + factual injector. The tokenizer
+    pickled blob is small (~1 MB) and thread-safe per tiktoken docs.
+    """
+
+    def __init__(
+        self,
+        split: str,
+        seq_len: int,
+        *,
+        base_seed: int = 0,
+        doc_buffer_size: int = 1000,
+        tokenizer_batch: int = 128,
+    ):
+        super().__init__()
+        assert split in ("train", "val"), split
+        self.split = split
+        self.seq_len = seq_len
+        self.row_capacity = seq_len + 1
+        self.base_seed = base_seed
+        self.doc_buffer_size = doc_buffer_size
+        self.tokenizer_batch = tokenizer_batch
+
+    def _pick_weights(self) -> dict[str, float]:
+        if self.split == "val":
+            if os.environ.get("HYDRA_USE_FULL_BLEND", "0") == "1":
+                return FULL_BLEND_WEIGHTS
+            return {"Nemotron-Pretraining-Multiple-Choice": 1.0}
+        if os.environ.get("HYDRA_USE_FULL_BLEND", "0") == "1":
+            return FULL_BLEND_WEIGHTS
+        phase = os.environ.get("HYDRA_NEMOTRON_PHASE", "phase1").strip().lower()
+        return PHASE2_WEIGHTS if phase == "phase2" else PHASE1_WEIGHTS
+
+    def __iter__(self) -> Iterator[torch.Tensor]:
+        info = get_worker_info()
+        worker_id = 0 if info is None else info.id
+
+        # Each worker builds its own tokenizer instance. tiktoken's Encoding
+        # object is pickleable and the underlying C++ BPE is thread-safe;
+        # per-worker instantiation avoids cross-process sharing headaches.
+        tokenizer = _prepare.Tokenizer.from_directory()
+        bos = tokenizer.get_bos_token_id()
+
+        # Each worker gets its own weighted HF stream. Seed offset ensures
+        # disjoint config-choice trajectories; HF's own shuffle buffer handles
+        # shard randomization.
+        val_seed = 12345  # deterministic val
+        seed = val_seed if self.split == "val" else self.base_seed
+        stream = _WorkerWeightedStream(
+            self._pick_weights(), base_seed=seed, worker_id=worker_id,
+        )
+
+        row_capacity = self.row_capacity
+        doc_buffer: list[list[int]] = []
+        doc_batch_size = self.tokenizer_batch
+
+        def refill_buffer() -> None:
+            # Collect doc_batch_size text strings, then batch-tokenize.
+            texts: list[str] = []
+            for _ in range(doc_batch_size):
+                text, _epoch = next(stream)
+                if text:
+                    texts.append(text)
+            if texts:
+                token_lists = tokenizer.encode(texts, prepend=bos)
+                doc_buffer.extend(token_lists)
+
+        while True:
+            pos = 0
+            row = torch.empty(row_capacity, dtype=torch.long)
+            while pos < row_capacity:
+                while len(doc_buffer) < self.doc_buffer_size:
+                    refill_buffer()
+
+                remaining = row_capacity - pos
+
+                # Best-fit packing: largest doc that fully fits.
+                best_idx = -1
+                best_len = 0
+                for i, doc in enumerate(doc_buffer):
+                    dlen = len(doc)
+                    if dlen <= remaining and dlen > best_len:
+                        best_idx = i
+                        best_len = dlen
+
+                if best_idx >= 0:
+                    doc = doc_buffer.pop(best_idx)
+                    row[pos : pos + len(doc)] = torch.tensor(doc, dtype=torch.long)
+                    pos += len(doc)
+                else:
+                    # No doc fits remaining space — crop shortest to fill.
+                    shortest_idx = min(
+                        range(len(doc_buffer)),
+                        key=lambda i: len(doc_buffer[i]),
+                    )
+                    doc = doc_buffer.pop(shortest_idx)
+                    row[pos : pos + remaining] = torch.tensor(
+                        doc[:remaining], dtype=torch.long,
+                    )
+                    pos += remaining
+
+            yield row
+
+
+# ---------------------------------------------------------------------------
+# LightningDataModule
+# ---------------------------------------------------------------------------
+
+
+class HydraDataModule(L.LightningDataModule):
+    def __init__(
+        self,
+        batch_size: int | None = None,
+        seq_len: int | None = None,
+        num_workers: int | None = None,
+        prefetch_factor: int | None = None,
+    ):
+        super().__init__()
+        self.batch_size = batch_size or int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
+        self.seq_len = seq_len or int(os.environ.get("HYDRA_SEQ_LEN", "512"))
+        self.num_workers = (
+            num_workers
+            if num_workers is not None
+            else int(os.environ.get("HYDRA_DATA_NUM_WORKERS", "2"))
+        )
+        self.prefetch_factor = (
+            prefetch_factor
+            if prefetch_factor is not None
+            else int(os.environ.get("HYDRA_DATA_PREFETCH", "4"))
+        )
+        self.doc_buffer = int(os.environ.get("HYDRA_DATA_BUFFER", "1000"))
+
+    def _make_loader(self, split: str, seed: int) -> DataLoader:
+        dataset = IterableStreamDataset(
+            split=split,
+            seq_len=self.seq_len,
+            base_seed=seed,
+            doc_buffer_size=self.doc_buffer,
+        )
+        # num_workers=0 → main-process iteration (useful for debugging). With
+        # IterableDataset the DataLoader batches the rows into (B, T+1) via
+        # default torch.stack-collate.
+        kw: dict = dict(
+            dataset=dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=True,
+        )
+        if self.num_workers > 0:
+            kw["prefetch_factor"] = self.prefetch_factor
+            kw["persistent_workers"] = True
+        return DataLoader(**kw)
+
+    def train_dataloader(self) -> DataLoader:
+        return self._make_loader("train", seed=0)
+
+    def val_dataloader(self) -> DataLoader:
+        return self._make_loader("val", seed=12345)

overlay/hydra/diffusion_loss.py

@@ -1,236 +1,236 @@
-"""MDLM Rao-Blackwellized Masked Diffusion Loss.
-
-Implements the masked-diffusion ELBO from:
-    Sahoo et al., "Simple and Effective Masked Diffusion Language Models" (MDLM),
-    NeurIPS 2024, arXiv:2406.07524.
-
-Equations referenced:
-    - Forward process: eq. 2  (per-token Bernoulli masking at rate 1 - alpha_t)
-    - Log-linear schedule:    alpha_t = 1 - t,  t ~ Uniform(0, 1)
-    - RB-ELBO:     eq. 7-8   L_RB = E_t E_q [ (1/alpha_t) * CE(x_theta(x_t), x_0) ]
-                              where the expectation over masked positions.
-
-Key insight: the Rao-Blackwellized estimate replaces an average over all masks
-(exponential) by a closed-form weighted CE that applies weight 1/alpha_t only
-on the positions that were masked, and 0 on unmasked positions. This gives an
-unbiased estimator with lower variance than a naive Monte Carlo over mask
-patterns.
-
-Reference implementation cross-checked against:
-    https://github.com/kuleshov-group/mdlm  (diffusion.py::DiffusionModel._loss)
-"""
-
-from __future__ import annotations
-
-from typing import Literal
-
-import torch
-import torch.nn.functional as F
-
-
-# Clamping weight keeps gradients finite while still up-weighting high-noise
-# positions. Historical value 1/eps=1000 blew up HYDRA training on a 12h v2
-# launch (2026-04-22): loss 26 → 42 → NaN in 13 steps under Muon lr=7e-3
-# because per-token CE × 1000 saturated the 100-unit FAIL guard. The MDLM
-# paper reports stable training at Adam lr=1e-4; HYDRA uses Muon at 7e-3
-# (70× larger), so the weight clamp needs to compensate.
-#
-# Tunable via HYDRA_MDLM_MAX_WEIGHT (default 5.0). Set =1.0 to disable
-# weighting entirely (flat masked-LM CE, no RB reweighting — simpler and
-# more stable, sacrifices the theoretical ELBO property).
-import os as _os
-_MAX_WEIGHT: float = float(_os.environ.get("HYDRA_MDLM_MAX_WEIGHT", "5.0"))
-_MIN_ALPHA: float = 1.0 / _MAX_WEIGHT  # so clamp(alpha, min=_MIN_ALPHA) gives 1/alpha <= _MAX_WEIGHT
-
-
-# ---------------------------------------------------------------------------
-# Public API
-# ---------------------------------------------------------------------------
-
-def mdlm_masked_forward_process(
-    targets: torch.Tensor,
-    mask_token_id: int,
-    t: torch.Tensor | None = None,
-    alpha_schedule: Literal["linear", "loglinear"] = "loglinear",
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """MDLM forward (noising) process: mask tokens and compute RB weights.
-
-    Args:
-        targets: (B, T) int64 token ids — the clean sequence x_0.
-        mask_token_id: The special token id used to represent a masked token.
-        t: (B,) float in (0, 1). If None, samples Uniform(0, 1) per batch
-            element. t=0 means fully clean; t=1 means fully masked.
-        alpha_schedule: Noise schedule.
-            "loglinear" (MDLM default): alpha_t = 1 - t
-            "linear": identical formula — both are provided for completeness
-            since the paper calls the 1-t schedule "log-linear" in the context
-            of the ELBO derivation.
-
-    Returns:
-        x_t           : (B, T) int64 — noised sequence; masked positions hold
-                        mask_token_id, unmasked positions equal targets.
-        mask_positions: (B, T) bool  — True where the token was masked.
-        loss_weights  : (B, T) float32 — RB weighting factor. On masked
-                        positions: 1/alpha_t (clamped to _MAX_WEIGHT). On
-                        unmasked positions: 0.0. Summing
-                        (CE * loss_weights * mask_positions).sum() / mask.sum()
-                        gives the per-sample RB-ELBO estimator.
-    """
-    B, T = targets.shape
-    device = targets.device
-    dtype = torch.float32
-
-    # --- sample or validate t ---
-    if t is None:
-        # Uniform(0, 1) per batch element; avoid exactly 0 and 1.
-        t = torch.rand(B, device=device, dtype=dtype)
-    else:
-        t = t.to(device=device, dtype=dtype)
-        if t.shape != (B,):
-            raise ValueError(f"t must be shape (B,)={(B,)}, got {t.shape}")
-        if (t < 0).any() or (t > 1).any():
-            raise ValueError("t must be in [0, 1]")
-
-    # --- noise schedule: alpha_t = probability that a token is NOT masked ---
-    # Both "linear" and "loglinear" in MDLM use alpha_t = 1 - t; the paper
-    # refers to "log-linear" because the schedule is linear in the *log* domain
-    # of the forward process probability. We expose both names for clarity.
-    if alpha_schedule in ("linear", "loglinear"):
-        alpha_t = 1.0 - t          # (B,) float, in [0, 1]
-    else:
-        raise ValueError(f"Unknown alpha_schedule: {alpha_schedule!r}. Use 'linear' or 'loglinear'.")
-
-    # --- per-token Bernoulli mask ---
-    # alpha_t[:, None] broadcasts to (B, T).
-    alpha_t_expanded = alpha_t[:, None]                # (B, 1)
-    # Bernoulli(1 - alpha_t) = 1 means "mask this token".
-    # We sample independently per token, per batch element.
-    rand = torch.rand(B, T, device=device, dtype=dtype)
-    mask_positions = rand > alpha_t_expanded           # (B, T) bool
-    # True  → masked position
-    # False → unmasked (kept as original)
-
-    # --- build x_t ---
-    x_t = targets.clone()
-    x_t = torch.where(mask_positions, torch.full_like(x_t, mask_token_id), x_t)
-
-    # --- RB loss weights: 1/alpha_t on masked positions, 0 elsewhere ---
-    # Clamp alpha_t so weights stay finite near t→1.
-    safe_alpha = alpha_t.clamp(min=_MIN_ALPHA)         # (B,)
-    weight_per_sample = 1.0 / safe_alpha               # (B,)
-    # Broadcast to (B, T) and zero out unmasked positions.
-    loss_weights = weight_per_sample[:, None].expand(B, T).to(dtype=dtype)  # (B, T)
-    loss_weights = loss_weights * mask_positions.float()
-
-    return x_t, mask_positions, loss_weights
-
-
-def mdlm_rb_loss(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-    mask_positions: torch.Tensor,
-    loss_weights: torch.Tensor,
-    ignore_index: int = -100,
-) -> torch.Tensor:
-    """Rao-Blackwellized negative ELBO.
-
-    Applies the MDLM loss: cross-entropy on masked positions only, weighted
-    per-token by loss_weights, averaged over the batch.
-
-    The formula (eq. 7-8 of arXiv:2406.07524):
-        L_RB = mean_B [ sum_T (weight_t * CE(logits_i, target_i) * mask_i)
-                        / max(sum_T(mask_i), 1) ]
-
-    Args:
-        logits        : (B, T, V) raw logits. May be bf16; internally cast to
-                        float32 for CE computation.
-        targets       : (B, T) int64 true token ids (x_0).
-        mask_positions: (B, T) bool — True = masked position.
-        loss_weights  : (B, T) float32 — 1/alpha_t on masked positions, 0 elsewhere.
-        ignore_index  : Passed to F.cross_entropy; positions with this label
-                        are excluded from the loss.
-
-    Returns:
-        Scalar float32 loss. Returns 0.0 tensor if no positions are masked.
-    """
-    B, T, V = logits.shape
-
-    # Ensure float32 for numerical stability; F.cross_entropy accepts fp16/bf16
-    # logits but accumulates in float internally anyway. Being explicit avoids
-    # silent precision surprises.
-    logits_f = logits.float()                          # (B, T, V)
-
-    # Build targets with ignore_index on UNmasked positions so CE only fires
-    # where mask_positions is True. We also honour any pre-existing -100 values
-    # (e.g. doc-separator masking upstream).
-    targets_masked = torch.where(
-        mask_positions & (targets != ignore_index),
-        targets,
-        torch.full_like(targets, ignore_index),
-    )
-
-    # Per-token CE; shape (B, T). Positions with ignore_index → 0 from CE.
-    per_tok_ce = F.cross_entropy(
-        logits_f.reshape(B * T, V),
-        targets_masked.reshape(B * T),
-        ignore_index=ignore_index,
-        reduction="none",
-    ).reshape(B, T)                                    # (B, T) float32
-
-    # Apply RB weight. loss_weights already has 0 on unmasked positions.
-    weighted = per_tok_ce * loss_weights               # (B, T)
-
-    # Per-sample mean over masked positions, then average over batch.
-    mask_f = mask_positions.float()                    # (B, T)
-    per_sample_mask_count = mask_f.sum(dim=1).clamp(min=1)   # (B,)
-    per_sample_loss = weighted.sum(dim=1) / per_sample_mask_count  # (B,)
-
-    return per_sample_loss.mean()                      # scalar float32
-
-
-def mdlm_loss(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-    mask_token_id: int,
-    t: torch.Tensor | None = None,
-    alpha_schedule: Literal["linear", "loglinear"] = "loglinear",
-    ignore_index: int = -100,
-) -> torch.Tensor:
-    """Convenience wrapper: forward process + RB-ELBO in one call.
-
-    Suitable for the common case where the caller has full-vocab logits and
-    wants a drop-in replacement for a standard masked-LM CE loss.
-
-    Args:
-        logits        : (B, T, V) raw logits.
-        targets       : (B, T) int64 clean token ids.
-        mask_token_id : The MASK token id used to corrupt the input.
-        t             : Optional (B,) timestep in (0, 1). Sampled if None.
-        alpha_schedule: "loglinear" (default) or "linear".
-        ignore_index  : Token id to ignore in the loss (e.g. padding).
-
-    Returns:
-        Scalar float32 MDLM RB-ELBO loss.
-
-    Note on sampled-softmax / partial logits:
-        If your model only computes logits for a subset of vocab positions
-        (e.g. HYDRA's sampled-softmax head), call mdlm_masked_forward_process
-        and mdlm_rb_loss separately. mdlm_rb_loss expects full-vocab logits.
-    """
-    x_t, mask_positions, loss_weights = mdlm_masked_forward_process(
-        targets=targets,
-        mask_token_id=mask_token_id,
-        t=t,
-        alpha_schedule=alpha_schedule,
-    )
-    # x_t is produced for the model's input (not used by this convenience
-    # wrapper since logits are already provided by the caller). In a real
-    # training loop the caller feeds x_t into the model to get logits, THEN
-    # calls this function. See the orchestrator wiring note in training.py.
-    return mdlm_rb_loss(
-        logits=logits,
-        targets=targets,
-        mask_positions=mask_positions,
-        loss_weights=loss_weights,
-        ignore_index=ignore_index,
-    )
+"""MDLM Rao-Blackwellized Masked Diffusion Loss.
+
+Implements the masked-diffusion ELBO from:
+    Sahoo et al., "Simple and Effective Masked Diffusion Language Models" (MDLM),
+    NeurIPS 2024, arXiv:2406.07524.
+
+Equations referenced:
+    - Forward process: eq. 2  (per-token Bernoulli masking at rate 1 - alpha_t)
+    - Log-linear schedule:    alpha_t = 1 - t,  t ~ Uniform(0, 1)
+    - RB-ELBO:     eq. 7-8   L_RB = E_t E_q [ (1/alpha_t) * CE(x_theta(x_t), x_0) ]
+                              where the expectation over masked positions.
+
+Key insight: the Rao-Blackwellized estimate replaces an average over all masks
+(exponential) by a closed-form weighted CE that applies weight 1/alpha_t only
+on the positions that were masked, and 0 on unmasked positions. This gives an
+unbiased estimator with lower variance than a naive Monte Carlo over mask
+patterns.
+
+Reference implementation cross-checked against:
+    https://github.com/kuleshov-group/mdlm  (diffusion.py::DiffusionModel._loss)
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+
+
+# Clamping weight keeps gradients finite while still up-weighting high-noise
+# positions. Historical value 1/eps=1000 blew up HYDRA training on a 12h v2
+# launch (2026-04-22): loss 26 → 42 → NaN in 13 steps under Muon lr=7e-3
+# because per-token CE × 1000 saturated the 100-unit FAIL guard. The MDLM
+# paper reports stable training at Adam lr=1e-4; HYDRA uses Muon at 7e-3
+# (70× larger), so the weight clamp needs to compensate.
+#
+# Tunable via HYDRA_MDLM_MAX_WEIGHT (default 5.0). Set =1.0 to disable
+# weighting entirely (flat masked-LM CE, no RB reweighting — simpler and
+# more stable, sacrifices the theoretical ELBO property).
+import os as _os
+_MAX_WEIGHT: float = float(_os.environ.get("HYDRA_MDLM_MAX_WEIGHT", "5.0"))
+_MIN_ALPHA: float = 1.0 / _MAX_WEIGHT  # so clamp(alpha, min=_MIN_ALPHA) gives 1/alpha <= _MAX_WEIGHT
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def mdlm_masked_forward_process(
+    targets: torch.Tensor,
+    mask_token_id: int,
+    t: torch.Tensor | None = None,
+    alpha_schedule: Literal["linear", "loglinear"] = "loglinear",
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """MDLM forward (noising) process: mask tokens and compute RB weights.
+
+    Args:
+        targets: (B, T) int64 token ids — the clean sequence x_0.
+        mask_token_id: The special token id used to represent a masked token.
+        t: (B,) float in (0, 1). If None, samples Uniform(0, 1) per batch
+            element. t=0 means fully clean; t=1 means fully masked.
+        alpha_schedule: Noise schedule.
+            "loglinear" (MDLM default): alpha_t = 1 - t
+            "linear": identical formula — both are provided for completeness
+            since the paper calls the 1-t schedule "log-linear" in the context
+            of the ELBO derivation.
+
+    Returns:
+        x_t           : (B, T) int64 — noised sequence; masked positions hold
+                        mask_token_id, unmasked positions equal targets.
+        mask_positions: (B, T) bool  — True where the token was masked.
+        loss_weights  : (B, T) float32 — RB weighting factor. On masked
+                        positions: 1/alpha_t (clamped to _MAX_WEIGHT). On
+                        unmasked positions: 0.0. Summing
+                        (CE * loss_weights * mask_positions).sum() / mask.sum()
+                        gives the per-sample RB-ELBO estimator.
+    """
+    B, T = targets.shape
+    device = targets.device
+    dtype = torch.float32
+
+    # --- sample or validate t ---
+    if t is None:
+        # Uniform(0, 1) per batch element; avoid exactly 0 and 1.
+        t = torch.rand(B, device=device, dtype=dtype)
+    else:
+        t = t.to(device=device, dtype=dtype)
+        if t.shape != (B,):
+            raise ValueError(f"t must be shape (B,)={(B,)}, got {t.shape}")
+        if (t < 0).any() or (t > 1).any():
+            raise ValueError("t must be in [0, 1]")
+
+    # --- noise schedule: alpha_t = probability that a token is NOT masked ---
+    # Both "linear" and "loglinear" in MDLM use alpha_t = 1 - t; the paper
+    # refers to "log-linear" because the schedule is linear in the *log* domain
+    # of the forward process probability. We expose both names for clarity.
+    if alpha_schedule in ("linear", "loglinear"):
+        alpha_t = 1.0 - t          # (B,) float, in [0, 1]
+    else:
+        raise ValueError(f"Unknown alpha_schedule: {alpha_schedule!r}. Use 'linear' or 'loglinear'.")
+
+    # --- per-token Bernoulli mask ---
+    # alpha_t[:, None] broadcasts to (B, T).
+    alpha_t_expanded = alpha_t[:, None]                # (B, 1)
+    # Bernoulli(1 - alpha_t) = 1 means "mask this token".
+    # We sample independently per token, per batch element.
+    rand = torch.rand(B, T, device=device, dtype=dtype)
+    mask_positions = rand > alpha_t_expanded           # (B, T) bool
+    # True  → masked position
+    # False → unmasked (kept as original)
+
+    # --- build x_t ---
+    x_t = targets.clone()
+    x_t = torch.where(mask_positions, torch.full_like(x_t, mask_token_id), x_t)
+
+    # --- RB loss weights: 1/alpha_t on masked positions, 0 elsewhere ---
+    # Clamp alpha_t so weights stay finite near t→1.
+    safe_alpha = alpha_t.clamp(min=_MIN_ALPHA)         # (B,)
+    weight_per_sample = 1.0 / safe_alpha               # (B,)
+    # Broadcast to (B, T) and zero out unmasked positions.
+    loss_weights = weight_per_sample[:, None].expand(B, T).to(dtype=dtype)  # (B, T)
+    loss_weights = loss_weights * mask_positions.float()
+
+    return x_t, mask_positions, loss_weights
+
+
+def mdlm_rb_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    mask_positions: torch.Tensor,
+    loss_weights: torch.Tensor,
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Rao-Blackwellized negative ELBO.
+
+    Applies the MDLM loss: cross-entropy on masked positions only, weighted
+    per-token by loss_weights, averaged over the batch.
+
+    The formula (eq. 7-8 of arXiv:2406.07524):
+        L_RB = mean_B [ sum_T (weight_t * CE(logits_i, target_i) * mask_i)
+                        / max(sum_T(mask_i), 1) ]
+
+    Args:
+        logits        : (B, T, V) raw logits. May be bf16; internally cast to
+                        float32 for CE computation.
+        targets       : (B, T) int64 true token ids (x_0).
+        mask_positions: (B, T) bool — True = masked position.
+        loss_weights  : (B, T) float32 — 1/alpha_t on masked positions, 0 elsewhere.
+        ignore_index  : Passed to F.cross_entropy; positions with this label
+                        are excluded from the loss.
+
+    Returns:
+        Scalar float32 loss. Returns 0.0 tensor if no positions are masked.
+    """
+    B, T, V = logits.shape
+
+    # Ensure float32 for numerical stability; F.cross_entropy accepts fp16/bf16
+    # logits but accumulates in float internally anyway. Being explicit avoids
+    # silent precision surprises.
+    logits_f = logits.float()                          # (B, T, V)
+
+    # Build targets with ignore_index on UNmasked positions so CE only fires
+    # where mask_positions is True. We also honour any pre-existing -100 values
+    # (e.g. doc-separator masking upstream).
+    targets_masked = torch.where(
+        mask_positions & (targets != ignore_index),
+        targets,
+        torch.full_like(targets, ignore_index),
+    )
+
+    # Per-token CE; shape (B, T). Positions with ignore_index → 0 from CE.
+    per_tok_ce = F.cross_entropy(
+        logits_f.reshape(B * T, V),
+        targets_masked.reshape(B * T),
+        ignore_index=ignore_index,
+        reduction="none",
+    ).reshape(B, T)                                    # (B, T) float32
+
+    # Apply RB weight. loss_weights already has 0 on unmasked positions.
+    weighted = per_tok_ce * loss_weights               # (B, T)
+
+    # Per-sample mean over masked positions, then average over batch.
+    mask_f = mask_positions.float()                    # (B, T)
+    per_sample_mask_count = mask_f.sum(dim=1).clamp(min=1)   # (B,)
+    per_sample_loss = weighted.sum(dim=1) / per_sample_mask_count  # (B,)
+
+    return per_sample_loss.mean()                      # scalar float32
+
+
+def mdlm_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    mask_token_id: int,
+    t: torch.Tensor | None = None,
+    alpha_schedule: Literal["linear", "loglinear"] = "loglinear",
+    ignore_index: int = -100,
+) -> torch.Tensor:
+    """Convenience wrapper: forward process + RB-ELBO in one call.
+
+    Suitable for the common case where the caller has full-vocab logits and
+    wants a drop-in replacement for a standard masked-LM CE loss.
+
+    Args:
+        logits        : (B, T, V) raw logits.
+        targets       : (B, T) int64 clean token ids.
+        mask_token_id : The MASK token id used to corrupt the input.
+        t             : Optional (B,) timestep in (0, 1). Sampled if None.
+        alpha_schedule: "loglinear" (default) or "linear".
+        ignore_index  : Token id to ignore in the loss (e.g. padding).
+
+    Returns:
+        Scalar float32 MDLM RB-ELBO loss.
+
+    Note on sampled-softmax / partial logits:
+        If your model only computes logits for a subset of vocab positions
+        (e.g. HYDRA's sampled-softmax head), call mdlm_masked_forward_process
+        and mdlm_rb_loss separately. mdlm_rb_loss expects full-vocab logits.
+    """
+    x_t, mask_positions, loss_weights = mdlm_masked_forward_process(
+        targets=targets,
+        mask_token_id=mask_token_id,
+        t=t,
+        alpha_schedule=alpha_schedule,
+    )
+    # x_t is produced for the model's input (not used by this convenience
+    # wrapper since logits are already provided by the caller). In a real
+    # training loop the caller feeds x_t into the model to get logits, THEN
+    # calls this function. See the orchestrator wiring note in training.py.
+    return mdlm_rb_loss(
+        logits=logits,
+        targets=targets,
+        mask_positions=mask_positions,
+        loss_weights=loss_weights,
+        ignore_index=ignore_index,
+    )

overlay/hydra/engram.py

@@ -1,177 +1,177 @@
-"""GPU Engram — Sparse Modern Hopfield retrieval path.
-
-## What changed (scatter-gather → Hopfield matmul)
-
-The original forward used `self.memory[indices]` (scatter-gather), which misses
-L2 cache at n_columns > 4096 and creates a hard tps ceiling.
-
-The replacement uses:
-    scores   = x @ self.memory.T          # (B, T, n_columns) — coalesced matmul
-    weights  = entmax15(scores, dim=-1)   # sparse attention; 95%+ exact zeros
-    retrieved = weights @ self.memory     # (B, T, d_model)   — coalesced matmul
-
-Both matmuls are tile-friendly (cuBLAS GEMM), so L2 reuse is high regardless of
-n_columns. Gradient flows through both matmuls so `self.memory` learns via
-autograd in addition to (or instead of) the Hebbian EMA writes.
-
-## Sparsity mechanism
-
-alpha-entmax with alpha=1.5 (entmax15) is a sparse attention operator that maps
-logit vectors to distributions where many entries are *exactly* zero (not merely
-small). It generalises softmax (alpha=1) and argmax (alpha→∞). At n_columns=1024
-with d_model=64 a random batch typically hits ≥95% zero entries — the key
-property that keeps bandwidth proportional to *attended* columns, not all columns.
-
-Fallback: if `entmax` is not pip-installed, top-k softmax (k=32) is used instead.
-This is chosen at module-import time — NO runtime branching per forward call.
-
-## token_ids argument
-
-token_ids is accepted for API compatibility with the rest of the hydra stack
-(train.py, lightning_module.py call `engram(x, token_ids)`). It is NOT used in
-the retrieval path — the Hopfield path computes dense similarity over the whole
-memory bank, which subsumes any hash-based column selection. Documented here to
-prevent confusion.
-
-## Hebbian writes (hebbian_boost=False by default)
-
-With Hopfield retrieval, gradient signals reach self.memory through autograd, so
-Hebbian EMA writes are no longer critical. They are preserved as an *optional*
-boost (hebbian_boost=True) for experiments that want both signals. Default is off.
-
-## Checkpoint compatibility
-
-`self.memory` shape (n_columns, d_model) is unchanged, so existing .pt / .ckpt
-files load without modification.
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-# ---------------------------------------------------------------------------
-# Sparse-attention backend — chosen ONCE at import time, no runtime branching.
-# ---------------------------------------------------------------------------
-
-try:
-    from entmax import entmax15 as _entmax15  # type: ignore[import]
-
-    def _sparse_attention(scores: torch.Tensor) -> torch.Tensor:
-        """alpha-entmax (alpha=1.5): truly sparse distribution over last dim."""
-        return _entmax15(scores, dim=-1).to(dtype=scores.dtype)
-
-    _BACKEND = "entmax15"
-
-except ImportError:  # pragma: no cover — entmax always installed in CI
-    _K = 32  # top-k for fallback
-
-    def _sparse_attention(scores: torch.Tensor) -> torch.Tensor:  # type: ignore[misc]
-        """Top-k softmax fallback: zero outside the k highest-scoring columns."""
-        topk_vals, topk_idx = scores.topk(_K, dim=-1)
-        topk_w = torch.softmax(topk_vals, dim=-1)
-        weights = torch.zeros_like(scores)
-        weights.scatter_(-1, topk_idx, topk_w.to(dtype=weights.dtype))
-        return weights
-
-    _BACKEND = "topk32"
-
-
-class GPUEngram(nn.Module):
-    """GPU Engram: Sparse Modern Hopfield retrieval.
-
-    Args:
-        d_model:       Model dimension — must match the surrounding transformer.
-        n_columns:     Number of memory columns (key-value pairs). Safe at 32 768
-                       with the matmul path; the old scatter-gather had an L2
-                       cliff above ~4 096.
-        max_ngram:     Retained for API compatibility; unused in retrieval path.
-        hebbian_boost: If True, also run a Hebbian EMA write on the memory bank
-                       during training (old behaviour, now optional). Default False.
-    """
-
-    def __init__(
-        self,
-        d_model: int,
-        n_columns: int = 1024,
-        max_ngram: int = 3,
-        hebbian_boost: bool = False,
-    ) -> None:
-        super().__init__()
-        self.n_columns = n_columns
-        self.max_ngram = max_ngram
-        self.hebbian_boost = hebbian_boost
-        # Shape unchanged from original — existing checkpoints load cleanly.
-        self.memory = nn.Parameter(torch.randn(n_columns, d_model) * 0.01)
-        self.gate = nn.Linear(d_model, 1, bias=True)
-        nn.init.constant_(self.gate.bias, 0.0)  # START OPEN
-        # Retained for any external code that reads these attrs.
-        self.primes = [2654435761, 2246822519, 3266489917]
-        self.hebbian_lr = 0.01
-
-    # ------------------------------------------------------------------
-    # _hash: retained for API/checkpoint compat; unused in forward below.
-    # ------------------------------------------------------------------
-
-    def _hash(self, token_ids: torch.Tensor) -> torch.Tensor:
-        """N-gram hash → column index (kept for backward-compat; not used in retrieval)."""
-        B, T = token_ids.shape
-        h = token_ids * self.primes[0]
-        if T > 1:
-            shifted1 = torch.roll(token_ids, 1, dims=1)
-            shifted1[:, 0] = 0
-            h = h ^ (shifted1 * self.primes[1])
-        if T > 2:
-            shifted2 = torch.roll(token_ids, 2, dims=1)
-            shifted2[:, :2] = 0
-            h = h ^ (shifted2 * self.primes[2])
-        return h % self.n_columns
-
-    # ------------------------------------------------------------------
-    # forward
-    # ------------------------------------------------------------------
-
-    def forward(self, x: torch.Tensor, token_ids: torch.Tensor):
-        """Hopfield retrieve + soft gate + residual.
-
-        Args:
-            x:         (B, T, d_model) — input activations.
-            token_ids: (B, T) — token indices. Accepted for API compatibility;
-                       NOT used in the retrieval path (see module docstring).
-
-        Returns:
-            (x + alpha * retrieved, hit_rate)
-            - x + alpha * retrieved: (B, T, d_model)
-            - hit_rate: scalar tensor — fraction of gate values > 0.1
-        """
-        # ---- 1. Similarity scores (coalesced GEMM) ----------------------
-        # scores[b, t, c] = dot(x[b,t], memory[c])
-        scores = x @ self.memory.T  # (B, T, n_columns)
-
-        # ---- 2. Sparse attention weights --------------------------------
-        # _sparse_attention is fixed at import time (entmax15 or top-k).
-        weights = _sparse_attention(scores)  # (B, T, n_columns), many exact zeros
-
-        # ---- 3. Retrieved vector (coalesced GEMM) -----------------------
-        retrieved = weights @ self.memory  # (B, T, d_model)
-
-        # ---- 4. Soft gate (unchanged) -----------------------------------
-        alpha = torch.sigmoid(self.gate(x))  # (B, T, 1)
-
-        # ---- 5. Optional Hebbian EMA write ------------------------------
-        if self.training and self.hebbian_boost:
-            with torch.no_grad():
-                # Reuse the hash-based indices for the write target (sparse update).
-                indices = self._hash(token_ids)
-                flat_idx = indices.reshape(-1)           # (B*T,)
-                flat_x = x.detach().reshape(-1, x.shape[-1])  # (B*T, d_model)
-                mem_dtype = self.memory.data.dtype
-                updates = (
-                    self.hebbian_lr * flat_x
-                    - self.hebbian_lr * self.memory.data[flat_idx]
-                ).to(mem_dtype)
-                self.memory.data.index_add_(0, flat_idx, updates)
-
-        # ---- 6. Residual + hit_rate -------------------------------------
-        hit_rate = (alpha.detach() > 0.1).float().mean()
-        return x + alpha * retrieved, hit_rate
+"""GPU Engram — Sparse Modern Hopfield retrieval path.
+
+## What changed (scatter-gather → Hopfield matmul)
+
+The original forward used `self.memory[indices]` (scatter-gather), which misses
+L2 cache at n_columns > 4096 and creates a hard tps ceiling.
+
+The replacement uses:
+    scores   = x @ self.memory.T          # (B, T, n_columns) — coalesced matmul
+    weights  = entmax15(scores, dim=-1)   # sparse attention; 95%+ exact zeros
+    retrieved = weights @ self.memory     # (B, T, d_model)   — coalesced matmul
+
+Both matmuls are tile-friendly (cuBLAS GEMM), so L2 reuse is high regardless of
+n_columns. Gradient flows through both matmuls so `self.memory` learns via
+autograd in addition to (or instead of) the Hebbian EMA writes.
+
+## Sparsity mechanism
+
+alpha-entmax with alpha=1.5 (entmax15) is a sparse attention operator that maps
+logit vectors to distributions where many entries are *exactly* zero (not merely
+small). It generalises softmax (alpha=1) and argmax (alpha→∞). At n_columns=1024
+with d_model=64 a random batch typically hits ≥95% zero entries — the key
+property that keeps bandwidth proportional to *attended* columns, not all columns.
+
+Fallback: if `entmax` is not pip-installed, top-k softmax (k=32) is used instead.
+This is chosen at module-import time — NO runtime branching per forward call.
+
+## token_ids argument
+
+token_ids is accepted for API compatibility with the rest of the hydra stack
+(train.py, lightning_module.py call `engram(x, token_ids)`). It is NOT used in
+the retrieval path — the Hopfield path computes dense similarity over the whole
+memory bank, which subsumes any hash-based column selection. Documented here to
+prevent confusion.
+
+## Hebbian writes (hebbian_boost=False by default)
+
+With Hopfield retrieval, gradient signals reach self.memory through autograd, so
+Hebbian EMA writes are no longer critical. They are preserved as an *optional*
+boost (hebbian_boost=True) for experiments that want both signals. Default is off.
+
+## Checkpoint compatibility
+
+`self.memory` shape (n_columns, d_model) is unchanged, so existing .pt / .ckpt
+files load without modification.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+# ---------------------------------------------------------------------------
+# Sparse-attention backend — chosen ONCE at import time, no runtime branching.
+# ---------------------------------------------------------------------------
+
+try:
+    from entmax import entmax15 as _entmax15  # type: ignore[import]
+
+    def _sparse_attention(scores: torch.Tensor) -> torch.Tensor:
+        """alpha-entmax (alpha=1.5): truly sparse distribution over last dim."""
+        return _entmax15(scores, dim=-1).to(dtype=scores.dtype)
+
+    _BACKEND = "entmax15"
+
+except ImportError:  # pragma: no cover — entmax always installed in CI
+    _K = 32  # top-k for fallback
+
+    def _sparse_attention(scores: torch.Tensor) -> torch.Tensor:  # type: ignore[misc]
+        """Top-k softmax fallback: zero outside the k highest-scoring columns."""
+        topk_vals, topk_idx = scores.topk(_K, dim=-1)
+        topk_w = torch.softmax(topk_vals, dim=-1)
+        weights = torch.zeros_like(scores)
+        weights.scatter_(-1, topk_idx, topk_w.to(dtype=weights.dtype))
+        return weights
+
+    _BACKEND = "topk32"
+
+
+class GPUEngram(nn.Module):
+    """GPU Engram: Sparse Modern Hopfield retrieval.
+
+    Args:
+        d_model:       Model dimension — must match the surrounding transformer.
+        n_columns:     Number of memory columns (key-value pairs). Safe at 32 768
+                       with the matmul path; the old scatter-gather had an L2
+                       cliff above ~4 096.
+        max_ngram:     Retained for API compatibility; unused in retrieval path.
+        hebbian_boost: If True, also run a Hebbian EMA write on the memory bank
+                       during training (old behaviour, now optional). Default False.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_columns: int = 1024,
+        max_ngram: int = 3,
+        hebbian_boost: bool = False,
+    ) -> None:
+        super().__init__()
+        self.n_columns = n_columns
+        self.max_ngram = max_ngram
+        self.hebbian_boost = hebbian_boost
+        # Shape unchanged from original — existing checkpoints load cleanly.
+        self.memory = nn.Parameter(torch.randn(n_columns, d_model) * 0.01)
+        self.gate = nn.Linear(d_model, 1, bias=True)
+        nn.init.constant_(self.gate.bias, 0.0)  # START OPEN
+        # Retained for any external code that reads these attrs.
+        self.primes = [2654435761, 2246822519, 3266489917]
+        self.hebbian_lr = 0.01
+
+    # ------------------------------------------------------------------
+    # _hash: retained for API/checkpoint compat; unused in forward below.
+    # ------------------------------------------------------------------
+
+    def _hash(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """N-gram hash → column index (kept for backward-compat; not used in retrieval)."""
+        B, T = token_ids.shape
+        h = token_ids * self.primes[0]
+        if T > 1:
+            shifted1 = torch.roll(token_ids, 1, dims=1)
+            shifted1[:, 0] = 0
+            h = h ^ (shifted1 * self.primes[1])
+        if T > 2:
+            shifted2 = torch.roll(token_ids, 2, dims=1)
+            shifted2[:, :2] = 0
+            h = h ^ (shifted2 * self.primes[2])
+        return h % self.n_columns
+
+    # ------------------------------------------------------------------
+    # forward
+    # ------------------------------------------------------------------
+
+    def forward(self, x: torch.Tensor, token_ids: torch.Tensor):
+        """Hopfield retrieve + soft gate + residual.
+
+        Args:
+            x:         (B, T, d_model) — input activations.
+            token_ids: (B, T) — token indices. Accepted for API compatibility;
+                       NOT used in the retrieval path (see module docstring).
+
+        Returns:
+            (x + alpha * retrieved, hit_rate)
+            - x + alpha * retrieved: (B, T, d_model)
+            - hit_rate: scalar tensor — fraction of gate values > 0.1
+        """
+        # ---- 1. Similarity scores (coalesced GEMM) ----------------------
+        # scores[b, t, c] = dot(x[b,t], memory[c])
+        scores = x @ self.memory.T  # (B, T, n_columns)
+
+        # ---- 2. Sparse attention weights --------------------------------
+        # _sparse_attention is fixed at import time (entmax15 or top-k).
+        weights = _sparse_attention(scores)  # (B, T, n_columns), many exact zeros
+
+        # ---- 3. Retrieved vector (coalesced GEMM) -----------------------
+        retrieved = weights @ self.memory  # (B, T, d_model)
+
+        # ---- 4. Soft gate (unchanged) -----------------------------------
+        alpha = torch.sigmoid(self.gate(x))  # (B, T, 1)
+
+        # ---- 5. Optional Hebbian EMA write ------------------------------
+        if self.training and self.hebbian_boost:
+            with torch.no_grad():
+                # Reuse the hash-based indices for the write target (sparse update).
+                indices = self._hash(token_ids)
+                flat_idx = indices.reshape(-1)           # (B*T,)
+                flat_x = x.detach().reshape(-1, x.shape[-1])  # (B*T, d_model)
+                mem_dtype = self.memory.data.dtype
+                updates = (
+                    self.hebbian_lr * flat_x
+                    - self.hebbian_lr * self.memory.data[flat_idx]
+                ).to(mem_dtype)
+                self.memory.data.index_add_(0, flat_idx, updates)
+
+        # ---- 6. Residual + hit_rate -------------------------------------
+        hit_rate = (alpha.detach() > 0.1).float().mean()
+        return x + alpha * retrieved, hit_rate

overlay/hydra/eval.py

@@ -1,210 +1,210 @@
-"""Evaluation: factual probes + sampled factual English scoring.
-
-Extracted from train.py (W1 modularization). Semantics unchanged.
-
-Perf optimizations (eval_perf_fix):
-- Probe mode: single forward per prompt instead of autoregressive gen
-- Batch decode: all GPU work first, all CPU decode after
-- Batched factual probes: single padded forward instead of N sequential
-"""
-
-from __future__ import annotations
-
-import os
-import re as _re
-
-import torch
-
-from hydra.config import FACTUAL_SAMPLES, FACTUAL_BATCH, FACTUAL_GEN_TOKENS
-
-# Default to probe mode (1 forward per prompt); set HYDRA_FACTUAL_MODE=gen for
-# the original autoregressive generation path.
-FACTUAL_MODE = os.environ.get("HYDRA_FACTUAL_MODE", "probe")
-
-FACTUAL_EVAL = [
-    # Hard factual recall — requires specific knowledge memorization
-    ("The capital of France is", ["Paris", "paris"]),
-    ("Water boils at", ["100", "boiling"]),
-    ("The largest planet in our solar system is", ["Jupiter", "jupiter"]),
-    # Easier completions — common collocations / patterns the model may pick up
-    ("Once upon a", ["time"]),
-    ("Hello, my name", ["is", "'s"]),
-    ("The cat sat on the", ["mat", "floor", "rug", "table", "couch", "chair", "ground"]),
-    ("She opened the door and", ["walked", "saw", "found", "stepped", "looked", "went", "ran"]),
-    # Original hard ones kept for completeness
-    ("The speed of light is approximately", ["299", "300", "186,000", "light speed"]),
-    ("Two plus two equals", ["4", "four"]),
-]
-
-_FACTUAL_PROBES = [
-    "The capital of France is",
-    "Water boils at",
-    "The largest planet in our solar system is",
-    "The speed of light is approximately",
-    "Shakespeare wrote",
-]
-
-
-def run_factual_probes(model, tokenizer, device, autocast_ctx) -> None:
-    """Top-5 next-token predictions for canonical factual prompts.
-
-    Batched: pads all prompts into a single forward pass instead of N
-    sequential passes.
-    """
-    print("\n--- Factual Probes ---")
-    model.eval()
-
-    # Process probes one at a time to avoid cooperative launch limit
-    # (batched forward with B=len(probes) can exceed SM residency cap).
-    for prompt_text in _FACTUAL_PROBES:
-        ids = tokenizer.encode(prompt_text)
-        x = torch.tensor([ids], device=device)
-        with torch.no_grad(), autocast_ctx:
-            logits = model(x)
-        probs = torch.softmax(logits[0, -1].float(), dim=-1)
-        top5 = torch.topk(probs, 5)
-        completions = [tokenizer.decode([idx.item()]) for idx in top5.indices]
-        probs_list = [f"{p:.4f}" for p in top5.values[:3].tolist()]
-        print(f'  "{prompt_text}" -> {completions[:3]} (p={probs_list})')
-    print("--- End Factual Probes ---\n")
-
-
-# ---------------------------------------------------------------------------
-# Probe mode: single forward per prompt (Fix D)
-# ---------------------------------------------------------------------------
-
-def _run_factual_english_probe(model, tokenizer, max_seq_len: int):
-    """Fast probe mode: for each (prompt, answers), encode prompt + each answer
-    candidate as a single sequence, do ONE forward pass, and check if the model's
-    argmax at the last prompt token matches the first answer token.
-
-    Falls back to checking top-K predictions to be generous (same as gen mode
-    which samples multiple temperatures).
-    """
-    print("---")
-    print("factual_english_samples: (probe mode)")
-    model.eval()
-    hits = 0
-
-    with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-        for prompt, answers in FACTUAL_EVAL:
-            prompt_ids = tokenizer.encode(prompt)
-            prompt_len = len(prompt_ids)
-            x = torch.tensor([prompt_ids], device="cuda", dtype=torch.long)
-            logits = model(x, targets=None)
-            # logits shape: [1, seq_len, vocab] or [1, vocab]
-            if logits.dim() == 3:
-                last_logits = logits[0, -1, :]
-            else:
-                last_logits = logits[0]
-
-            probs = torch.softmax(last_logits.float(), dim=-1)
-            # Check top-K predictions (generous: K=20 to match multi-sample gen)
-            top_k = min(20, probs.shape[-1])
-            top_ids = torch.topk(probs, top_k).indices.tolist()
-            top_tokens = [tokenizer.decode([tid]).strip().lower() for tid in top_ids]
-
-            answers_lower = [a.lower() for a in answers]
-            any_hit = any(
-                any(a in tok for a in answers_lower)
-                for tok in top_tokens
-            )
-            if any_hit:
-                hits += 1
-
-            best_completion = tokenizer.decode([top_ids[0]])
-            print(f"  prompt: {prompt!r}")
-            print(f"  output: {(prompt + best_completion).replace(chr(10), ' ')!r}")
-            print(f"  hit:    {any_hit} (probe top-{top_k})")
-
-    score = hits / len(FACTUAL_EVAL)
-    print("---")
-    print(f"factual_english_score: {score:.4f}")
-    print(f"factual_english_hits:  {hits}/{len(FACTUAL_EVAL)}")
-    return score, hits, len(FACTUAL_EVAL)
-
-
-# ---------------------------------------------------------------------------
-# Gen mode: original autoregressive path (Fix F: batch decode)
-# ---------------------------------------------------------------------------
-
-def _run_factual_english_gen(model, tokenizer, max_seq_len: int):
-    """Original autoregressive generation path with batch decode optimization:
-    all GPU work runs first, then all CPU decoding happens after."""
-    print("---")
-    print("factual_english_samples: (gen mode)")
-    model.eval()
-
-    num_samples = FACTUAL_SAMPLES
-    batch = FACTUAL_BATCH
-    gen_tokens = FACTUAL_GEN_TOKENS
-    temps = [0.7, 0.9, 1.1]
-    hits = 0
-
-    with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-        for prompt, answers in FACTUAL_EVAL:
-            ids = tokenizer.encode(prompt)
-            answers_lower = [a.lower() for a in answers]
-            # Collect all generated token sequences on GPU first
-            all_rows: list[list[int]] = []
-            samples_done = 0
-            batch_idx = 0
-            while samples_done < num_samples:
-                b = min(batch, num_samples - samples_done)
-                temp = temps[batch_idx % len(temps)]
-                batch_idx += 1
-                ctx = torch.tensor([ids] * b, device="cuda", dtype=torch.long)
-                for _ in range(gen_tokens):
-                    logits = model(ctx, targets=None)
-                    next_logits = logits[:, -1, :] if logits.dim() == 3 else logits
-                    probs = torch.softmax(next_logits.float() / temp, dim=-1)
-                    next_id = torch.multinomial(probs, num_samples=1)
-                    ctx = torch.cat([ctx, next_id], dim=1)
-                    if ctx.size(1) >= max_seq_len:
-                        break
-                # Transfer to CPU in one shot, no per-row sync
-                all_rows.extend(ctx.cpu().tolist())
-                samples_done += b
-
-            # CPU-side batch decode — no GPU sync between decodes
-            any_hit = False
-            first_gen = None
-            hit_gen = None
-            for row in all_rows:
-                generated = tokenizer.decode(row)
-                continuation = generated[len(prompt):].strip()
-                _words = set(w.lower() for w in _re.findall(r"\b[\w'-]+\b", continuation))
-                hit = any(a in _words for a in answers_lower)
-                if first_gen is None:
-                    first_gen = generated
-                if hit:
-                    any_hit = True
-                    if hit_gen is None:
-                        hit_gen = generated
-            if any_hit:
-                hits += 1
-            print(f"  prompt: {prompt!r}")
-            print(f"  output: {(first_gen or '').replace(chr(10), ' ')!r}")
-            print(f"  hit:    {any_hit} (any of {num_samples} samples, temps={temps}, gen={gen_tokens}tok)")
-            if hit_gen is not None and hit_gen != first_gen:
-                print(f"  hit_sample: {hit_gen.replace(chr(10), ' ')!r}")
-
-    score = hits / len(FACTUAL_EVAL)
-    print("---")
-    print(f"factual_english_score: {score:.4f}")
-    print(f"factual_english_hits:  {hits}/{len(FACTUAL_EVAL)}")
-    return score, hits, len(FACTUAL_EVAL)
-
-
-# ---------------------------------------------------------------------------
-# Public entry point
-# ---------------------------------------------------------------------------
-
-def run_factual_english(model, tokenizer, max_seq_len: int):
-    """Dispatch to probe (fast, default) or gen (original) mode.
-
-    Set HYDRA_FACTUAL_MODE=gen to use the autoregressive path.
-    """
-    if FACTUAL_MODE == "gen":
-        return _run_factual_english_gen(model, tokenizer, max_seq_len)
-    return _run_factual_english_probe(model, tokenizer, max_seq_len)
+"""Evaluation: factual probes + sampled factual English scoring.
+
+Extracted from train.py (W1 modularization). Semantics unchanged.
+
+Perf optimizations (eval_perf_fix):
+- Probe mode: single forward per prompt instead of autoregressive gen
+- Batch decode: all GPU work first, all CPU decode after
+- Batched factual probes: single padded forward instead of N sequential
+"""
+
+from __future__ import annotations
+
+import os
+import re as _re
+
+import torch
+
+from hydra.config import FACTUAL_SAMPLES, FACTUAL_BATCH, FACTUAL_GEN_TOKENS
+
+# Default to probe mode (1 forward per prompt); set HYDRA_FACTUAL_MODE=gen for
+# the original autoregressive generation path.
+FACTUAL_MODE = os.environ.get("HYDRA_FACTUAL_MODE", "probe")
+
+FACTUAL_EVAL = [
+    # Hard factual recall — requires specific knowledge memorization
+    ("The capital of France is", ["Paris", "paris"]),
+    ("Water boils at", ["100", "boiling"]),
+    ("The largest planet in our solar system is", ["Jupiter", "jupiter"]),
+    # Easier completions — common collocations / patterns the model may pick up
+    ("Once upon a", ["time"]),
+    ("Hello, my name", ["is", "'s"]),
+    ("The cat sat on the", ["mat", "floor", "rug", "table", "couch", "chair", "ground"]),
+    ("She opened the door and", ["walked", "saw", "found", "stepped", "looked", "went", "ran"]),
+    # Original hard ones kept for completeness
+    ("The speed of light is approximately", ["299", "300", "186,000", "light speed"]),
+    ("Two plus two equals", ["4", "four"]),
+]
+
+_FACTUAL_PROBES = [
+    "The capital of France is",
+    "Water boils at",
+    "The largest planet in our solar system is",
+    "The speed of light is approximately",
+    "Shakespeare wrote",
+]
+
+
+def run_factual_probes(model, tokenizer, device, autocast_ctx) -> None:
+    """Top-5 next-token predictions for canonical factual prompts.
+
+    Batched: pads all prompts into a single forward pass instead of N
+    sequential passes.
+    """
+    print("\n--- Factual Probes ---")
+    model.eval()
+
+    # Process probes one at a time to avoid cooperative launch limit
+    # (batched forward with B=len(probes) can exceed SM residency cap).
+    for prompt_text in _FACTUAL_PROBES:
+        ids = tokenizer.encode(prompt_text)
+        x = torch.tensor([ids], device=device)
+        with torch.no_grad(), autocast_ctx:
+            logits = model(x)
+        probs = torch.softmax(logits[0, -1].float(), dim=-1)
+        top5 = torch.topk(probs, 5)
+        completions = [tokenizer.decode([idx.item()]) for idx in top5.indices]
+        probs_list = [f"{p:.4f}" for p in top5.values[:3].tolist()]
+        print(f'  "{prompt_text}" -> {completions[:3]} (p={probs_list})')
+    print("--- End Factual Probes ---\n")
+
+
+# ---------------------------------------------------------------------------
+# Probe mode: single forward per prompt (Fix D)
+# ---------------------------------------------------------------------------
+
+def _run_factual_english_probe(model, tokenizer, max_seq_len: int):
+    """Fast probe mode: for each (prompt, answers), encode prompt + each answer
+    candidate as a single sequence, do ONE forward pass, and check if the model's
+    argmax at the last prompt token matches the first answer token.
+
+    Falls back to checking top-K predictions to be generous (same as gen mode
+    which samples multiple temperatures).
+    """
+    print("---")
+    print("factual_english_samples: (probe mode)")
+    model.eval()
+    hits = 0
+
+    with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+        for prompt, answers in FACTUAL_EVAL:
+            prompt_ids = tokenizer.encode(prompt)
+            prompt_len = len(prompt_ids)
+            x = torch.tensor([prompt_ids], device="cuda", dtype=torch.long)
+            logits = model(x, targets=None)
+            # logits shape: [1, seq_len, vocab] or [1, vocab]
+            if logits.dim() == 3:
+                last_logits = logits[0, -1, :]
+            else:
+                last_logits = logits[0]
+
+            probs = torch.softmax(last_logits.float(), dim=-1)
+            # Check top-K predictions (generous: K=20 to match multi-sample gen)
+            top_k = min(20, probs.shape[-1])
+            top_ids = torch.topk(probs, top_k).indices.tolist()
+            top_tokens = [tokenizer.decode([tid]).strip().lower() for tid in top_ids]
+
+            answers_lower = [a.lower() for a in answers]
+            any_hit = any(
+                any(a in tok for a in answers_lower)
+                for tok in top_tokens
+            )
+            if any_hit:
+                hits += 1
+
+            best_completion = tokenizer.decode([top_ids[0]])
+            print(f"  prompt: {prompt!r}")
+            print(f"  output: {(prompt + best_completion).replace(chr(10), ' ')!r}")
+            print(f"  hit:    {any_hit} (probe top-{top_k})")
+
+    score = hits / len(FACTUAL_EVAL)
+    print("---")
+    print(f"factual_english_score: {score:.4f}")
+    print(f"factual_english_hits:  {hits}/{len(FACTUAL_EVAL)}")
+    return score, hits, len(FACTUAL_EVAL)
+
+
+# ---------------------------------------------------------------------------
+# Gen mode: original autoregressive path (Fix F: batch decode)
+# ---------------------------------------------------------------------------
+
+def _run_factual_english_gen(model, tokenizer, max_seq_len: int):
+    """Original autoregressive generation path with batch decode optimization:
+    all GPU work runs first, then all CPU decoding happens after."""
+    print("---")
+    print("factual_english_samples: (gen mode)")
+    model.eval()
+
+    num_samples = FACTUAL_SAMPLES
+    batch = FACTUAL_BATCH
+    gen_tokens = FACTUAL_GEN_TOKENS
+    temps = [0.7, 0.9, 1.1]
+    hits = 0
+
+    with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+        for prompt, answers in FACTUAL_EVAL:
+            ids = tokenizer.encode(prompt)
+            answers_lower = [a.lower() for a in answers]
+            # Collect all generated token sequences on GPU first
+            all_rows: list[list[int]] = []
+            samples_done = 0
+            batch_idx = 0
+            while samples_done < num_samples:
+                b = min(batch, num_samples - samples_done)
+                temp = temps[batch_idx % len(temps)]
+                batch_idx += 1
+                ctx = torch.tensor([ids] * b, device="cuda", dtype=torch.long)
+                for _ in range(gen_tokens):
+                    logits = model(ctx, targets=None)
+                    next_logits = logits[:, -1, :] if logits.dim() == 3 else logits
+                    probs = torch.softmax(next_logits.float() / temp, dim=-1)
+                    next_id = torch.multinomial(probs, num_samples=1)
+                    ctx = torch.cat([ctx, next_id], dim=1)
+                    if ctx.size(1) >= max_seq_len:
+                        break
+                # Transfer to CPU in one shot, no per-row sync
+                all_rows.extend(ctx.cpu().tolist())
+                samples_done += b
+
+            # CPU-side batch decode — no GPU sync between decodes
+            any_hit = False
+            first_gen = None
+            hit_gen = None
+            for row in all_rows:
+                generated = tokenizer.decode(row)
+                continuation = generated[len(prompt):].strip()
+                _words = set(w.lower() for w in _re.findall(r"\b[\w'-]+\b", continuation))
+                hit = any(a in _words for a in answers_lower)
+                if first_gen is None:
+                    first_gen = generated
+                if hit:
+                    any_hit = True
+                    if hit_gen is None:
+                        hit_gen = generated
+            if any_hit:
+                hits += 1
+            print(f"  prompt: {prompt!r}")
+            print(f"  output: {(first_gen or '').replace(chr(10), ' ')!r}")
+            print(f"  hit:    {any_hit} (any of {num_samples} samples, temps={temps}, gen={gen_tokens}tok)")
+            if hit_gen is not None and hit_gen != first_gen:
+                print(f"  hit_sample: {hit_gen.replace(chr(10), ' ')!r}")
+
+    score = hits / len(FACTUAL_EVAL)
+    print("---")
+    print(f"factual_english_score: {score:.4f}")
+    print(f"factual_english_hits:  {hits}/{len(FACTUAL_EVAL)}")
+    return score, hits, len(FACTUAL_EVAL)
+
+
+# ---------------------------------------------------------------------------
+# Public entry point
+# ---------------------------------------------------------------------------
+
+def run_factual_english(model, tokenizer, max_seq_len: int):
+    """Dispatch to probe (fast, default) or gen (original) mode.
+
+    Set HYDRA_FACTUAL_MODE=gen to use the autoregressive path.
+    """
+    if FACTUAL_MODE == "gen":
+        return _run_factual_english_gen(model, tokenizer, max_seq_len)
+    return _run_factual_english_probe(model, tokenizer, max_seq_len)

overlay/hydra/gdn_block.py

@@ -1,126 +1,126 @@
-"""GDNBlock — Gated Delta Net block, drop-in shape-compatible with Mamba3Block and HyenaBlock.
-
-GatedDeltaNet (GDN) reference: arXiv:2412.06464 (ICLR 2025, NVLabs).
-Implementation: flash-linear-attention (fla) library, Triton kernels, sm86-compatible.
-
-Interface contract (MUST match how Mamba3/Hyena are called in hydra/model.py):
-    block = GDNBlock(d_model, ...)
-    y = block(x)    # x: [B, T, d_model]  ->  y: [B, T, d_model]
-
-The surrounding mHC layer does NOT pre-norm before calling this block (the
-raw hidden state is passed in); the block itself applies no input normalization,
-same as HyenaBlock.  We return the raw operator output; the mHC layer adds it
-as a residual stream contribution.
-
-NO attention, NO softmax-over-sequence-dim.  All state is stateless between
-.forward() calls by default (use_cache=False, past_key_values=None).
-"""
-
-from __future__ import annotations
-
-try:
-    from fla.layers.gated_deltanet import GatedDeltaNet as _GatedDeltaNet
-except ImportError as _fla_err:
-    raise ImportError(
-        "flash-linear-attention (fla) is required for GDNBlock but could not be imported. "
-        "Install it with:\n"
-        "    pip install flash-linear-attention\n"
-        "or from source:\n"
-        "    pip install git+https://github.com/fla-org/flash-linear-attention.git\n"
-        f"Original error: {_fla_err}"
-    ) from _fla_err
-
-import torch
-import torch.nn as nn
-
-
-class GDNBlock(nn.Module):
-    """Gated Delta Net block, drop-in shape-compatible with HYDRA's Mamba3Block and HyenaBlock.
-
-    Wraps `fla.layers.GatedDeltaNet` with the same external API that
-    `hydra.hyena_block.HyenaBlock` exposes:
-
-        forward(x: Tensor[B, T, d_model]) -> Tensor[B, T, d_model]
-
-    Internal GatedDeltaNet.forward returns a 3-tuple
-    (hidden_states, attn_weights, past_key_values); we extract [0] and
-    return only the hidden states, keeping the residual stream unchanged.
-
-    GDN outperforms Mamba-2 on in-context retrieval benchmarks (MQAR, etc.)
-    at equal or faster compute, making it a targeted fix for HYDRA's factual
-    plateau.
-
-    Parameter counts are deliberately kept within 2x of a Mamba3 block at the
-    same d_model/n_heads to be drop-in affordable.
-    """
-
-    def __init__(
-        self,
-        d_model: int,
-        n_heads: int = 6,
-        mode: str = "chunk",       # 'chunk' for training, 'fused_recurrent' for inference
-        expand_v: float = 2.0,     # value-projection expansion; controls KV memory
-        use_short_conv: bool = True,
-        conv_size: int = 4,
-    ):
-        super().__init__()
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.mode = mode
-
-        # head_dim must divide d_model.  GDN uses separate q/k head_dim from v;
-        # we set head_dim for q/k such that n_heads * head_dim == d_model.
-        if d_model % n_heads != 0:
-            raise ValueError(
-                f"d_model={d_model} must be divisible by n_heads={n_heads} "
-                "so that head_dim = d_model // n_heads is an integer."
-            )
-        head_dim = d_model // n_heads
-
-        self.gdn = _GatedDeltaNet(
-            hidden_size=d_model,
-            expand_v=expand_v,
-            head_dim=head_dim,
-            num_heads=n_heads,
-            mode=mode,
-            use_gate=True,          # gating is the key architectural feature of GDN
-            use_short_conv=use_short_conv,
-            conv_size=conv_size,
-            layer_idx=None,         # no KV-cache layer indexing; we manage state ourselves
-        )
-
-    # ------------------------------------------------------------------
-    # Forward
-    # ------------------------------------------------------------------
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """x: [B, T, d_model]  ->  y: [B, T, d_model].
-
-        Passes through GatedDeltaNet with use_cache=False so no recurrent
-        state leaks between independent forward() calls (important for
-        gradient-accumulation loops and eval).
-        """
-        # GatedDeltaNet.forward signature:
-        #   (hidden_states, attention_mask=None, past_key_values=None,
-        #    use_cache=False, output_attentions=False)
-        # Returns: tuple(hidden_states, attn_weights|None, past_kv|None)
-        out, _, _ = self.gdn(
-            hidden_states=x,
-            attention_mask=None,
-            past_key_values=None,
-            use_cache=False,
-            output_attentions=False,
-        )
-        return out
-
-    # ------------------------------------------------------------------
-    # API parity with HyenaBlock and Mamba3Block
-    # ------------------------------------------------------------------
-
-    def invalidate_caches(self) -> None:
-        """No-op — GDNBlock holds no persistent filter cache.
-
-        Provided for API parity with HyenaBlock, which invalidates its
-        Hyena filter cache here.  Calling this is always safe.
-        """
-        pass
+"""GDNBlock — Gated Delta Net block, drop-in shape-compatible with Mamba3Block and HyenaBlock.
+
+GatedDeltaNet (GDN) reference: arXiv:2412.06464 (ICLR 2025, NVLabs).
+Implementation: flash-linear-attention (fla) library, Triton kernels, sm86-compatible.
+
+Interface contract (MUST match how Mamba3/Hyena are called in hydra/model.py):
+    block = GDNBlock(d_model, ...)
+    y = block(x)    # x: [B, T, d_model]  ->  y: [B, T, d_model]
+
+The surrounding mHC layer does NOT pre-norm before calling this block (the
+raw hidden state is passed in); the block itself applies no input normalization,
+same as HyenaBlock.  We return the raw operator output; the mHC layer adds it
+as a residual stream contribution.
+
+NO attention, NO softmax-over-sequence-dim.  All state is stateless between
+.forward() calls by default (use_cache=False, past_key_values=None).
+"""
+
+from __future__ import annotations
+
+try:
+    from fla.layers.gated_deltanet import GatedDeltaNet as _GatedDeltaNet
+except ImportError as _fla_err:
+    raise ImportError(
+        "flash-linear-attention (fla) is required for GDNBlock but could not be imported. "
+        "Install it with:\n"
+        "    pip install flash-linear-attention\n"
+        "or from source:\n"
+        "    pip install git+https://github.com/fla-org/flash-linear-attention.git\n"
+        f"Original error: {_fla_err}"
+    ) from _fla_err
+
+import torch
+import torch.nn as nn
+
+
+class GDNBlock(nn.Module):
+    """Gated Delta Net block, drop-in shape-compatible with HYDRA's Mamba3Block and HyenaBlock.
+
+    Wraps `fla.layers.GatedDeltaNet` with the same external API that
+    `hydra.hyena_block.HyenaBlock` exposes:
+
+        forward(x: Tensor[B, T, d_model]) -> Tensor[B, T, d_model]
+
+    Internal GatedDeltaNet.forward returns a 3-tuple
+    (hidden_states, attn_weights, past_key_values); we extract [0] and
+    return only the hidden states, keeping the residual stream unchanged.
+
+    GDN outperforms Mamba-2 on in-context retrieval benchmarks (MQAR, etc.)
+    at equal or faster compute, making it a targeted fix for HYDRA's factual
+    plateau.
+
+    Parameter counts are deliberately kept within 2x of a Mamba3 block at the
+    same d_model/n_heads to be drop-in affordable.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int = 6,
+        mode: str = "chunk",       # 'chunk' for training, 'fused_recurrent' for inference
+        expand_v: float = 2.0,     # value-projection expansion; controls KV memory
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.mode = mode
+
+        # head_dim must divide d_model.  GDN uses separate q/k head_dim from v;
+        # we set head_dim for q/k such that n_heads * head_dim == d_model.
+        if d_model % n_heads != 0:
+            raise ValueError(
+                f"d_model={d_model} must be divisible by n_heads={n_heads} "
+                "so that head_dim = d_model // n_heads is an integer."
+            )
+        head_dim = d_model // n_heads
+
+        self.gdn = _GatedDeltaNet(
+            hidden_size=d_model,
+            expand_v=expand_v,
+            head_dim=head_dim,
+            num_heads=n_heads,
+            mode=mode,
+            use_gate=True,          # gating is the key architectural feature of GDN
+            use_short_conv=use_short_conv,
+            conv_size=conv_size,
+            layer_idx=None,         # no KV-cache layer indexing; we manage state ourselves
+        )
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, T, d_model]  ->  y: [B, T, d_model].
+
+        Passes through GatedDeltaNet with use_cache=False so no recurrent
+        state leaks between independent forward() calls (important for
+        gradient-accumulation loops and eval).
+        """
+        # GatedDeltaNet.forward signature:
+        #   (hidden_states, attention_mask=None, past_key_values=None,
+        #    use_cache=False, output_attentions=False)
+        # Returns: tuple(hidden_states, attn_weights|None, past_kv|None)
+        out, _, _ = self.gdn(
+            hidden_states=x,
+            attention_mask=None,
+            past_key_values=None,
+            use_cache=False,
+            output_attentions=False,
+        )
+        return out
+
+    # ------------------------------------------------------------------
+    # API parity with HyenaBlock and Mamba3Block
+    # ------------------------------------------------------------------
+
+    def invalidate_caches(self) -> None:
+        """No-op — GDNBlock holds no persistent filter cache.
+
+        Provided for API parity with HyenaBlock, which invalidates its
+        Hyena filter cache here.  Calling this is always safe.
+        """
+        pass

overlay/hydra/hyena_block.py

@@ -1,68 +1,68 @@
-"""HyenaBlock — drop-in block for HYDRA, supplement to Mamba3.
-
-Wraps `subsystems.hyena_pure.HyenaOperator` with a pre-norm + residual scheme
-consistent with how the mHC stack wraps Mamba3 in `hydra/model.py`.
-
-Interface contract (MUST match how Mamba3 is called in model.py):
-    block = HyenaBlock(d_model, seq_len)
-    y = block(x)   # x: [B, T, d_model]  ->  y: [B, T, d_model]
-
-The surrounding mHC layer does the pre-norm (`norm(h)`) BEFORE calling the
-block, so the block itself should NOT re-normalize at input — same as Mamba3
-in the current model. We return the raw operator output; the mHC layer then
-adds it as a residual stream contribution.
-
-NO attention, NO softmax-over-sequence-dim, NO KV-cache. All forbidden
-imports enumerated in tests/test_hyena.py (test #7) are absent.
-"""
-
-from __future__ import annotations
-
-import os
-
-import torch
-import torch.nn as nn
-
-from subsystems.hyena_pure import HyenaOperator
-
-
-class HyenaBlock(nn.Module):
-    """Single Hyena block, shape-compatible with Mamba3 in HYDRA."""
-
-    def __init__(
-        self,
-        d_model: int,
-        seq_len: int,
-        order: int | None = None,
-        filter_order: int | None = None,
-        dropout: float = 0.0,
-        filter_dropout: float = 0.0,
-        short_filter_order: int = 3,
-        activation: str = "id",
-    ):
-        super().__init__()
-        # Env overrides (documented in hydra/config.py).
-        if order is None:
-            order = int(os.environ.get("HYDRA_HYENA_ORDER", "2"))
-        if filter_order is None:
-            filter_order = int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64"))
-
-        self.d_model = d_model
-        self.seq_len = seq_len
-        self.order = order
-        self.filter_order = filter_order
-
-        self.operator = HyenaOperator(
-            d_model=d_model,
-            l_max=seq_len,
-            order=order,
-            filter_order=filter_order,
-            dropout=dropout,
-            filter_dropout=filter_dropout,
-            short_filter_order=short_filter_order,
-            activation=activation,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """x: [B, T, d_model]  ->  y: [B, T, d_model]."""
-        return self.operator(x)
+"""HyenaBlock — drop-in block for HYDRA, supplement to Mamba3.
+
+Wraps `subsystems.hyena_pure.HyenaOperator` with a pre-norm + residual scheme
+consistent with how the mHC stack wraps Mamba3 in `hydra/model.py`.
+
+Interface contract (MUST match how Mamba3 is called in model.py):
+    block = HyenaBlock(d_model, seq_len)
+    y = block(x)   # x: [B, T, d_model]  ->  y: [B, T, d_model]
+
+The surrounding mHC layer does the pre-norm (`norm(h)`) BEFORE calling the
+block, so the block itself should NOT re-normalize at input — same as Mamba3
+in the current model. We return the raw operator output; the mHC layer then
+adds it as a residual stream contribution.
+
+NO attention, NO softmax-over-sequence-dim, NO KV-cache. All forbidden
+imports enumerated in tests/test_hyena.py (test #7) are absent.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+
+from subsystems.hyena_pure import HyenaOperator
+
+
+class HyenaBlock(nn.Module):
+    """Single Hyena block, shape-compatible with Mamba3 in HYDRA."""
+
+    def __init__(
+        self,
+        d_model: int,
+        seq_len: int,
+        order: int | None = None,
+        filter_order: int | None = None,
+        dropout: float = 0.0,
+        filter_dropout: float = 0.0,
+        short_filter_order: int = 3,
+        activation: str = "id",
+    ):
+        super().__init__()
+        # Env overrides (documented in hydra/config.py).
+        if order is None:
+            order = int(os.environ.get("HYDRA_HYENA_ORDER", "2"))
+        if filter_order is None:
+            filter_order = int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64"))
+
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.order = order
+        self.filter_order = filter_order
+
+        self.operator = HyenaOperator(
+            d_model=d_model,
+            l_max=seq_len,
+            order=order,
+            filter_order=filter_order,
+            dropout=dropout,
+            filter_dropout=filter_dropout,
+            short_filter_order=short_filter_order,
+            activation=activation,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, T, d_model]  ->  y: [B, T, d_model]."""
+        return self.operator(x)

overlay/hydra/lightning_module.py

@@ -1,326 +1,326 @@
-"""LightningModule wrapping PostSemClawModel.
-
-Thin adapter. The model and the MuonAdamW optimizer are unchanged. This
-module implements:
-
-  • configure_optimizers — returns the existing MuonAdamW (subclass of
-    torch.optim.Optimizer) built by model.setup_optimizer. Lightning accepts
-    this directly.
-  • training_step — splits (B, T+1) batches into (x, y), forwards through
-    the model, logs loss / bpb / tps / mfu / vram. Preserves the
-    sampled-softmax path inside PostSemClawModel (no changes there).
-  • optimizer_step — before each step we update LR + muon momentum + WD
-    using the same time-progress schedule as hydra/training.py
-    (get_lr_multiplier / get_muon_momentum / get_weight_decay). Lightning
-    handles grad accumulation via Trainer(accumulate_grad_batches=N).
-
-The SDR SOM update and Hestia QAT snap are called at the same cadence as
-the legacy loop, but inline on the main thread (Lightning provides its own
-callbacks for async work if we need to extract them later — keeping it
-simple for now).
-
-Env vars respected:
-  HYDRA_TIME_BUDGET          — wall-clock budget (s) used for LR schedule
-                                and as Trainer max_time
-  HYDRA_HESTIA_INTERVAL      — steps between Hestia snaps (default 100)
-  HYDRA_BATCH_SIZE           — device batch size (for throughput calc)
-  HYDRA_SEQ_LEN              — sequence length (for throughput calc)
-"""
-from __future__ import annotations
-
-import math
-import os
-import time
-
-import torch
-import lightning as L
-
-from hydra.config import (
-    ADAM_BETAS,
-    EMBEDDING_LR,
-    FINAL_LR_FRAC,
-    GPU_BF16_PEAK_FLOPS,
-    MATRIX_LR,
-    SCALAR_LR,
-    UNEMBEDDING_LR,
-    WARMUP_RATIO,
-    WEIGHT_DECAY,
-    PostSemClawConfig,
-)
-from hydra.model import PostSemClawModel
-
-
-# ---------------------------------------------------------------------------
-# LR / momentum / wd schedules — verbatim copy of hydra/training.py so the
-# curves match exactly. Kept here to avoid import cycles.
-# ---------------------------------------------------------------------------
-
-
-def _lr_multiplier(progress: float) -> float:
-    if progress < WARMUP_RATIO:
-        return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
-    decay_progress = (progress - WARMUP_RATIO) / max(1.0 - WARMUP_RATIO, 1e-9)
-    return FINAL_LR_FRAC + 0.5 * (1.0 - FINAL_LR_FRAC) * (
-        1 + math.cos(math.pi * decay_progress)
-    )
-
-
-def _muon_momentum(step: int) -> float:
-    frac = min(step / 300.0, 1.0)
-    return (1 - frac) * 0.85 + frac * 0.95
-
-
-def _weight_decay(progress: float) -> float:
-    return WEIGHT_DECAY * (1 - progress)
-
-
-# ---------------------------------------------------------------------------
-
-
-class HydraLightningModule(L.LightningModule):
-    """Lightning wrapper. Public attrs: self.model, self.config."""
-
-    def __init__(self, config: PostSemClawConfig):
-        super().__init__()
-        self.config = config
-        self.model = PostSemClawModel(config)
-        # Model weights init must be deferred to the correct device; done by
-        # caller after construction (to match the meta-device + to_empty()
-        # pattern used in the legacy loop).
-
-        # Time-based progress tracks the legacy loop's semantics: LR cosine
-        # is driven by wall-clock, not step count. We capture training start
-        # in on_train_start and TIME_BUDGET from env.
-        self.time_budget = float(
-            int(os.environ.get("HYDRA_TIME_BUDGET", "300"))
-        )
-        self._train_start_time: float | None = None
-        self._total_training_time = 0.0
-        self._last_step_end: float | None = None
-        self._hestia_interval = int(os.environ.get("HYDRA_HESTIA_INTERVAL", "100"))
-        self._flops_per_token = 0
-        self._tokens_per_step = 0
-
-        # Smoothed loss for the header-line log (matches legacy format).
-        self._ema_beta = 0.9
-        self._smooth_loss = 0.0
-        self._bpt_ema = 0.0
-        self._token_bytes: torch.Tensor | None = None
-
-    # ------------------------------------------------------------------
-    # Lifecycle
-    # ------------------------------------------------------------------
-
-    def on_train_start(self) -> None:
-        self._train_start_time = time.time()
-        self._last_step_end = self._train_start_time
-        self._flops_per_token = self.model.estimate_flops()
-        # Tokens processed per optimizer step (pre-accum).
-        B = int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
-        T = int(os.environ.get("HYDRA_SEQ_LEN", "512"))
-        self._tokens_per_step = B * T
-
-        # Build/cache token_bytes LUT (for bits-per-byte live metric).
-        import prepare as _p
-        self._token_bytes = _p.get_token_bytes(device=self.device)
-
-    def configure_optimizers(self):
-        optimizer = self.model.setup_optimizer(
-            unembedding_lr=UNEMBEDDING_LR,
-            embedding_lr=EMBEDDING_LR,
-            scalar_lr=SCALAR_LR,
-            adam_betas=ADAM_BETAS,
-            matrix_lr=MATRIX_LR,
-            weight_decay=WEIGHT_DECAY,
-        )
-        return optimizer
-
-    # ------------------------------------------------------------------
-    # Training step. Lightning auto-handles: autocast (via precision flag
-    # on Trainer), backward, grad-accum, zero_grad. We only:
-    #   - split batch into (x, y)
-    #   - forward through model (autocast is established by Trainer)
-    #   - return loss (grads flow from return)
-    # ------------------------------------------------------------------
-
-    def training_step(self, batch: torch.Tensor, batch_idx: int):
-        # DataLoader produces (B, T+1) rows; split into input/target.
-        # Lightning's default collate already moved batch to self.device via
-        # the accelerator callback when pin_memory=True and device != cpu.
-        if batch.dim() != 2:
-            raise RuntimeError(f"Expected (B, T+1) batch, got shape {tuple(batch.shape)}")
-        x = batch[:, :-1].contiguous()
-        y = batch[:, 1:].contiguous()
-
-        loss = self.model(x, y)
-        # Lightning applies the grad-accum divisor automatically; we just
-        # return the raw loss. loss.detach() is stored for logging.
-        self._log_step(loss.detach(), y)
-        return loss
-
-    # ------------------------------------------------------------------
-    # Optimizer step hook: update LR / momentum / WD using time-progress.
-    # Runs once per optimizer step (after all accum micro-batches).
-    # ------------------------------------------------------------------
-
-    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
-        # Update schedules from wall-clock progress.
-        now = time.time()
-        if self._train_start_time is None:
-            self._train_start_time = now
-            self._last_step_end = now
-        progress = min(self._total_training_time / max(self.time_budget, 1.0), 1.0)
-
-        step = self.global_step
-        lrm = _lr_multiplier(progress)
-        mom = _muon_momentum(step)
-        wd = _weight_decay(progress)
-        for group in optimizer.param_groups:
-            group["lr"] = group["initial_lr"] * lrm
-            if group.get("kind") == "muon":
-                group["momentum"] = mom
-                group["weight_decay"] = wd
-
-        # Grad clip (matches legacy loop). Lightning provides this via
-        # Trainer(gradient_clip_val=1.0) but we want the exact call-site.
-        torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)
-
-        # Hyena train-cache: we must flush accumulated micro-batch grads BACK
-        # into the filter MLP params AFTER the accum-backward closure has run
-        # but BEFORE the optimizer actually consumes the grads. Lightning
-        # composes these so the closure runs inside optimizer.step(). We wrap
-        # the closure to insert our flush at the exact right moment.
-        #
-        # Ordering within the wrapped closure:
-        #   1. optimizer_closure() — runs all micro-batch forwards + backwards.
-        #      Each Hyena micro-batch backward accumulates into _k_leaf.grad.
-        #   2. flush_hyena_pending_grads() — one-shot
-        #      torch.autograd.backward(_k_graph, _k_leaf.grad) per HyenaFilter.
-        #      Now filter MLP / pos_emb / bias params have their correct grads.
-        #
-        # No-op when HYDRA_HYENA_TRAIN_CACHE=0 or no Hyena blocks exist.
-        _has_flush = hasattr(self.model, "flush_hyena_pending_grads")
-        if _has_flush:
-            _orig_closure = optimizer_closure
-
-            def _wrapped_closure():
-                result = _orig_closure()
-                self.model.flush_hyena_pending_grads()
-                return result
-
-            effective_closure = _wrapped_closure
-        else:
-            effective_closure = optimizer_closure
-
-        # Run the step (this is what Lightning would have done for us).
-        optimizer.step(closure=effective_closure)
-        self.model.zero_grad(set_to_none=True)
-
-        # Hyena filter-rfft cache invalidation. No-op if:
-        #   (a) no Hyena layers are in the model, or
-        #   (b) HYDRA_HYENA_FILTER_CACHE=0 and HYDRA_HYENA_TRAIN_CACHE=0
-        #       (the operators never populated either cache)
-        # In either case this is a handful of Python attribute resets.
-        if hasattr(self.model, "invalidate_hyena_caches"):
-            self.model.invalidate_hyena_caches()
-
-        # Hestia QAT snap every N steps. Temperature anneals every step.
-        progress_now = (now - self._train_start_time) / max(self.time_budget, 1.0)
-        self.model.hestia.anneal_temperature(progress_now)
-        if self._hestia_interval > 0 and step % self._hestia_interval == 0:
-            self.model.hestia.apply_to(self.model)
-
-        # SDR SOM update when the model stashed an sdr in the last forward.
-        _last_sdr = getattr(self.model, "_last_sdr", None)
-        if _last_sdr is not None and hasattr(self.model.sdr_semantic, "maybe_som_update"):
-            # x from the last training_step is not available here without
-            # captured state; the legacy loop passed (x, _last_sdr). To keep
-            # the interface clean we pass the last batch's x via a buffer.
-            # Since _last_sdr is derived from idx, we reuse self._last_x.
-            if getattr(self, "_last_x", None) is not None:
-                self.model.sdr_semantic.maybe_som_update(self._last_x, _last_sdr)
-
-        # Advance the wall-clock counter for LR schedule (matches legacy
-        # behavior which incremented only after the first warm-up step).
-        dt = now - (self._last_step_end or now)
-        self._last_step_end = now
-        if step > 10:
-            self._total_training_time += dt
-
-    # ------------------------------------------------------------------
-    # Logging — mirrors the step=NNNNN line format of the legacy loop so
-    # grep/tee pipelines keep working.
-    # ------------------------------------------------------------------
-
-    def _log_step(self, loss: torch.Tensor, y: torch.Tensor) -> None:
-        # Stash the current x so optimizer_step can drive SOM update.
-        self._last_x = None  # reset; we will set it below.
-        # We don't have x here (already discarded); emit a None marker that
-        # the SOM hook will silently skip if absent.
-
-        loss_f = float(loss.item())
-        if not math.isfinite(loss_f) or loss_f > 100:
-            # Let Lightning raise / the trainer callbacks handle this.
-            self.log("train_loss_nan", 1.0)
-            return
-
-        step = self.global_step
-        self._smooth_loss = (
-            self._ema_beta * self._smooth_loss + (1 - self._ema_beta) * loss_f
-        )
-        debiased = self._smooth_loss / max(1 - self._ema_beta ** (step + 1), 1e-9)
-        dt = max(time.time() - (self._last_step_end or time.time()), 1e-6)
-        tps = int(self._tokens_per_step / dt) if dt > 0 else 0
-        mfu = (
-            100.0
-            * self._flops_per_token
-            * self._tokens_per_step
-            / dt
-            / GPU_BF16_PEAK_FLOPS
-            if dt > 0
-            else 0.0
-        )
-
-        # bpb live: y flat -> token_bytes LUT -> avg bytes/token
-        bpt = debiased / math.log(2)
-        if self._token_bytes is not None:
-            with torch.no_grad():
-                y_flat = y.reshape(-1)
-                nbytes = self._token_bytes[y_flat]
-                mask = nbytes > 0
-                denom = mask.sum().clamp(min=1).float()
-                avg_bpt = (nbytes.float() * mask.float()).sum() / denom
-                bpt_batch = float(avg_bpt.item())
-            if step == 0 or self._bpt_ema <= 0.0:
-                self._bpt_ema = bpt_batch
-            else:
-                self._bpt_ema = 0.98 * self._bpt_ema + 0.02 * bpt_batch
-        bpb = bpt / max(self._bpt_ema, 1e-6)
-        vram = (
-            torch.cuda.memory_allocated() / 1024 / 1024
-            if torch.cuda.is_available()
-            else 0.0
-        )
-
-        self.log_dict(
-            {
-                "train/loss": debiased,
-                "train/bpb": bpb,
-                "train/bpt": bpt,
-                "train/tps": float(tps),
-                "train/mfu": float(mfu),
-                "train/vram_mib": float(vram),
-            },
-            prog_bar=False,
-            on_step=True,
-            on_epoch=False,
-        )
-
-        # Match legacy one-line format: "step=NNNNN loss=x bpb=y tps=z ..."
-        print(
-            f"step={step:05d} loss={debiased:.4f} bpb={bpb:.4f} "
-            f"bpt={bpt:.3f} bpt_div={self._bpt_ema:.2f} "
-            f"tps={tps} dt_ms={dt*1000:.0f} mfu={mfu:.1f} "
-            f"vram={vram:.0f}MiB",
-            flush=True,
-        )
+"""LightningModule wrapping PostSemClawModel.
+
+Thin adapter. The model and the MuonAdamW optimizer are unchanged. This
+module implements:
+
+  • configure_optimizers — returns the existing MuonAdamW (subclass of
+    torch.optim.Optimizer) built by model.setup_optimizer. Lightning accepts
+    this directly.
+  • training_step — splits (B, T+1) batches into (x, y), forwards through
+    the model, logs loss / bpb / tps / mfu / vram. Preserves the
+    sampled-softmax path inside PostSemClawModel (no changes there).
+  • optimizer_step — before each step we update LR + muon momentum + WD
+    using the same time-progress schedule as hydra/training.py
+    (get_lr_multiplier / get_muon_momentum / get_weight_decay). Lightning
+    handles grad accumulation via Trainer(accumulate_grad_batches=N).
+
+The SDR SOM update and Hestia QAT snap are called at the same cadence as
+the legacy loop, but inline on the main thread (Lightning provides its own
+callbacks for async work if we need to extract them later — keeping it
+simple for now).
+
+Env vars respected:
+  HYDRA_TIME_BUDGET          — wall-clock budget (s) used for LR schedule
+                                and as Trainer max_time
+  HYDRA_HESTIA_INTERVAL      — steps between Hestia snaps (default 100)
+  HYDRA_BATCH_SIZE           — device batch size (for throughput calc)
+  HYDRA_SEQ_LEN              — sequence length (for throughput calc)
+"""
+from __future__ import annotations
+
+import math
+import os
+import time
+
+import torch
+import lightning as L
+
+from hydra.config import (
+    ADAM_BETAS,
+    EMBEDDING_LR,
+    FINAL_LR_FRAC,
+    GPU_BF16_PEAK_FLOPS,
+    MATRIX_LR,
+    SCALAR_LR,
+    UNEMBEDDING_LR,
+    WARMUP_RATIO,
+    WEIGHT_DECAY,
+    PostSemClawConfig,
+)
+from hydra.model import PostSemClawModel
+
+
+# ---------------------------------------------------------------------------
+# LR / momentum / wd schedules — verbatim copy of hydra/training.py so the
+# curves match exactly. Kept here to avoid import cycles.
+# ---------------------------------------------------------------------------
+
+
+def _lr_multiplier(progress: float) -> float:
+    if progress < WARMUP_RATIO:
+        return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
+    decay_progress = (progress - WARMUP_RATIO) / max(1.0 - WARMUP_RATIO, 1e-9)
+    return FINAL_LR_FRAC + 0.5 * (1.0 - FINAL_LR_FRAC) * (
+        1 + math.cos(math.pi * decay_progress)
+    )
+
+
+def _muon_momentum(step: int) -> float:
+    frac = min(step / 300.0, 1.0)
+    return (1 - frac) * 0.85 + frac * 0.95
+
+
+def _weight_decay(progress: float) -> float:
+    return WEIGHT_DECAY * (1 - progress)
+
+
+# ---------------------------------------------------------------------------
+
+
+class HydraLightningModule(L.LightningModule):
+    """Lightning wrapper. Public attrs: self.model, self.config."""
+
+    def __init__(self, config: PostSemClawConfig):
+        super().__init__()
+        self.config = config
+        self.model = PostSemClawModel(config)
+        # Model weights init must be deferred to the correct device; done by
+        # caller after construction (to match the meta-device + to_empty()
+        # pattern used in the legacy loop).
+
+        # Time-based progress tracks the legacy loop's semantics: LR cosine
+        # is driven by wall-clock, not step count. We capture training start
+        # in on_train_start and TIME_BUDGET from env.
+        self.time_budget = float(
+            int(os.environ.get("HYDRA_TIME_BUDGET", "300"))
+        )
+        self._train_start_time: float | None = None
+        self._total_training_time = 0.0
+        self._last_step_end: float | None = None
+        self._hestia_interval = int(os.environ.get("HYDRA_HESTIA_INTERVAL", "100"))
+        self._flops_per_token = 0
+        self._tokens_per_step = 0
+
+        # Smoothed loss for the header-line log (matches legacy format).
+        self._ema_beta = 0.9
+        self._smooth_loss = 0.0
+        self._bpt_ema = 0.0
+        self._token_bytes: torch.Tensor | None = None
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def on_train_start(self) -> None:
+        self._train_start_time = time.time()
+        self._last_step_end = self._train_start_time
+        self._flops_per_token = self.model.estimate_flops()
+        # Tokens processed per optimizer step (pre-accum).
+        B = int(os.environ.get("HYDRA_BATCH_SIZE", "1"))
+        T = int(os.environ.get("HYDRA_SEQ_LEN", "512"))
+        self._tokens_per_step = B * T
+
+        # Build/cache token_bytes LUT (for bits-per-byte live metric).
+        import prepare as _p
+        self._token_bytes = _p.get_token_bytes(device=self.device)
+
+    def configure_optimizers(self):
+        optimizer = self.model.setup_optimizer(
+            unembedding_lr=UNEMBEDDING_LR,
+            embedding_lr=EMBEDDING_LR,
+            scalar_lr=SCALAR_LR,
+            adam_betas=ADAM_BETAS,
+            matrix_lr=MATRIX_LR,
+            weight_decay=WEIGHT_DECAY,
+        )
+        return optimizer
+
+    # ------------------------------------------------------------------
+    # Training step. Lightning auto-handles: autocast (via precision flag
+    # on Trainer), backward, grad-accum, zero_grad. We only:
+    #   - split batch into (x, y)
+    #   - forward through model (autocast is established by Trainer)
+    #   - return loss (grads flow from return)
+    # ------------------------------------------------------------------
+
+    def training_step(self, batch: torch.Tensor, batch_idx: int):
+        # DataLoader produces (B, T+1) rows; split into input/target.
+        # Lightning's default collate already moved batch to self.device via
+        # the accelerator callback when pin_memory=True and device != cpu.
+        if batch.dim() != 2:
+            raise RuntimeError(f"Expected (B, T+1) batch, got shape {tuple(batch.shape)}")
+        x = batch[:, :-1].contiguous()
+        y = batch[:, 1:].contiguous()
+
+        loss = self.model(x, y)
+        # Lightning applies the grad-accum divisor automatically; we just
+        # return the raw loss. loss.detach() is stored for logging.
+        self._log_step(loss.detach(), y)
+        return loss
+
+    # ------------------------------------------------------------------
+    # Optimizer step hook: update LR / momentum / WD using time-progress.
+    # Runs once per optimizer step (after all accum micro-batches).
+    # ------------------------------------------------------------------
+
+    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
+        # Update schedules from wall-clock progress.
+        now = time.time()
+        if self._train_start_time is None:
+            self._train_start_time = now
+            self._last_step_end = now
+        progress = min(self._total_training_time / max(self.time_budget, 1.0), 1.0)
+
+        step = self.global_step
+        lrm = _lr_multiplier(progress)
+        mom = _muon_momentum(step)
+        wd = _weight_decay(progress)
+        for group in optimizer.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+            if group.get("kind") == "muon":
+                group["momentum"] = mom
+                group["weight_decay"] = wd
+
+        # Grad clip (matches legacy loop). Lightning provides this via
+        # Trainer(gradient_clip_val=1.0) but we want the exact call-site.
+        torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)
+
+        # Hyena train-cache: we must flush accumulated micro-batch grads BACK
+        # into the filter MLP params AFTER the accum-backward closure has run
+        # but BEFORE the optimizer actually consumes the grads. Lightning
+        # composes these so the closure runs inside optimizer.step(). We wrap
+        # the closure to insert our flush at the exact right moment.
+        #
+        # Ordering within the wrapped closure:
+        #   1. optimizer_closure() — runs all micro-batch forwards + backwards.
+        #      Each Hyena micro-batch backward accumulates into _k_leaf.grad.
+        #   2. flush_hyena_pending_grads() — one-shot
+        #      torch.autograd.backward(_k_graph, _k_leaf.grad) per HyenaFilter.
+        #      Now filter MLP / pos_emb / bias params have their correct grads.
+        #
+        # No-op when HYDRA_HYENA_TRAIN_CACHE=0 or no Hyena blocks exist.
+        _has_flush = hasattr(self.model, "flush_hyena_pending_grads")
+        if _has_flush:
+            _orig_closure = optimizer_closure
+
+            def _wrapped_closure():
+                result = _orig_closure()
+                self.model.flush_hyena_pending_grads()
+                return result
+
+            effective_closure = _wrapped_closure
+        else:
+            effective_closure = optimizer_closure
+
+        # Run the step (this is what Lightning would have done for us).
+        optimizer.step(closure=effective_closure)
+        self.model.zero_grad(set_to_none=True)
+
+        # Hyena filter-rfft cache invalidation. No-op if:
+        #   (a) no Hyena layers are in the model, or
+        #   (b) HYDRA_HYENA_FILTER_CACHE=0 and HYDRA_HYENA_TRAIN_CACHE=0
+        #       (the operators never populated either cache)
+        # In either case this is a handful of Python attribute resets.
+        if hasattr(self.model, "invalidate_hyena_caches"):
+            self.model.invalidate_hyena_caches()
+
+        # Hestia QAT snap every N steps. Temperature anneals every step.
+        progress_now = (now - self._train_start_time) / max(self.time_budget, 1.0)
+        self.model.hestia.anneal_temperature(progress_now)
+        if self._hestia_interval > 0 and step % self._hestia_interval == 0:
+            self.model.hestia.apply_to(self.model)
+
+        # SDR SOM update when the model stashed an sdr in the last forward.
+        _last_sdr = getattr(self.model, "_last_sdr", None)
+        if _last_sdr is not None and hasattr(self.model.sdr_semantic, "maybe_som_update"):
+            # x from the last training_step is not available here without
+            # captured state; the legacy loop passed (x, _last_sdr). To keep
+            # the interface clean we pass the last batch's x via a buffer.
+            # Since _last_sdr is derived from idx, we reuse self._last_x.
+            if getattr(self, "_last_x", None) is not None:
+                self.model.sdr_semantic.maybe_som_update(self._last_x, _last_sdr)
+
+        # Advance the wall-clock counter for LR schedule (matches legacy
+        # behavior which incremented only after the first warm-up step).
+        dt = now - (self._last_step_end or now)
+        self._last_step_end = now
+        if step > 10:
+            self._total_training_time += dt
+
+    # ------------------------------------------------------------------
+    # Logging — mirrors the step=NNNNN line format of the legacy loop so
+    # grep/tee pipelines keep working.
+    # ------------------------------------------------------------------
+
+    def _log_step(self, loss: torch.Tensor, y: torch.Tensor) -> None:
+        # Stash the current x so optimizer_step can drive SOM update.
+        self._last_x = None  # reset; we will set it below.
+        # We don't have x here (already discarded); emit a None marker that
+        # the SOM hook will silently skip if absent.
+
+        loss_f = float(loss.item())
+        if not math.isfinite(loss_f) or loss_f > 100:
+            # Let Lightning raise / the trainer callbacks handle this.
+            self.log("train_loss_nan", 1.0)
+            return
+
+        step = self.global_step
+        self._smooth_loss = (
+            self._ema_beta * self._smooth_loss + (1 - self._ema_beta) * loss_f
+        )
+        debiased = self._smooth_loss / max(1 - self._ema_beta ** (step + 1), 1e-9)
+        dt = max(time.time() - (self._last_step_end or time.time()), 1e-6)
+        tps = int(self._tokens_per_step / dt) if dt > 0 else 0
+        mfu = (
+            100.0
+            * self._flops_per_token
+            * self._tokens_per_step
+            / dt
+            / GPU_BF16_PEAK_FLOPS
+            if dt > 0
+            else 0.0
+        )
+
+        # bpb live: y flat -> token_bytes LUT -> avg bytes/token
+        bpt = debiased / math.log(2)
+        if self._token_bytes is not None:
+            with torch.no_grad():
+                y_flat = y.reshape(-1)
+                nbytes = self._token_bytes[y_flat]
+                mask = nbytes > 0
+                denom = mask.sum().clamp(min=1).float()
+                avg_bpt = (nbytes.float() * mask.float()).sum() / denom
+                bpt_batch = float(avg_bpt.item())
+            if step == 0 or self._bpt_ema <= 0.0:
+                self._bpt_ema = bpt_batch
+            else:
+                self._bpt_ema = 0.98 * self._bpt_ema + 0.02 * bpt_batch
+        bpb = bpt / max(self._bpt_ema, 1e-6)
+        vram = (
+            torch.cuda.memory_allocated() / 1024 / 1024
+            if torch.cuda.is_available()
+            else 0.0
+        )
+
+        self.log_dict(
+            {
+                "train/loss": debiased,
+                "train/bpb": bpb,
+                "train/bpt": bpt,
+                "train/tps": float(tps),
+                "train/mfu": float(mfu),
+                "train/vram_mib": float(vram),
+            },
+            prog_bar=False,
+            on_step=True,
+            on_epoch=False,
+        )
+
+        # Match legacy one-line format: "step=NNNNN loss=x bpb=y tps=z ..."
+        print(
+            f"step={step:05d} loss={debiased:.4f} bpb={bpb:.4f} "
+            f"bpt={bpt:.3f} bpt_div={self._bpt_ema:.2f} "
+            f"tps={tps} dt_ms={dt*1000:.0f} mfu={mfu:.1f} "
+            f"vram={vram:.0f}MiB",
+            flush=True,
+        )

overlay/hydra/model.py

@@ -1,894 +1,894 @@
-"""PostSemClawModel — full-architecture model assembly.
-
-Extracted from the monolithic train.py (W1 modularization). Semantics
-unchanged. Imports `GPUEngram` from `hydra.engram` and `MuonAdamW` from
-`hydra.optimizer`.
-
-Triton kernel integration status (Phase 2):
-  HYDRA_FUSED_BCNORM — DEFERRED. The bcnorm_fused Triton kernel fuses
-    LayerNorm + RoPE on B/C projections. However, mamba-ssm's Mamba3 block
-    uses RMSNormGated (not LayerNorm) for B/C, and RoPE is applied inside
-    the mamba3_siso_combined CUDA kernel via the Angles parameter. Replacing
-    would require either (a) monkey-patching RMSNormGated + intercepting the
-    fused CUDA scan — invasive, 50+ lines, high breakage risk — or (b) a
-    full custom Mamba3Block reimplementation. Both are out of scope for
-    Phase 2. The kernel is validated standalone; integration deferred to
-    Phase 3 when HYDRA moves to a custom SSM block.
-
-  HYDRA_FUSED_SSD — DEFERRED. The ssd_exp_trap Triton kernel implements
-    exponential-trapezoidal discretization as a sequential scan. mamba-ssm's
-    Mamba3 block delegates the entire scan + gating + output projection to
-    mamba3_siso_combined (a compiled CUDA kernel with tilelang). Replacing
-    it would require decomposing the combined kernel into constituent ops
-    and substituting only the scan — not feasible without a custom block.
-    Same Phase 3 gate as above.
-
-Both env vars are accepted but currently no-ops (gates read, logged, but
-the code path is unchanged). This avoids silent regression if someone
-sets them expecting a speedup.
-"""
-
-from __future__ import annotations
-
-import os
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from mamba_ssm import Mamba3
-
-from subsystems.hestia_mini import HestiaQAT
-from subsystems.htm import HTMLayer
-from subsystems.mhc_mini import ManifoldHyperConnection
-from subsystems.sdr_semantic import SemanticFoldingSDR
-
-from hydra.engram import GPUEngram
-from hydra.hyena_block import HyenaBlock
-# GDNBlock is imported lazily inside __init__ so the `fla` dependency is
-# only required when HYDRA_GDN_LAYERS is actually non-empty. Baseline
-# pure-Mamba3 runs continue to work without flash-linear-attention installed.
-from hydra.optimizer import MuonAdamW
-
-
-def norm(x: torch.Tensor) -> torch.Tensor:
-    """RMSNorm over the last dim — stateless, autocast-friendly."""
-    return F.rms_norm(x, (x.size(-1),))
-
-
-class PostSemClawModel(nn.Module):
-    """Full Post-SEM-Claw model assembly.
-
-    Architecture:
-        Token Embedding -> [Mamba3 + residual] x n_layer
-        -> SDR + Engram (at configured layer) -> norm -> LM head
-
-    Interface (must match prepare.py evaluate_bpb):
-        model(x, y, reduction='none').view(-1)  -> per-token losses
-        model(x, y, reduction='mean')           -> scalar loss
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-
-        # Token embedding
-        self.wte = nn.Embedding(config.vocab_size, config.d_model)
-
-        # Mamba-3 blocks — official mamba-ssm fused CUDA kernel. No fallbacks.
-        # RoPE is applied internally by the Mamba3 CUDA kernel via the Angles
-        # parameter; external cos/sin buffers are not needed.
-        #
-        # Hyena supplement: layers whose index appears in `config.hyena_layers`
-        # are instantiated as HyenaBlock instead of Mamba3. The config field
-        # is populated from HYDRA_HYENA_LAYERS at construction time and then
-        # persisted to checkpoints, so resume is safe even when the env var
-        # is unset. Empty tuple → all-Mamba3, byte-identical to pre-port.
-        _hyena_layer_set = set(getattr(config, "hyena_layers", ()) or ())
-        _gdn_layer_set = set(getattr(config, "gdn_layers", ()) or ())
-        # Hyena wins on overlap; conflict is logged at construction time.
-        _both = _hyena_layer_set & _gdn_layer_set
-        if _both:
-            print(f"[WARN] layers in both hyena_layers and gdn_layers; using Hyena: {sorted(_both)}", flush=True)
-            _gdn_layer_set -= _hyena_layer_set
-
-        if _gdn_layer_set:
-            from hydra.gdn_block import GDNBlock  # requires `fla` package
-
-        def _build_block(i: int) -> nn.Module:
-            if i in _hyena_layer_set:
-                return HyenaBlock(
-                    d_model=config.d_model,
-                    seq_len=config.sequence_len,
-                    order=int(os.environ.get("HYDRA_HYENA_ORDER", "2")),
-                    filter_order=int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64")),
-                )
-            if i in _gdn_layer_set:
-                return GDNBlock(
-                    d_model=config.d_model,
-                    n_heads=config.n_heads,
-                )
-            return Mamba3(
-                d_model=config.d_model,
-                d_state=config.d_state,
-                expand=config.expand,
-                headdim=config.headdim,
-                is_mimo=False,          # SISO path uses stable mamba3_siso_combined kernel
-                chunk_size=64,          # upstream-recommended SISO chunk; 16 violated tl.dot M>=16 constraint
-                is_outproj_norm=False,
-                dtype=torch.bfloat16,
-            )
-
-        self.blocks = nn.ModuleList([_build_block(i) for i in range(config.n_layer)])
-
-        # Full-architecture SDR: offline semantic retina + STE (no-bypass).
-        self.sdr_semantic = SemanticFoldingSDR(
-            vocab_size=config.vocab_size,
-            n_bits=config.sdr_n_bits,
-            target_active=config.sdr_target_active,
-            delta_rank=config.sdr_delta_rank,
-            som_warmup_steps=config.sdr_som_warmup,
-            som_update_interval=config.sdr_som_interval,
-        )
-
-        # HTM spatial pooler + temporal memory (Rust, Hebbian).
-        self.htm = HTMLayer(
-            input_bits=config.sdr_n_bits,
-            n_columns=config.htm_n_columns,
-            cells_per_column=config.htm_cells_per_column,
-            batch_size=1,           # grows lazily to actual B on first forward
-            seed=42,
-            learn=True,
-            reset_each_forward=True,
-        )
-
-        # Gradient bridge: (n_columns + anomaly) -> d_model.
-        self.htm_proj = nn.Linear(config.htm_n_columns + 1, config.d_model, bias=False)
-
-        # GPU Engram with Hebbian writes — runs EVERY step.
-        self.engram = GPUEngram(
-            d_model=config.d_model,
-            n_columns=config.engram_n_columns,
-            max_ngram=3,
-        )
-        self.engram_layer_idx = config.engram_layer_idx
-
-        # Manifold-Constrained Hyper-Connections (one per Mamba-3 block).
-        self.mhc = nn.ModuleList([
-            ManifoldHyperConnection(d_model=config.d_model, n_streams=2, sinkhorn_iters=3)
-            for _ in range(config.n_layer)
-        ])
-
-        # Hestia QAT — ternary weight quantization applied post-optimizer-step.
-        self.hestia = HestiaQAT(enabled=True, bits=1.58)
-
-        # LM head
-        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
-
-        # Learnability knob 1: Multi-Token Prediction (Llama-3 style).
-        #   MTP_K=1 -> standard next-token. MTP_K>1 -> extra heads predict
-        #   tokens at positions t+1, t+2, ..., t+K. Heads are weight-tied to
-        #   lm_head (we share Parameters), so the only extra compute is
-        #   additional CE losses; no new params. Activated via HYDRA_MTP_K.
-        self._mtp_k = max(1, int(os.environ.get("HYDRA_MTP_K", "1")))
-
-        # Learnability knob 3: gradient checkpointing on Mamba3 blocks.
-        self._grad_ckpt = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
-
-        # Learnability knob 4: doc-separator BOS masking in packed sequences.
-        self._doc_sep_mask = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
-        # BOS token id is looked up lazily on first forward (requires tokenizer
-        # load); -1 means uninitialized.
-        self._bos_token_id = -1
-
-        # Learnability knob 5: explicit stop-grad on HTM tensor (htm_rust
-        # outputs already have requires_grad=False; this is defense-in-depth).
-        self._htm_stop_grad = os.environ.get("HYDRA_HTM_STOP_GRAD", "0") == "1"
-
-        # Learnability knob 6: entropy penalty coefficient on LM logits.
-        self._entropy_penalty = float(os.environ.get("HYDRA_ENTROPY_PENALTY", "0.0"))
-
-        # Residual dropout
-        self.drop = nn.Dropout(float(os.environ.get("HYDRA_DROPOUT", "0.2")))
-
-        # Logits soft-capping
-        self.softcap = 15.0
-
-        # Secondary metrics storage
-        self._metrics = {}
-
-        # Per-layer diagnostic panel. Env-gated; zero overhead when off.
-        # Emits residual-contribution (delta_ratio), feature std, effective rank,
-        # gradient norm per layer; used to identify minimum viable n_layer + find
-        # entropy leakage / dead layers. See docs/depth-sweep.md.
-        self._diag_enabled = os.environ.get("HYDRA_LAYER_DIAGNOSTICS", "0") == "1"
-        self._diag_step = 0
-        self._diag_svd_every = int(os.environ.get("HYDRA_LAYER_DIAG_SVD_EVERY", "100"))
-        if self._diag_enabled:
-            # Gradient-norm backward hooks on each Mamba3 block output.
-            for _i, _block in enumerate(self.blocks):
-                def _mk_grad_hook(_layer_idx):
-                    def _hook(module, grad_input, grad_output):
-                        if grad_output and grad_output[0] is not None:
-                            g = grad_output[0].detach()
-                            self._metrics[f'layer_{_layer_idx}_grad_norm'] = float(
-                                g.pow(2).mean().sqrt().item()
-                            )
-                    return _hook
-                _block.register_full_backward_hook(_mk_grad_hook(_i))
-
-            # Forward hooks on each Mamba3 block capture the block's OUTPUT
-            # directly. This is the clean measurement: unlike merge_streams()
-            # sampling which sees (streams + M*block_output) in bf16 — where
-            # small block contributions round to zero against unit-norm
-            # residuals — this captures `block_output` itself as produced.
-            # Reports both its absolute RMS norm and its ratio to the block
-            # INPUT's RMS norm (contribution magnitude relative to the
-            # residual it's added to).
-            for _i, _block in enumerate(self.blocks):
-                def _mk_fwd_hook(_layer_idx):
-                    def _hook(module, inputs, output):
-                        with torch.no_grad():
-                            inp = inputs[0].detach().float() if inputs else None
-                            out = output.detach().float() if isinstance(output, torch.Tensor) else None
-                            if out is not None:
-                                out_rms = out.pow(2).mean().sqrt().item()
-                                self._metrics[f'layer_{_layer_idx}_block_out_rms'] = float(out_rms)
-                                if inp is not None:
-                                    in_rms = inp.pow(2).mean().sqrt().item()
-                                    self._metrics[f'layer_{_layer_idx}_block_in_rms'] = float(in_rms)
-                                    self._metrics[f'layer_{_layer_idx}_contrib_ratio'] = float(
-                                        out_rms / (in_rms + 1e-8)
-                                    )
-                    return _hook
-                _block.register_forward_hook(_mk_fwd_hook(_i))
-
-        # Triton kernel integration gates (Phase 2 — deferred, see module docstring).
-        self._fused_bcnorm = os.environ.get("HYDRA_FUSED_BCNORM", "0") == "1"
-        self._fused_ssd = os.environ.get("HYDRA_FUSED_SSD", "0") == "1"
-        if self._fused_bcnorm or self._fused_ssd:
-            import sys
-            _active = []
-            if self._fused_bcnorm:
-                _active.append("HYDRA_FUSED_BCNORM")
-            if self._fused_ssd:
-                _active.append("HYDRA_FUSED_SSD")
-            print(
-                f"[HYDRA] Triton kernel gates set: {', '.join(_active)}. "
-                f"NOTE: Both are DEFERRED (mamba-ssm Mamba3 uses internal "
-                f"CUDA kernels). Gates accepted but currently no-ops.",
-                file=sys.stderr,
-            )
-
-        # R6 optional torch.compile on the impl forward. Gated (default OFF).
-        if os.environ.get("HYDRA_MODEL_COMPILE", "0") == "1":
-            self._forward_impl = torch.compile(
-                self._forward_impl,
-                fullgraph=False,
-                dynamic=True,
-                mode="default",
-            )
-
-    @torch.no_grad()
-    def init_weights(self) -> None:
-        s = 3 ** 0.5 * self.config.d_model ** -0.5
-
-        # Move SDR retina indices (plain attribute, not buffer) to same device as params.
-        # Required because to_empty() only moves params/buffers, and _retina_indices
-        # is loaded from numpy (always CPU) by SemanticFoldingSDR.__init__.
-        device = self.wte.weight.device
-        if hasattr(self.sdr_semantic, '_retina_indices'):
-            self.sdr_semantic._retina_indices = self.sdr_semantic._retina_indices.to(device)
-
-        # Embedding init: GPT-2 / LLaMA convention. std=1.0 was chosen for
-        # vocab=8192; at larger vocabs, smaller std prevents logit blowup.
-        # Use std = 1/sqrt(d_model) which scales sensibly with model width.
-        import math as _math
-        _d_model = self.wte.weight.shape[1]
-        wte_std = float(os.environ.get("HYDRA_WTE_STD", str(1.0 / _math.sqrt(_d_model))))
-        nn.init.normal_(self.wte.weight, mean=0.0, std=wte_std)
-        # LM head init: was std=0.001 — PATHOLOGICAL at vocab>=32k because
-        # logits collapse to zero, loss locks at log(V)~=11, gradient through
-        # head ∝ 1/V is too small to escape. GPT-2 uses std=0.02; LLaMA uses
-        # std=1/sqrt(d_model). Pick 0.02 as robust default, env-overridable.
-        lm_head_std = float(os.environ.get("HYDRA_LM_HEAD_STD", "0.02"))
-        nn.init.normal_(self.lm_head.weight, mean=0.0, std=lm_head_std)
-        # F8 (NOT APPLIED): Weight tying would save V*D params but current LR
-        # groups have embedding_lr=1.0 and unembedding_lr=0.005 × d_model_scale
-        # — tying forces the shared tensor under a single LR group and either
-        # the embeddings learn 200x too slow (under unembed LR) or the LM head
-        # becomes unstable (under embed LR). Short 15-step smoke with tying +
-        # embed-group update showed initial loss jump 9 -> 20. Deferred until
-        # LR groups are re-tuned; see docs/OPTIMIZATION_PLAN.md Post-plan.
-
-        for li, block in enumerate(self.blocks):
-            if hasattr(block, 'in_proj') and hasattr(block.in_proj, 'weight'):
-                nn.init.uniform_(block.in_proj.weight, -s, s)
-            if hasattr(block, 'out_proj') and hasattr(block.out_proj, 'weight'):
-                # GPT-2 residual init: std = 0.02 / sqrt(2 * n_layer).
-                # NOT zeros — zero init makes the block a permanent pass-through
-                # (block_out_rms=0, zero gradient flow to SSM internals).
-                # With non-zero init the block contributes to the residual stream
-                # from step 1, giving the SSM scan actual gradient signal.
-                n_layer = self.config.n_layer
-                out_std = float(os.environ.get(
-                    "HYDRA_OUT_PROJ_STD",
-                    str(0.02 / (2 * n_layer) ** 0.5),
-                ))
-                nn.init.normal_(block.out_proj.weight, mean=0.0, std=out_std)
-
-        nn.init.normal_(self.htm_proj.weight, mean=0.0, std=s)
-
-        # Cast to bf16 to match Mamba3 dtype; Muon groups by shape so mixed
-        # dtypes in the same shape group would break lerp_ dtype checks.
-        self.wte.to(dtype=torch.bfloat16)
-        self.htm_proj.to(dtype=torch.bfloat16)
-        self.engram.to(dtype=torch.bfloat16)
-
-    def set_bos_token_id(self, bos_id: int) -> None:
-        """Inform the model of the tokenizer's BOS id so doc-separator
-        masking (learnability #4) knows which positions to skip. Called from
-        training setup once the tokenizer is loaded."""
-        self._bos_token_id = int(bos_id)
-
-    def invalidate_hyena_caches(self) -> None:
-        """Invalidate filter-rfft caches on all Hyena blocks.
-
-        MUST be called after each `optimizer.step()` when
-        `HYDRA_HYENA_FILTER_CACHE=1` is set, otherwise cached rfft values
-        will be reused with stale filter parameters.
-
-        No-op for blocks that are not HyenaBlock (Mamba3, etc.).
-        """
-        for block in self.blocks:
-            if hasattr(block, "operator") and hasattr(block.operator, "invalidate_filter_cache"):
-                block.operator.invalidate_filter_cache()
-
-    def flush_hyena_pending_grads(self) -> None:
-        """Push pending train-cache filter gradients into filter params.
-
-        Used ONLY when HYDRA_HYENA_TRAIN_CACHE=1. Must be called exactly once
-        per optimizer step, BEFORE `optimizer.step()` and BEFORE
-        `invalidate_hyena_caches()`. The lightning_module wires this in
-        `optimizer_step` around the existing optimizer.step() call.
-
-        No-op if:
-          * No HyenaBlocks are in the model, OR
-          * No micro-batch ever ran with grad enabled (e.g. all-eval step).
-        """
-        for block in self.blocks:
-            if hasattr(block, "operator") and hasattr(block.operator, "flush_pending_filter_grads"):
-                block.operator.flush_pending_filter_grads()
-
-    def estimate_flops(self) -> int:
-        nparams = sum(p.numel() for p in self.parameters())
-        embed_params = self.wte.weight.numel()
-        return 6 * (nparams - embed_params)
-
-    def num_scaling_params(self) -> dict:
-        wte = sum(p.numel() for p in self.wte.parameters())
-        lm_head = sum(p.numel() for p in self.lm_head.parameters())
-        blocks = sum(p.numel() for p in self.blocks.parameters())
-        sdr = sum(p.numel() for p in self.sdr_semantic.parameters())
-        htm_proj = sum(p.numel() for p in self.htm_proj.parameters())
-        engram = sum(p.numel() for p in self.engram.parameters())
-        total = sum(p.numel() for p in self.parameters())
-        return {
-            'wte': wte, 'lm_head': lm_head, 'blocks': blocks,
-            'sdr_semantic': sdr, 'htm_proj': htm_proj,
-            'engram': engram, 'total': total,
-        }
-
-    def get_secondary_metrics(self) -> dict:
-        """Flush any lingering CUDA tensors to host (single sync)."""
-        flushed = {}
-        for k, v in self._metrics.items():
-            if hasattr(v, 'item'):
-                try:
-                    flushed[k] = float(v.item())
-                except Exception:
-                    flushed[k] = v
-            else:
-                flushed[k] = v
-        return flushed
-
-    def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.6, matrix_lr=0.04,
-                        weight_decay=0.2, adam_betas=(0.8, 0.95), scalar_lr=0.5):
-        """Setup MuonAdamW optimizer with per-component LR groups."""
-        model_dim = self.config.d_model
-
-        embedding_params = list(self.wte.parameters())
-        lm_head_params = list(self.lm_head.parameters())
-
-        # Muon routing guard: 2D parameters are NOT automatically matrices.
-        # Exclude:
-        #   (a) params whose name ends in `.freq` — Sin frequency vectors used
-        #       by Hyena's implicit filter MLP. Shape (1, dim) is nominally 2D
-        #       but semantically a per-dim scalar. Muon's polar-express
-        #       orthogonalization would force it toward an orthogonal matrix,
-        #       destroying the learned modulation frequencies.
-        #   (b) 2-D params with min(shape) < MUON_MIN_DIM. Tiny projections
-        #       (e.g. HyenaFilter.implicit_filter.0.weight of shape (64, 3))
-        #       get collapsed toward near-identity by orthogonalization on the
-        #       narrow axis, damaging expressivity. These belong in AdamW.
-        # These exclusions route the params into the AdamW scalar/vector group.
-        MUON_MIN_DIM = 8
-
-        def _muon_eligible(name: str, p: torch.Tensor) -> bool:
-            if p.dim() != 2:
-                return False
-            if name.endswith(".freq"):
-                return False
-            if min(p.shape) < MUON_MIN_DIM:
-                return False
-            return True
-
-        # Matrix params -> Muon (2D weight matrices passing the routing guard).
-        matrix_params = []
-        for name, p in self.blocks.named_parameters():
-            if _muon_eligible(name, p):
-                matrix_params.append(p)
-        # NOTE (W1 audit REG-2): SemanticFoldingSDR.delta_u / delta_v are
-        # currently GRADIENT-DEAD. The forward path uses `binary_only(idx)` for
-        # HTM and stores it as `self._last_sdr`, but does NOT route the STE
-        # output through any downstream op. Including them in the Muon group
-        # burns compute (stack + orthogonalize + lerp) on zero-grad params
-        # every step. Excluded here; a later W5 pass can reconnect STE via a
-        # gated residual if the SDR signal is wanted back in-graph. The
-        # parameters still exist, so no state_dict break.
-        # for p in self.sdr_semantic.parameters():
-        #     if p.dim() == 2:
-        #         matrix_params.append(p)
-        for name, p in self.htm_proj.named_parameters():
-            if _muon_eligible(name, p):
-                matrix_params.append(p)
-        for name, p in self.engram.named_parameters():
-            if _muon_eligible(name, p):
-                matrix_params.append(p)
-
-        # SDR params are intentionally not in any optimizer group — they
-        # receive no gradient in the current forward, so any update would be
-        # pure noise (weight_decay × lr on a zero-grad param).
-        sdr_param_ids = set(id(p) for p in self.sdr_semantic.parameters())
-        assigned = set(id(p) for p in embedding_params + lm_head_params + matrix_params)
-        scalar_params = [
-            p for p in self.parameters()
-            if id(p) not in assigned and id(p) not in sdr_param_ids
-        ]
-
-        total_assigned = len(embedding_params) + len(lm_head_params) + len(matrix_params) + len(scalar_params)
-        total_params = len(list(self.parameters()))
-        sdr_excluded = len(list(self.sdr_semantic.parameters()))
-        assert total_assigned + sdr_excluded == total_params, (
-            f"Parameter count mismatch: assigned {total_assigned} + sdr_excluded "
-            f"{sdr_excluded} vs total {total_params}"
-        )
-
-        dmodel_lr_scale = (model_dim / 768) ** -0.5
-        print(f"Scaling AdamW LRs by 1/sqrt({model_dim}/768) = {dmodel_lr_scale:.6f}")
-
-        param_groups = [
-            dict(kind='adamw', params=lm_head_params,
-                 lr=unembedding_lr * dmodel_lr_scale, betas=adam_betas,
-                 eps=1e-10, weight_decay=0.0),
-            dict(kind='adamw', params=embedding_params,
-                 lr=embedding_lr * dmodel_lr_scale, betas=adam_betas,
-                 eps=1e-10, weight_decay=0.0),
-        ]
-
-        if scalar_params:
-            param_groups.append(
-                dict(kind='adamw', params=scalar_params,
-                     lr=scalar_lr * dmodel_lr_scale, betas=adam_betas,
-                     eps=1e-10, weight_decay=0.0)
-            )
-
-        for shape in sorted({p.shape for p in matrix_params}):
-            group_params = [p for p in matrix_params if p.shape == shape]
-            param_groups.append(dict(
-                kind='muon', params=group_params, lr=matrix_lr,
-                momentum=0.95, ns_steps=5, beta2=0.95, weight_decay=weight_decay,
-            ))
-
-        optimizer = MuonAdamW(param_groups)
-        for group in optimizer.param_groups:
-            group["initial_lr"] = group["lr"]
-        return optimizer
-
-    def forward(self, idx, targets=None, reduction='mean'):
-        """idx: (B, T) int64. Returns loss if targets given, else logits.
-
-        Nested bf16 autocast is a no-op when ambient autocast is already on;
-        when it's off (e.g. integration tests) we establish the dtype contract.
-        """
-        if torch.is_autocast_enabled():
-            return self._forward_impl(idx, targets=targets, reduction=reduction)
-        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-            return self._forward_impl(idx, targets=targets, reduction=reduction)
-
-    def _forward_impl(self, idx, targets=None, reduction='mean'):
-        B, T = idx.shape
-
-        # Diagnostic: per-subsystem CUDA event timing. Env-gated; zero overhead
-        # when disabled. Logs one timing line per forward call. Used to isolate
-        # which subsystem is the tps bottleneck on paid hardware.
-        _profile = os.environ.get("HYDRA_PROFILE_FORWARD", "0") == "1"
-        if _profile:
-            def _ev():
-                e = torch.cuda.Event(enable_timing=True)
-                e.record()
-                return e
-            _t0 = _ev()
-        else:
-            _t0 = None
-
-        # Compute SDR binary ONCE and reuse for both HTM input and the stash.
-        sdr_binary = self.sdr_semantic.binary_only(idx)
-        self._last_sdr = sdr_binary  # uint8 stash (not bf16 → 256MB avoidance)
-
-        # HTM subsampling: run HTM on 1 of every N micro-batches within a
-        # gradient accumulation step, reuse the cached result for the other
-        # N-1 micro-batches. Cooperative launch monopolizes all SMs (grid.sync
-        # requires full-grid residency), so HTM and mamba can't overlap via
-        # streams. Subsampling removes HTM from most micro-batches' critical
-        # path instead.
-        #
-        # Math: N=8, 64 accum steps → 8 HTM calls (10.6ms each) + 56 fast
-        # calls (4ms each). Total = 84.8 + 224 = 309ms → 106k tps.
-        #
-        # HYDRA_HTM_SUBSAMPLE=N (default 8). Set =1 for every-microbatch HTM.
-        _htm_sub = int(os.environ.get("HYDRA_HTM_SUBSAMPLE", "8"))
-        if not hasattr(self, '_htm_call_idx'):
-            self._htm_call_idx = 0
-
-        _run_htm = (self._htm_call_idx % _htm_sub == 0)
-        self._htm_call_idx += 1
-
-        if _run_htm:
-            htm_handle = self.htm.forward_async(sdr_binary)
-        else:
-            htm_handle = None
-
-        if _profile: _t_htm_async = _ev()
-
-        dense_emb = self.wte(idx)  # (B, T, d_model) bf16
-
-        if _profile: _t_wte = _ev()
-
-        if _run_htm:
-            htm_out = self.htm.forward_await(htm_handle)
-            self._htm_cache = htm_out.detach()  # cache for non-HTM micro-batches
-        elif hasattr(self, '_htm_cache') and self._htm_cache is not None \
-                and self._htm_cache.shape[0] == B and self._htm_cache.shape[1] == T:
-            htm_out = self._htm_cache
-        else:
-            # Very first call with subsample > 1: run HTM anyway.
-            htm_handle = self.htm.forward_async(sdr_binary)
-            htm_out = self.htm.forward_await(htm_handle)
-            self._htm_cache = htm_out.detach()
-
-        if _profile: _t_htm_await = _ev()
-        with torch.no_grad():
-            sdr_active_bits = float(self.sdr_semantic.target_active)
-            htm_anomaly = htm_out[..., -1].mean()
-
-        # Learnability #5: explicit stop-grad on HTM output. htm_rust already
-        # produces a detached tensor, but making it explicit here hardens the
-        # contract against future refactors that might route HTM through a
-        # grad-enabled op.
-        if self._htm_stop_grad:
-            htm_out = htm_out.detach()
-
-        # Gradient bridge: HTM columns+anomaly -> d_model.
-        htm_proj_out = self.htm_proj(htm_out.to(dense_emb.dtype))
-        x = dense_emb + htm_proj_out
-        x = norm(x)
-
-        if _profile: _t_htm_proj = _ev()
-
-        # mHC-routed Mamba-3 stack with Engram injection at configured layer.
-        streams = self.mhc[0].init_streams(x)
-        _engram_ev = None
-
-        # Per-layer diagnostic panel. The pre-layer merged state h_pre lets us
-        # measure residual contribution of each layer: delta_N = h_post - h_pre.
-        # All reads are detached no-grad to avoid autograd graph pollution.
-        _diag = self._diag_enabled
-        if _diag:
-            # Cast to float32 for the diagnostic arithmetic: the layer's
-            # residual contribution is small (~0.5 × rms-normed block output),
-            # which underflows in bf16 subtraction (3-digit mantissa) and
-            # reports delta_ratio=0 at the boundaries. float32 snapshot is
-            # ~3.8 MB extra memory per diag sample (B=1, T=2048, d=96) —
-            # negligible vs peak VRAM.
-            with torch.no_grad():
-                h_pre = self.mhc[0].merge_streams(streams).detach().float()
-            _run_svd = (self._diag_step % self._diag_svd_every) == 0
-
-        for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
-            def _block_fn(h, _block=block):
-                return self.drop(_block(norm(h)))
-
-            # Learnability #3: gradient checkpointing. Wrap the block-fn so
-            # the mhc layer's internal uses of it re-run the block in backward
-            # (trading compute for activation memory). use_reentrant=False is
-            # the modern API and works cleanly under autocast.
-            if self._grad_ckpt and self.training:
-                import torch.utils.checkpoint as _ckpt
-                _raw_fn = _block_fn
-                def _block_fn(h, _raw=_raw_fn):  # noqa: E731
-                    return _ckpt.checkpoint(_raw, h, use_reentrant=False)
-
-            streams = mhc_layer(streams, _block_fn)
-
-            if i == self.engram_layer_idx:
-                if _profile: _t_pre_engram = _ev()
-                x_mid = mhc_layer.merge_streams(streams)
-                x_mid, hit_rate = self.engram(x_mid, idx)
-                streams = mhc_layer.init_streams(x_mid)
-                self._metrics['engram_hit_rate'] = hit_rate
-                if _profile: _engram_ev = _ev()
-
-            if _diag:
-                with torch.no_grad():
-                    h_post = mhc_layer.merge_streams(streams).detach().float()
-                    in_n  = h_pre.pow(2).mean().sqrt()
-                    out_n = h_post.pow(2).mean().sqrt()
-                    d_n   = (h_post - h_pre).pow(2).mean().sqrt()
-                    self._metrics[f'layer_{i}_in_norm']     = float(in_n.item())
-                    self._metrics[f'layer_{i}_out_norm']    = float(out_n.item())
-                    self._metrics[f'layer_{i}_delta_ratio'] = float((d_n / (in_n + 1e-6)).item())
-                    self._metrics[f'layer_{i}_feat_std']    = float(h_post.std(dim=-1).mean().item())
-                    if _run_svd:
-                        # Effective rank via participation ratio of singular values.
-                        # eff_rank = (Σσ)^2 / Σσ² — smooth rank proxy, bounded by d_model.
-                        # Sampled to keep overhead low (SVD is O(min(B*T, D)^2·D)).
-                        flat = h_post.reshape(-1, h_post.shape[-1])[:512].float()
-                        try:
-                            s = torch.linalg.svdvals(flat)
-                            eff_rank = float(((s.sum() ** 2) / (s.pow(2).sum() + 1e-6)).item())
-                            self._metrics[f'layer_{i}_eff_rank'] = eff_rank
-                        except Exception:
-                            pass
-                    h_pre = h_post
-
-        if _diag:
-            self._diag_step += 1
-
-        if _profile: _t_blocks = _ev()
-
-        self._metrics['sdr_active_bits'] = sdr_active_bits
-        self._metrics['htm_anomaly'] = htm_anomaly
-
-        x = self.mhc[-1].merge_streams(streams)
-        x = norm(x)
-
-        if _profile: _t_merge = _ev()
-
-        softcap = self.softcap
-        _softcap_clamp = os.environ.get("HYDRA_SOFTCAP_CLAMP", "0") == "1"
-        if targets is not None:
-            smoothing = self.config.label_smoothing
-            V = self.config.vocab_size
-
-            # Learnability #4: doc-separator masking. In packed rows,
-            # tokenizer.encode(..., prepend=bos_token) places a BOS at every
-            # document boundary. Without masking, the model is penalized for
-            # failing to predict "doc B's BOS" from the last tokens of doc A
-            # — pure noise. We set targets==bos to -1 (ignore_index). Done
-            # BEFORE MTP/entropy/sampled-softmax branches so all downstream
-            # losses inherit the mask.
-            if self._doc_sep_mask and self._bos_token_id >= 0:
-                targets = torch.where(
-                    targets == self._bos_token_id,
-                    torch.full_like(targets, -1),
-                    targets,
-                )
-
-            # Sampled softmax: instead of computing logits for ALL V tokens,
-            # compute only for the target + K random negatives. Reduces the
-            # lm_head matmul from (B*T, d) × (d, V) to (B*T, d) × (d, K+1).
-            # At V=65536 and K=4096: 16× less compute, ~4× tps improvement.
-            # The log-sum-exp correction adjusts for the sampling bias.
-            # Set HYDRA_SAMPLED_SOFTMAX=0 to disable (full softmax).
-            K_neg = int(os.environ.get("HYDRA_SAMPLED_SOFTMAX", "4096"))
-            use_sampled = K_neg > 0 and K_neg < V and self.training
-
-            if use_sampled:
-                # Flatten hidden states + targets
-                h_flat = x.reshape(-1, x.shape[-1])            # (B*T, d)
-                t_flat = targets.reshape(-1)                    # (B*T,)
-                n = h_flat.shape[0]
-
-                # Learnability #4 hardening: sampled-softmax gather crashes on
-                # negative ids (-1 from doc-sep mask). Replace -1 with 0 for
-                # gather; the actual loss is masked below.
-                valid_mask_flat = (t_flat >= 0)
-                t_flat_safe = torch.where(valid_mask_flat, t_flat, torch.zeros_like(t_flat))
-
-                # Sample K negatives uniformly from [0, V)
-                neg_ids = torch.randint(0, V, (K_neg,), device=x.device)
-                # Gather lm_head weights for target + negatives
-                all_ids = torch.cat([t_flat_safe, neg_ids])     # (B*T + K,)
-                sampled_w = self.lm_head.weight[all_ids]        # (B*T + K, d)
-
-                # Compute sampled logits: for each position, dot with its
-                # target weight and all K negative weights.
-                # Target logit: dot product of h[i] with w[target[i]]
-                target_w = sampled_w[:n]                        # (B*T, d)
-                neg_w = sampled_w[n:]                           # (K, d)
-                target_logit = (h_flat * target_w).sum(-1)      # (B*T,)
-                neg_logits = h_flat @ neg_w.t()                 # (B*T, K)
-
-                if not _softcap_clamp:
-                    target_logit = softcap * torch.tanh(target_logit / softcap)
-                    neg_logits = softcap * torch.tanh(neg_logits / softcap)
-
-                # Sampled softmax loss: -log(exp(target) / (exp(target) + sum(exp(neg))))
-                # With log-sum-exp correction for sampling K of V negatives.
-                # Correction: add log(V/K) to negative logits to account for
-                # the fact that we're only seeing K of V possible negatives.
-                log_correction = torch.tensor(V / K_neg, device=x.device).log()
-                all_logits = torch.cat([
-                    target_logit.unsqueeze(-1),                 # (B*T, 1)
-                    neg_logits + log_correction,                # (B*T, K)
-                ], dim=-1).float()                              # (B*T, K+1)
-
-                # CE with target always at index 0
-                ce_targets = torch.zeros(n, dtype=torch.long, device=x.device)
-                if reduction == 'none':
-                    per_tok = F.cross_entropy(all_logits, ce_targets, reduction='none')
-                    if self._doc_sep_mask and self._bos_token_id >= 0:
-                        per_tok = torch.where(valid_mask_flat, per_tok, torch.zeros_like(per_tok))
-                    return per_tok
-                per_tok_ce = F.cross_entropy(
-                    all_logits, ce_targets, reduction='none',
-                    label_smoothing=smoothing,
-                )
-                # Mask doc-separator positions. valid_mask_flat is always
-                # computed; when doc_sep_mask is off every token is valid so
-                # this reduces to a plain mean.
-                valid_f = valid_mask_flat.float()
-                valid_n = valid_f.sum().clamp(min=1)
-                out = (per_tok_ce * valid_f).sum() / valid_n
-            else:
-                # Full softmax path (eval or HYDRA_SAMPLED_SOFTMAX=0)
-                chunk_size = int(os.environ.get("HYDRA_CE_CHUNK", "1024"))
-                if chunk_size <= 0:
-                    MAX_LOGITS_BYTES = 256 * 1024 * 1024
-                    tokens_per_chunk = max(V, MAX_LOGITS_BYTES // (V * 4))
-                    chunk_size = max(1, tokens_per_chunk // max(1, B))
-                chunk_size = min(chunk_size, T)
-
-                if reduction == 'none':
-                    loss_parts = []
-                    for start in range(0, T, chunk_size):
-                        end = min(start + chunk_size, T)
-                        chunk_logits = self.lm_head(x[:, start:end, :]).float()
-                        if _softcap_clamp:
-                            chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
-                        else:
-                            chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
-                        chunk_targets = targets[:, start:end].reshape(-1)
-                        chunk_loss = F.cross_entropy(
-                            chunk_logits.view(-1, chunk_logits.size(-1)),
-                            chunk_targets, ignore_index=-1, reduction='none',
-                        )
-                        loss_parts.append(chunk_loss)
-                    return torch.cat(loss_parts)
-
-                total_loss = 0.0
-                total_tokens = 0
-                for start in range(0, T, chunk_size):
-                    end = min(start + chunk_size, T)
-                    chunk_logits = self.lm_head(x[:, start:end, :]).float()
-                    if _softcap_clamp:
-                        chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
-                    else:
-                        chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
-                    chunk_targets = targets[:, start:end].reshape(-1)
-                    chunk_loss = F.cross_entropy(
-                        chunk_logits.view(-1, chunk_logits.size(-1)),
-                        chunk_targets, ignore_index=-1, reduction='sum',
-                        label_smoothing=smoothing,
-                    )
-                    total_loss = total_loss + chunk_loss
-                    total_tokens += (chunk_targets != -1).sum()
-                out = total_loss / total_tokens
-
-            # -----------------------------------------------------------
-            # Learnability #1: Multi-Token Prediction.
-            # For k in {2..K}, add a CE loss at position (t) predicting
-            # the token at position (t+k), using the SAME lm_head weights
-            # (weight-tied). Cost: K-1 extra CEs on a subset of positions.
-            # Only triggered in reduction='mean' path, training only.
-            # -----------------------------------------------------------
-            if reduction == 'mean' and self._mtp_k > 1 and self.training and use_sampled:
-                # TRUE zero-cost MTP: reuse primary's neg_logits (B*T, K_neg)
-                # entirely. Only cost per extra head: O(B*T*d) target-weight
-                # gather + dot product. neg_logits is sliced (view) to match.
-                mtp_loss_sum = out.new_tensor(0.0)
-                mtp_terms = 0
-                # Reshape primary neg_logits back to (B, T, K_neg) so we can slice positions
-                neg_logits_bt = neg_logits.view(B, T, K_neg)
-                for k in range(2, self._mtp_k + 1):
-                    shift = k - 1
-                    if T <= shift:
-                        continue
-                    n_k = B * (T - shift)
-                    h_k_flat = x[:, :T - shift, :].reshape(n_k, -1)  # (n_k, d)
-                    t_k = targets[:, shift:].reshape(-1)             # (n_k,)
-                    mask_k = (t_k >= 0)
-                    t_k_safe = torch.where(mask_k, t_k, torch.zeros_like(t_k))
-                    tgt_w_k = self.lm_head.weight[t_k_safe]          # (n_k, d)
-                    tgt_logit_k = (h_k_flat * tgt_w_k).sum(-1)       # (n_k,)
-                    if not _softcap_clamp:
-                        tgt_logit_k = softcap * torch.tanh(tgt_logit_k / softcap)
-                    # REUSE primary neg_logits — slice positions [:T-shift]
-                    neg_logits_k = neg_logits_bt[:, :T - shift, :].reshape(n_k, K_neg)
-                    all_logits_k = torch.cat([
-                        tgt_logit_k.unsqueeze(-1),
-                        neg_logits_k + log_correction,
-                    ], dim=-1).float()
-                    ce_targets_k = torch.zeros(n_k, dtype=torch.long, device=x.device)
-                    per_tok_ce_k = F.cross_entropy(
-                        all_logits_k, ce_targets_k, reduction='none',
-                        label_smoothing=smoothing,
-                    )
-                    per_tok_ce_k = torch.where(mask_k, per_tok_ce_k, torch.zeros_like(per_tok_ce_k))
-                    n_valid_k = mask_k.sum().clamp(min=1)
-                    mtp_loss_sum = mtp_loss_sum + per_tok_ce_k.sum() / n_valid_k
-                    mtp_terms += 1
-                if mtp_terms > 0:
-                    out = (out + mtp_loss_sum) / float(mtp_terms + 1)
-
-            # -----------------------------------------------------------
-            # Learnability #6: output entropy penalty.
-            # L += -lambda * H(softmax(logits)). Negative entropy penalizes
-            # peaked distributions; encourages diverse predictions and
-            # breaks repetition loops. Computed on a small subset of
-            # positions to keep V-sized logits cost bounded.
-            # -----------------------------------------------------------
-            if reduction == 'mean' and self._entropy_penalty > 0.0 and self.training:
-                # Sample up to 64 random positions. V-sized logits on 64
-                # positions = 64 * V * 4 bytes (~50 MB at V=200k) — fits
-                # on the 3060 and adds ~2 ms.
-                h_flat = x.reshape(-1, x.shape[-1])
-                n_pos = h_flat.shape[0]
-                n_sample = min(64, n_pos)
-                idx_sample = torch.randint(0, n_pos, (n_sample,), device=x.device)
-                h_sample = h_flat[idx_sample]
-                logits_s = F.linear(h_sample, self.lm_head.weight).float()
-                if _softcap_clamp:
-                    logits_s = torch.clamp(logits_s, -softcap, softcap)
-                else:
-                    logits_s = softcap * torch.tanh(logits_s / softcap)
-                log_probs = F.log_softmax(logits_s, dim=-1)
-                probs = log_probs.exp()
-                entropy = -(probs * log_probs).sum(-1).mean()   # scalar, nats
-                out = out - self._entropy_penalty * entropy
-
-            if _profile:
-                _t_end = _ev()
-                torch.cuda.synchronize()
-                def _ms(a, b): return a.elapsed_time(b)
-                print(
-                    f"[PROFILE B={B} T={T}] "
-                    f"htm_launch={_ms(_t0, _t_htm_async):.2f} "
-                    f"wte={_ms(_t_htm_async, _t_wte):.2f} "
-                    f"htm_await={_ms(_t_wte, _t_htm_await):.2f} "
-                    f"htm_proj={_ms(_t_htm_await, _t_htm_proj):.2f} "
-                    f"mamba_mhc_engram={_ms(_t_htm_proj, _t_blocks):.2f} "
-                    f"merge={_ms(_t_blocks, _t_merge):.2f} "
-                    f"lm_head_loss={_ms(_t_merge, _t_end):.2f} "
-                    f"total={_ms(_t0, _t_end):.2f} ms",
-                    flush=True,
-                )
-            return out
-
-        logits = self.lm_head(x).float()
-        if _softcap_clamp:
-            logits = torch.clamp(logits, -softcap, softcap)
-        else:
-            logits = softcap * torch.tanh(logits / softcap)
-        return logits
+"""PostSemClawModel — full-architecture model assembly.
+
+Extracted from the monolithic train.py (W1 modularization). Semantics
+unchanged. Imports `GPUEngram` from `hydra.engram` and `MuonAdamW` from
+`hydra.optimizer`.
+
+Triton kernel integration status (Phase 2):
+  HYDRA_FUSED_BCNORM — DEFERRED. The bcnorm_fused Triton kernel fuses
+    LayerNorm + RoPE on B/C projections. However, mamba-ssm's Mamba3 block
+    uses RMSNormGated (not LayerNorm) for B/C, and RoPE is applied inside
+    the mamba3_siso_combined CUDA kernel via the Angles parameter. Replacing
+    would require either (a) monkey-patching RMSNormGated + intercepting the
+    fused CUDA scan — invasive, 50+ lines, high breakage risk — or (b) a
+    full custom Mamba3Block reimplementation. Both are out of scope for
+    Phase 2. The kernel is validated standalone; integration deferred to
+    Phase 3 when HYDRA moves to a custom SSM block.
+
+  HYDRA_FUSED_SSD — DEFERRED. The ssd_exp_trap Triton kernel implements
+    exponential-trapezoidal discretization as a sequential scan. mamba-ssm's
+    Mamba3 block delegates the entire scan + gating + output projection to
+    mamba3_siso_combined (a compiled CUDA kernel with tilelang). Replacing
+    it would require decomposing the combined kernel into constituent ops
+    and substituting only the scan — not feasible without a custom block.
+    Same Phase 3 gate as above.
+
+Both env vars are accepted but currently no-ops (gates read, logged, but
+the code path is unchanged). This avoids silent regression if someone
+sets them expecting a speedup.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from mamba_ssm import Mamba3
+
+from subsystems.hestia_mini import HestiaQAT
+from subsystems.htm import HTMLayer
+from subsystems.mhc_mini import ManifoldHyperConnection
+from subsystems.sdr_semantic import SemanticFoldingSDR
+
+from hydra.engram import GPUEngram
+from hydra.hyena_block import HyenaBlock
+# GDNBlock is imported lazily inside __init__ so the `fla` dependency is
+# only required when HYDRA_GDN_LAYERS is actually non-empty. Baseline
+# pure-Mamba3 runs continue to work without flash-linear-attention installed.
+from hydra.optimizer import MuonAdamW
+
+
+def norm(x: torch.Tensor) -> torch.Tensor:
+    """RMSNorm over the last dim — stateless, autocast-friendly."""
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class PostSemClawModel(nn.Module):
+    """Full Post-SEM-Claw model assembly.
+
+    Architecture:
+        Token Embedding -> [Mamba3 + residual] x n_layer
+        -> SDR + Engram (at configured layer) -> norm -> LM head
+
+    Interface (must match prepare.py evaluate_bpb):
+        model(x, y, reduction='none').view(-1)  -> per-token losses
+        model(x, y, reduction='mean')           -> scalar loss
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        # Token embedding
+        self.wte = nn.Embedding(config.vocab_size, config.d_model)
+
+        # Mamba-3 blocks — official mamba-ssm fused CUDA kernel. No fallbacks.
+        # RoPE is applied internally by the Mamba3 CUDA kernel via the Angles
+        # parameter; external cos/sin buffers are not needed.
+        #
+        # Hyena supplement: layers whose index appears in `config.hyena_layers`
+        # are instantiated as HyenaBlock instead of Mamba3. The config field
+        # is populated from HYDRA_HYENA_LAYERS at construction time and then
+        # persisted to checkpoints, so resume is safe even when the env var
+        # is unset. Empty tuple → all-Mamba3, byte-identical to pre-port.
+        _hyena_layer_set = set(getattr(config, "hyena_layers", ()) or ())
+        _gdn_layer_set = set(getattr(config, "gdn_layers", ()) or ())
+        # Hyena wins on overlap; conflict is logged at construction time.
+        _both = _hyena_layer_set & _gdn_layer_set
+        if _both:
+            print(f"[WARN] layers in both hyena_layers and gdn_layers; using Hyena: {sorted(_both)}", flush=True)
+            _gdn_layer_set -= _hyena_layer_set
+
+        if _gdn_layer_set:
+            from hydra.gdn_block import GDNBlock  # requires `fla` package
+
+        def _build_block(i: int) -> nn.Module:
+            if i in _hyena_layer_set:
+                return HyenaBlock(
+                    d_model=config.d_model,
+                    seq_len=config.sequence_len,
+                    order=int(os.environ.get("HYDRA_HYENA_ORDER", "2")),
+                    filter_order=int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64")),
+                )
+            if i in _gdn_layer_set:
+                return GDNBlock(
+                    d_model=config.d_model,
+                    n_heads=config.n_heads,
+                )
+            return Mamba3(
+                d_model=config.d_model,
+                d_state=config.d_state,
+                expand=config.expand,
+                headdim=config.headdim,
+                is_mimo=False,          # SISO path uses stable mamba3_siso_combined kernel
+                chunk_size=64,          # upstream-recommended SISO chunk; 16 violated tl.dot M>=16 constraint
+                is_outproj_norm=False,
+                dtype=torch.bfloat16,
+            )
+
+        self.blocks = nn.ModuleList([_build_block(i) for i in range(config.n_layer)])
+
+        # Full-architecture SDR: offline semantic retina + STE (no-bypass).
+        self.sdr_semantic = SemanticFoldingSDR(
+            vocab_size=config.vocab_size,
+            n_bits=config.sdr_n_bits,
+            target_active=config.sdr_target_active,
+            delta_rank=config.sdr_delta_rank,
+            som_warmup_steps=config.sdr_som_warmup,
+            som_update_interval=config.sdr_som_interval,
+        )
+
+        # HTM spatial pooler + temporal memory (Rust, Hebbian).
+        self.htm = HTMLayer(
+            input_bits=config.sdr_n_bits,
+            n_columns=config.htm_n_columns,
+            cells_per_column=config.htm_cells_per_column,
+            batch_size=1,           # grows lazily to actual B on first forward
+            seed=42,
+            learn=True,
+            reset_each_forward=True,
+        )
+
+        # Gradient bridge: (n_columns + anomaly) -> d_model.
+        self.htm_proj = nn.Linear(config.htm_n_columns + 1, config.d_model, bias=False)
+
+        # GPU Engram with Hebbian writes — runs EVERY step.
+        self.engram = GPUEngram(
+            d_model=config.d_model,
+            n_columns=config.engram_n_columns,
+            max_ngram=3,
+        )
+        self.engram_layer_idx = config.engram_layer_idx
+
+        # Manifold-Constrained Hyper-Connections (one per Mamba-3 block).
+        self.mhc = nn.ModuleList([
+            ManifoldHyperConnection(d_model=config.d_model, n_streams=2, sinkhorn_iters=3)
+            for _ in range(config.n_layer)
+        ])
+
+        # Hestia QAT — ternary weight quantization applied post-optimizer-step.
+        self.hestia = HestiaQAT(enabled=True, bits=1.58)
+
+        # LM head
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        # Learnability knob 1: Multi-Token Prediction (Llama-3 style).
+        #   MTP_K=1 -> standard next-token. MTP_K>1 -> extra heads predict
+        #   tokens at positions t+1, t+2, ..., t+K. Heads are weight-tied to
+        #   lm_head (we share Parameters), so the only extra compute is
+        #   additional CE losses; no new params. Activated via HYDRA_MTP_K.
+        self._mtp_k = max(1, int(os.environ.get("HYDRA_MTP_K", "1")))
+
+        # Learnability knob 3: gradient checkpointing on Mamba3 blocks.
+        self._grad_ckpt = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
+
+        # Learnability knob 4: doc-separator BOS masking in packed sequences.
+        self._doc_sep_mask = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
+        # BOS token id is looked up lazily on first forward (requires tokenizer
+        # load); -1 means uninitialized.
+        self._bos_token_id = -1
+
+        # Learnability knob 5: explicit stop-grad on HTM tensor (htm_rust
+        # outputs already have requires_grad=False; this is defense-in-depth).
+        self._htm_stop_grad = os.environ.get("HYDRA_HTM_STOP_GRAD", "0") == "1"
+
+        # Learnability knob 6: entropy penalty coefficient on LM logits.
+        self._entropy_penalty = float(os.environ.get("HYDRA_ENTROPY_PENALTY", "0.0"))
+
+        # Residual dropout
+        self.drop = nn.Dropout(float(os.environ.get("HYDRA_DROPOUT", "0.2")))
+
+        # Logits soft-capping
+        self.softcap = 15.0
+
+        # Secondary metrics storage
+        self._metrics = {}
+
+        # Per-layer diagnostic panel. Env-gated; zero overhead when off.
+        # Emits residual-contribution (delta_ratio), feature std, effective rank,
+        # gradient norm per layer; used to identify minimum viable n_layer + find
+        # entropy leakage / dead layers. See docs/depth-sweep.md.
+        self._diag_enabled = os.environ.get("HYDRA_LAYER_DIAGNOSTICS", "0") == "1"
+        self._diag_step = 0
+        self._diag_svd_every = int(os.environ.get("HYDRA_LAYER_DIAG_SVD_EVERY", "100"))
+        if self._diag_enabled:
+            # Gradient-norm backward hooks on each Mamba3 block output.
+            for _i, _block in enumerate(self.blocks):
+                def _mk_grad_hook(_layer_idx):
+                    def _hook(module, grad_input, grad_output):
+                        if grad_output and grad_output[0] is not None:
+                            g = grad_output[0].detach()
+                            self._metrics[f'layer_{_layer_idx}_grad_norm'] = float(
+                                g.pow(2).mean().sqrt().item()
+                            )
+                    return _hook
+                _block.register_full_backward_hook(_mk_grad_hook(_i))
+
+            # Forward hooks on each Mamba3 block capture the block's OUTPUT
+            # directly. This is the clean measurement: unlike merge_streams()
+            # sampling which sees (streams + M*block_output) in bf16 — where
+            # small block contributions round to zero against unit-norm
+            # residuals — this captures `block_output` itself as produced.
+            # Reports both its absolute RMS norm and its ratio to the block
+            # INPUT's RMS norm (contribution magnitude relative to the
+            # residual it's added to).
+            for _i, _block in enumerate(self.blocks):
+                def _mk_fwd_hook(_layer_idx):
+                    def _hook(module, inputs, output):
+                        with torch.no_grad():
+                            inp = inputs[0].detach().float() if inputs else None
+                            out = output.detach().float() if isinstance(output, torch.Tensor) else None
+                            if out is not None:
+                                out_rms = out.pow(2).mean().sqrt().item()
+                                self._metrics[f'layer_{_layer_idx}_block_out_rms'] = float(out_rms)
+                                if inp is not None:
+                                    in_rms = inp.pow(2).mean().sqrt().item()
+                                    self._metrics[f'layer_{_layer_idx}_block_in_rms'] = float(in_rms)
+                                    self._metrics[f'layer_{_layer_idx}_contrib_ratio'] = float(
+                                        out_rms / (in_rms + 1e-8)
+                                    )
+                    return _hook
+                _block.register_forward_hook(_mk_fwd_hook(_i))
+
+        # Triton kernel integration gates (Phase 2 — deferred, see module docstring).
+        self._fused_bcnorm = os.environ.get("HYDRA_FUSED_BCNORM", "0") == "1"
+        self._fused_ssd = os.environ.get("HYDRA_FUSED_SSD", "0") == "1"
+        if self._fused_bcnorm or self._fused_ssd:
+            import sys
+            _active = []
+            if self._fused_bcnorm:
+                _active.append("HYDRA_FUSED_BCNORM")
+            if self._fused_ssd:
+                _active.append("HYDRA_FUSED_SSD")
+            print(
+                f"[HYDRA] Triton kernel gates set: {', '.join(_active)}. "
+                f"NOTE: Both are DEFERRED (mamba-ssm Mamba3 uses internal "
+                f"CUDA kernels). Gates accepted but currently no-ops.",
+                file=sys.stderr,
+            )
+
+        # R6 optional torch.compile on the impl forward. Gated (default OFF).
+        if os.environ.get("HYDRA_MODEL_COMPILE", "0") == "1":
+            self._forward_impl = torch.compile(
+                self._forward_impl,
+                fullgraph=False,
+                dynamic=True,
+                mode="default",
+            )
+
+    @torch.no_grad()
+    def init_weights(self) -> None:
+        s = 3 ** 0.5 * self.config.d_model ** -0.5
+
+        # Move SDR retina indices (plain attribute, not buffer) to same device as params.
+        # Required because to_empty() only moves params/buffers, and _retina_indices
+        # is loaded from numpy (always CPU) by SemanticFoldingSDR.__init__.
+        device = self.wte.weight.device
+        if hasattr(self.sdr_semantic, '_retina_indices'):
+            self.sdr_semantic._retina_indices = self.sdr_semantic._retina_indices.to(device)
+
+        # Embedding init: GPT-2 / LLaMA convention. std=1.0 was chosen for
+        # vocab=8192; at larger vocabs, smaller std prevents logit blowup.
+        # Use std = 1/sqrt(d_model) which scales sensibly with model width.
+        import math as _math
+        _d_model = self.wte.weight.shape[1]
+        wte_std = float(os.environ.get("HYDRA_WTE_STD", str(1.0 / _math.sqrt(_d_model))))
+        nn.init.normal_(self.wte.weight, mean=0.0, std=wte_std)
+        # LM head init: was std=0.001 — PATHOLOGICAL at vocab>=32k because
+        # logits collapse to zero, loss locks at log(V)~=11, gradient through
+        # head ∝ 1/V is too small to escape. GPT-2 uses std=0.02; LLaMA uses
+        # std=1/sqrt(d_model). Pick 0.02 as robust default, env-overridable.
+        lm_head_std = float(os.environ.get("HYDRA_LM_HEAD_STD", "0.02"))
+        nn.init.normal_(self.lm_head.weight, mean=0.0, std=lm_head_std)
+        # F8 (NOT APPLIED): Weight tying would save V*D params but current LR
+        # groups have embedding_lr=1.0 and unembedding_lr=0.005 × d_model_scale
+        # — tying forces the shared tensor under a single LR group and either
+        # the embeddings learn 200x too slow (under unembed LR) or the LM head
+        # becomes unstable (under embed LR). Short 15-step smoke with tying +
+        # embed-group update showed initial loss jump 9 -> 20. Deferred until
+        # LR groups are re-tuned; see docs/OPTIMIZATION_PLAN.md Post-plan.
+
+        for li, block in enumerate(self.blocks):
+            if hasattr(block, 'in_proj') and hasattr(block.in_proj, 'weight'):
+                nn.init.uniform_(block.in_proj.weight, -s, s)
+            if hasattr(block, 'out_proj') and hasattr(block.out_proj, 'weight'):
+                # GPT-2 residual init: std = 0.02 / sqrt(2 * n_layer).
+                # NOT zeros — zero init makes the block a permanent pass-through
+                # (block_out_rms=0, zero gradient flow to SSM internals).
+                # With non-zero init the block contributes to the residual stream
+                # from step 1, giving the SSM scan actual gradient signal.
+                n_layer = self.config.n_layer
+                out_std = float(os.environ.get(
+                    "HYDRA_OUT_PROJ_STD",
+                    str(0.02 / (2 * n_layer) ** 0.5),
+                ))
+                nn.init.normal_(block.out_proj.weight, mean=0.0, std=out_std)
+
+        nn.init.normal_(self.htm_proj.weight, mean=0.0, std=s)
+
+        # Cast to bf16 to match Mamba3 dtype; Muon groups by shape so mixed
+        # dtypes in the same shape group would break lerp_ dtype checks.
+        self.wte.to(dtype=torch.bfloat16)
+        self.htm_proj.to(dtype=torch.bfloat16)
+        self.engram.to(dtype=torch.bfloat16)
+
+    def set_bos_token_id(self, bos_id: int) -> None:
+        """Inform the model of the tokenizer's BOS id so doc-separator
+        masking (learnability #4) knows which positions to skip. Called from
+        training setup once the tokenizer is loaded."""
+        self._bos_token_id = int(bos_id)
+
+    def invalidate_hyena_caches(self) -> None:
+        """Invalidate filter-rfft caches on all Hyena blocks.
+
+        MUST be called after each `optimizer.step()` when
+        `HYDRA_HYENA_FILTER_CACHE=1` is set, otherwise cached rfft values
+        will be reused with stale filter parameters.
+
+        No-op for blocks that are not HyenaBlock (Mamba3, etc.).
+        """
+        for block in self.blocks:
+            if hasattr(block, "operator") and hasattr(block.operator, "invalidate_filter_cache"):
+                block.operator.invalidate_filter_cache()
+
+    def flush_hyena_pending_grads(self) -> None:
+        """Push pending train-cache filter gradients into filter params.
+
+        Used ONLY when HYDRA_HYENA_TRAIN_CACHE=1. Must be called exactly once
+        per optimizer step, BEFORE `optimizer.step()` and BEFORE
+        `invalidate_hyena_caches()`. The lightning_module wires this in
+        `optimizer_step` around the existing optimizer.step() call.
+
+        No-op if:
+          * No HyenaBlocks are in the model, OR
+          * No micro-batch ever ran with grad enabled (e.g. all-eval step).
+        """
+        for block in self.blocks:
+            if hasattr(block, "operator") and hasattr(block.operator, "flush_pending_filter_grads"):
+                block.operator.flush_pending_filter_grads()
+
+    def estimate_flops(self) -> int:
+        nparams = sum(p.numel() for p in self.parameters())
+        embed_params = self.wte.weight.numel()
+        return 6 * (nparams - embed_params)
+
+    def num_scaling_params(self) -> dict:
+        wte = sum(p.numel() for p in self.wte.parameters())
+        lm_head = sum(p.numel() for p in self.lm_head.parameters())
+        blocks = sum(p.numel() for p in self.blocks.parameters())
+        sdr = sum(p.numel() for p in self.sdr_semantic.parameters())
+        htm_proj = sum(p.numel() for p in self.htm_proj.parameters())
+        engram = sum(p.numel() for p in self.engram.parameters())
+        total = sum(p.numel() for p in self.parameters())
+        return {
+            'wte': wte, 'lm_head': lm_head, 'blocks': blocks,
+            'sdr_semantic': sdr, 'htm_proj': htm_proj,
+            'engram': engram, 'total': total,
+        }
+
+    def get_secondary_metrics(self) -> dict:
+        """Flush any lingering CUDA tensors to host (single sync)."""
+        flushed = {}
+        for k, v in self._metrics.items():
+            if hasattr(v, 'item'):
+                try:
+                    flushed[k] = float(v.item())
+                except Exception:
+                    flushed[k] = v
+            else:
+                flushed[k] = v
+        return flushed
+
+    def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.6, matrix_lr=0.04,
+                        weight_decay=0.2, adam_betas=(0.8, 0.95), scalar_lr=0.5):
+        """Setup MuonAdamW optimizer with per-component LR groups."""
+        model_dim = self.config.d_model
+
+        embedding_params = list(self.wte.parameters())
+        lm_head_params = list(self.lm_head.parameters())
+
+        # Muon routing guard: 2D parameters are NOT automatically matrices.
+        # Exclude:
+        #   (a) params whose name ends in `.freq` — Sin frequency vectors used
+        #       by Hyena's implicit filter MLP. Shape (1, dim) is nominally 2D
+        #       but semantically a per-dim scalar. Muon's polar-express
+        #       orthogonalization would force it toward an orthogonal matrix,
+        #       destroying the learned modulation frequencies.
+        #   (b) 2-D params with min(shape) < MUON_MIN_DIM. Tiny projections
+        #       (e.g. HyenaFilter.implicit_filter.0.weight of shape (64, 3))
+        #       get collapsed toward near-identity by orthogonalization on the
+        #       narrow axis, damaging expressivity. These belong in AdamW.
+        # These exclusions route the params into the AdamW scalar/vector group.
+        MUON_MIN_DIM = 8
+
+        def _muon_eligible(name: str, p: torch.Tensor) -> bool:
+            if p.dim() != 2:
+                return False
+            if name.endswith(".freq"):
+                return False
+            if min(p.shape) < MUON_MIN_DIM:
+                return False
+            return True
+
+        # Matrix params -> Muon (2D weight matrices passing the routing guard).
+        matrix_params = []
+        for name, p in self.blocks.named_parameters():
+            if _muon_eligible(name, p):
+                matrix_params.append(p)
+        # NOTE (W1 audit REG-2): SemanticFoldingSDR.delta_u / delta_v are
+        # currently GRADIENT-DEAD. The forward path uses `binary_only(idx)` for
+        # HTM and stores it as `self._last_sdr`, but does NOT route the STE
+        # output through any downstream op. Including them in the Muon group
+        # burns compute (stack + orthogonalize + lerp) on zero-grad params
+        # every step. Excluded here; a later W5 pass can reconnect STE via a
+        # gated residual if the SDR signal is wanted back in-graph. The
+        # parameters still exist, so no state_dict break.
+        # for p in self.sdr_semantic.parameters():
+        #     if p.dim() == 2:
+        #         matrix_params.append(p)
+        for name, p in self.htm_proj.named_parameters():
+            if _muon_eligible(name, p):
+                matrix_params.append(p)
+        for name, p in self.engram.named_parameters():
+            if _muon_eligible(name, p):
+                matrix_params.append(p)
+
+        # SDR params are intentionally not in any optimizer group — they
+        # receive no gradient in the current forward, so any update would be
+        # pure noise (weight_decay × lr on a zero-grad param).
+        sdr_param_ids = set(id(p) for p in self.sdr_semantic.parameters())
+        assigned = set(id(p) for p in embedding_params + lm_head_params + matrix_params)
+        scalar_params = [
+            p for p in self.parameters()
+            if id(p) not in assigned and id(p) not in sdr_param_ids
+        ]
+
+        total_assigned = len(embedding_params) + len(lm_head_params) + len(matrix_params) + len(scalar_params)
+        total_params = len(list(self.parameters()))
+        sdr_excluded = len(list(self.sdr_semantic.parameters()))
+        assert total_assigned + sdr_excluded == total_params, (
+            f"Parameter count mismatch: assigned {total_assigned} + sdr_excluded "
+            f"{sdr_excluded} vs total {total_params}"
+        )
+
+        dmodel_lr_scale = (model_dim / 768) ** -0.5
+        print(f"Scaling AdamW LRs by 1/sqrt({model_dim}/768) = {dmodel_lr_scale:.6f}")
+
+        param_groups = [
+            dict(kind='adamw', params=lm_head_params,
+                 lr=unembedding_lr * dmodel_lr_scale, betas=adam_betas,
+                 eps=1e-10, weight_decay=0.0),
+            dict(kind='adamw', params=embedding_params,
+                 lr=embedding_lr * dmodel_lr_scale, betas=adam_betas,
+                 eps=1e-10, weight_decay=0.0),
+        ]
+
+        if scalar_params:
+            param_groups.append(
+                dict(kind='adamw', params=scalar_params,
+                     lr=scalar_lr * dmodel_lr_scale, betas=adam_betas,
+                     eps=1e-10, weight_decay=0.0)
+            )
+
+        for shape in sorted({p.shape for p in matrix_params}):
+            group_params = [p for p in matrix_params if p.shape == shape]
+            param_groups.append(dict(
+                kind='muon', params=group_params, lr=matrix_lr,
+                momentum=0.95, ns_steps=5, beta2=0.95, weight_decay=weight_decay,
+            ))
+
+        optimizer = MuonAdamW(param_groups)
+        for group in optimizer.param_groups:
+            group["initial_lr"] = group["lr"]
+        return optimizer
+
+    def forward(self, idx, targets=None, reduction='mean'):
+        """idx: (B, T) int64. Returns loss if targets given, else logits.
+
+        Nested bf16 autocast is a no-op when ambient autocast is already on;
+        when it's off (e.g. integration tests) we establish the dtype contract.
+        """
+        if torch.is_autocast_enabled():
+            return self._forward_impl(idx, targets=targets, reduction=reduction)
+        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+            return self._forward_impl(idx, targets=targets, reduction=reduction)
+
+    def _forward_impl(self, idx, targets=None, reduction='mean'):
+        B, T = idx.shape
+
+        # Diagnostic: per-subsystem CUDA event timing. Env-gated; zero overhead
+        # when disabled. Logs one timing line per forward call. Used to isolate
+        # which subsystem is the tps bottleneck on paid hardware.
+        _profile = os.environ.get("HYDRA_PROFILE_FORWARD", "0") == "1"
+        if _profile:
+            def _ev():
+                e = torch.cuda.Event(enable_timing=True)
+                e.record()
+                return e
+            _t0 = _ev()
+        else:
+            _t0 = None
+
+        # Compute SDR binary ONCE and reuse for both HTM input and the stash.
+        sdr_binary = self.sdr_semantic.binary_only(idx)
+        self._last_sdr = sdr_binary  # uint8 stash (not bf16 → 256MB avoidance)
+
+        # HTM subsampling: run HTM on 1 of every N micro-batches within a
+        # gradient accumulation step, reuse the cached result for the other
+        # N-1 micro-batches. Cooperative launch monopolizes all SMs (grid.sync
+        # requires full-grid residency), so HTM and mamba can't overlap via
+        # streams. Subsampling removes HTM from most micro-batches' critical
+        # path instead.
+        #
+        # Math: N=8, 64 accum steps → 8 HTM calls (10.6ms each) + 56 fast
+        # calls (4ms each). Total = 84.8 + 224 = 309ms → 106k tps.
+        #
+        # HYDRA_HTM_SUBSAMPLE=N (default 8). Set =1 for every-microbatch HTM.
+        _htm_sub = int(os.environ.get("HYDRA_HTM_SUBSAMPLE", "8"))
+        if not hasattr(self, '_htm_call_idx'):
+            self._htm_call_idx = 0
+
+        _run_htm = (self._htm_call_idx % _htm_sub == 0)
+        self._htm_call_idx += 1
+
+        if _run_htm:
+            htm_handle = self.htm.forward_async(sdr_binary)
+        else:
+            htm_handle = None
+
+        if _profile: _t_htm_async = _ev()
+
+        dense_emb = self.wte(idx)  # (B, T, d_model) bf16
+
+        if _profile: _t_wte = _ev()
+
+        if _run_htm:
+            htm_out = self.htm.forward_await(htm_handle)
+            self._htm_cache = htm_out.detach()  # cache for non-HTM micro-batches
+        elif hasattr(self, '_htm_cache') and self._htm_cache is not None \
+                and self._htm_cache.shape[0] == B and self._htm_cache.shape[1] == T:
+            htm_out = self._htm_cache
+        else:
+            # Very first call with subsample > 1: run HTM anyway.
+            htm_handle = self.htm.forward_async(sdr_binary)
+            htm_out = self.htm.forward_await(htm_handle)
+            self._htm_cache = htm_out.detach()
+
+        if _profile: _t_htm_await = _ev()
+        with torch.no_grad():
+            sdr_active_bits = float(self.sdr_semantic.target_active)
+            htm_anomaly = htm_out[..., -1].mean()
+
+        # Learnability #5: explicit stop-grad on HTM output. htm_rust already
+        # produces a detached tensor, but making it explicit here hardens the
+        # contract against future refactors that might route HTM through a
+        # grad-enabled op.
+        if self._htm_stop_grad:
+            htm_out = htm_out.detach()
+
+        # Gradient bridge: HTM columns+anomaly -> d_model.
+        htm_proj_out = self.htm_proj(htm_out.to(dense_emb.dtype))
+        x = dense_emb + htm_proj_out
+        x = norm(x)
+
+        if _profile: _t_htm_proj = _ev()
+
+        # mHC-routed Mamba-3 stack with Engram injection at configured layer.
+        streams = self.mhc[0].init_streams(x)
+        _engram_ev = None
+
+        # Per-layer diagnostic panel. The pre-layer merged state h_pre lets us
+        # measure residual contribution of each layer: delta_N = h_post - h_pre.
+        # All reads are detached no-grad to avoid autograd graph pollution.
+        _diag = self._diag_enabled
+        if _diag:
+            # Cast to float32 for the diagnostic arithmetic: the layer's
+            # residual contribution is small (~0.5 × rms-normed block output),
+            # which underflows in bf16 subtraction (3-digit mantissa) and
+            # reports delta_ratio=0 at the boundaries. float32 snapshot is
+            # ~3.8 MB extra memory per diag sample (B=1, T=2048, d=96) —
+            # negligible vs peak VRAM.
+            with torch.no_grad():
+                h_pre = self.mhc[0].merge_streams(streams).detach().float()
+            _run_svd = (self._diag_step % self._diag_svd_every) == 0
+
+        for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
+            def _block_fn(h, _block=block):
+                return self.drop(_block(norm(h)))
+
+            # Learnability #3: gradient checkpointing. Wrap the block-fn so
+            # the mhc layer's internal uses of it re-run the block in backward
+            # (trading compute for activation memory). use_reentrant=False is
+            # the modern API and works cleanly under autocast.
+            if self._grad_ckpt and self.training:
+                import torch.utils.checkpoint as _ckpt
+                _raw_fn = _block_fn
+                def _block_fn(h, _raw=_raw_fn):  # noqa: E731
+                    return _ckpt.checkpoint(_raw, h, use_reentrant=False)
+
+            streams = mhc_layer(streams, _block_fn)
+
+            if i == self.engram_layer_idx:
+                if _profile: _t_pre_engram = _ev()
+                x_mid = mhc_layer.merge_streams(streams)
+                x_mid, hit_rate = self.engram(x_mid, idx)
+                streams = mhc_layer.init_streams(x_mid)
+                self._metrics['engram_hit_rate'] = hit_rate
+                if _profile: _engram_ev = _ev()
+
+            if _diag:
+                with torch.no_grad():
+                    h_post = mhc_layer.merge_streams(streams).detach().float()
+                    in_n  = h_pre.pow(2).mean().sqrt()
+                    out_n = h_post.pow(2).mean().sqrt()
+                    d_n   = (h_post - h_pre).pow(2).mean().sqrt()
+                    self._metrics[f'layer_{i}_in_norm']     = float(in_n.item())
+                    self._metrics[f'layer_{i}_out_norm']    = float(out_n.item())
+                    self._metrics[f'layer_{i}_delta_ratio'] = float((d_n / (in_n + 1e-6)).item())
+                    self._metrics[f'layer_{i}_feat_std']    = float(h_post.std(dim=-1).mean().item())
+                    if _run_svd:
+                        # Effective rank via participation ratio of singular values.
+                        # eff_rank = (Σσ)^2 / Σσ² — smooth rank proxy, bounded by d_model.
+                        # Sampled to keep overhead low (SVD is O(min(B*T, D)^2·D)).
+                        flat = h_post.reshape(-1, h_post.shape[-1])[:512].float()
+                        try:
+                            s = torch.linalg.svdvals(flat)
+                            eff_rank = float(((s.sum() ** 2) / (s.pow(2).sum() + 1e-6)).item())
+                            self._metrics[f'layer_{i}_eff_rank'] = eff_rank
+                        except Exception:
+                            pass
+                    h_pre = h_post
+
+        if _diag:
+            self._diag_step += 1
+
+        if _profile: _t_blocks = _ev()
+
+        self._metrics['sdr_active_bits'] = sdr_active_bits
+        self._metrics['htm_anomaly'] = htm_anomaly
+
+        x = self.mhc[-1].merge_streams(streams)
+        x = norm(x)
+
+        if _profile: _t_merge = _ev()
+
+        softcap = self.softcap
+        _softcap_clamp = os.environ.get("HYDRA_SOFTCAP_CLAMP", "0") == "1"
+        if targets is not None:
+            smoothing = self.config.label_smoothing
+            V = self.config.vocab_size
+
+            # Learnability #4: doc-separator masking. In packed rows,
+            # tokenizer.encode(..., prepend=bos_token) places a BOS at every
+            # document boundary. Without masking, the model is penalized for
+            # failing to predict "doc B's BOS" from the last tokens of doc A
+            # — pure noise. We set targets==bos to -1 (ignore_index). Done
+            # BEFORE MTP/entropy/sampled-softmax branches so all downstream
+            # losses inherit the mask.
+            if self._doc_sep_mask and self._bos_token_id >= 0:
+                targets = torch.where(
+                    targets == self._bos_token_id,
+                    torch.full_like(targets, -1),
+                    targets,
+                )
+
+            # Sampled softmax: instead of computing logits for ALL V tokens,
+            # compute only for the target + K random negatives. Reduces the
+            # lm_head matmul from (B*T, d) × (d, V) to (B*T, d) × (d, K+1).
+            # At V=65536 and K=4096: 16× less compute, ~4× tps improvement.
+            # The log-sum-exp correction adjusts for the sampling bias.
+            # Set HYDRA_SAMPLED_SOFTMAX=0 to disable (full softmax).
+            K_neg = int(os.environ.get("HYDRA_SAMPLED_SOFTMAX", "4096"))
+            use_sampled = K_neg > 0 and K_neg < V and self.training
+
+            if use_sampled:
+                # Flatten hidden states + targets
+                h_flat = x.reshape(-1, x.shape[-1])            # (B*T, d)
+                t_flat = targets.reshape(-1)                    # (B*T,)
+                n = h_flat.shape[0]
+
+                # Learnability #4 hardening: sampled-softmax gather crashes on
+                # negative ids (-1 from doc-sep mask). Replace -1 with 0 for
+                # gather; the actual loss is masked below.
+                valid_mask_flat = (t_flat >= 0)
+                t_flat_safe = torch.where(valid_mask_flat, t_flat, torch.zeros_like(t_flat))
+
+                # Sample K negatives uniformly from [0, V)
+                neg_ids = torch.randint(0, V, (K_neg,), device=x.device)
+                # Gather lm_head weights for target + negatives
+                all_ids = torch.cat([t_flat_safe, neg_ids])     # (B*T + K,)
+                sampled_w = self.lm_head.weight[all_ids]        # (B*T + K, d)
+
+                # Compute sampled logits: for each position, dot with its
+                # target weight and all K negative weights.
+                # Target logit: dot product of h[i] with w[target[i]]
+                target_w = sampled_w[:n]                        # (B*T, d)
+                neg_w = sampled_w[n:]                           # (K, d)
+                target_logit = (h_flat * target_w).sum(-1)      # (B*T,)
+                neg_logits = h_flat @ neg_w.t()                 # (B*T, K)
+
+                if not _softcap_clamp:
+                    target_logit = softcap * torch.tanh(target_logit / softcap)
+                    neg_logits = softcap * torch.tanh(neg_logits / softcap)
+
+                # Sampled softmax loss: -log(exp(target) / (exp(target) + sum(exp(neg))))
+                # With log-sum-exp correction for sampling K of V negatives.
+                # Correction: add log(V/K) to negative logits to account for
+                # the fact that we're only seeing K of V possible negatives.
+                log_correction = torch.tensor(V / K_neg, device=x.device).log()
+                all_logits = torch.cat([
+                    target_logit.unsqueeze(-1),                 # (B*T, 1)
+                    neg_logits + log_correction,                # (B*T, K)
+                ], dim=-1).float()                              # (B*T, K+1)
+
+                # CE with target always at index 0
+                ce_targets = torch.zeros(n, dtype=torch.long, device=x.device)
+                if reduction == 'none':
+                    per_tok = F.cross_entropy(all_logits, ce_targets, reduction='none')
+                    if self._doc_sep_mask and self._bos_token_id >= 0:
+                        per_tok = torch.where(valid_mask_flat, per_tok, torch.zeros_like(per_tok))
+                    return per_tok
+                per_tok_ce = F.cross_entropy(
+                    all_logits, ce_targets, reduction='none',
+                    label_smoothing=smoothing,
+                )
+                # Mask doc-separator positions. valid_mask_flat is always
+                # computed; when doc_sep_mask is off every token is valid so
+                # this reduces to a plain mean.
+                valid_f = valid_mask_flat.float()
+                valid_n = valid_f.sum().clamp(min=1)
+                out = (per_tok_ce * valid_f).sum() / valid_n
+            else:
+                # Full softmax path (eval or HYDRA_SAMPLED_SOFTMAX=0)
+                chunk_size = int(os.environ.get("HYDRA_CE_CHUNK", "1024"))
+                if chunk_size <= 0:
+                    MAX_LOGITS_BYTES = 256 * 1024 * 1024
+                    tokens_per_chunk = max(V, MAX_LOGITS_BYTES // (V * 4))
+                    chunk_size = max(1, tokens_per_chunk // max(1, B))
+                chunk_size = min(chunk_size, T)
+
+                if reduction == 'none':
+                    loss_parts = []
+                    for start in range(0, T, chunk_size):
+                        end = min(start + chunk_size, T)
+                        chunk_logits = self.lm_head(x[:, start:end, :]).float()
+                        if _softcap_clamp:
+                            chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
+                        else:
+                            chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
+                        chunk_targets = targets[:, start:end].reshape(-1)
+                        chunk_loss = F.cross_entropy(
+                            chunk_logits.view(-1, chunk_logits.size(-1)),
+                            chunk_targets, ignore_index=-1, reduction='none',
+                        )
+                        loss_parts.append(chunk_loss)
+                    return torch.cat(loss_parts)
+
+                total_loss = 0.0
+                total_tokens = 0
+                for start in range(0, T, chunk_size):
+                    end = min(start + chunk_size, T)
+                    chunk_logits = self.lm_head(x[:, start:end, :]).float()
+                    if _softcap_clamp:
+                        chunk_logits = torch.clamp(chunk_logits, -softcap, softcap)
+                    else:
+                        chunk_logits = softcap * torch.tanh(chunk_logits / softcap)
+                    chunk_targets = targets[:, start:end].reshape(-1)
+                    chunk_loss = F.cross_entropy(
+                        chunk_logits.view(-1, chunk_logits.size(-1)),
+                        chunk_targets, ignore_index=-1, reduction='sum',
+                        label_smoothing=smoothing,
+                    )
+                    total_loss = total_loss + chunk_loss
+                    total_tokens += (chunk_targets != -1).sum()
+                out = total_loss / total_tokens
+
+            # -----------------------------------------------------------
+            # Learnability #1: Multi-Token Prediction.
+            # For k in {2..K}, add a CE loss at position (t) predicting
+            # the token at position (t+k), using the SAME lm_head weights
+            # (weight-tied). Cost: K-1 extra CEs on a subset of positions.
+            # Only triggered in reduction='mean' path, training only.
+            # -----------------------------------------------------------
+            if reduction == 'mean' and self._mtp_k > 1 and self.training and use_sampled:
+                # TRUE zero-cost MTP: reuse primary's neg_logits (B*T, K_neg)
+                # entirely. Only cost per extra head: O(B*T*d) target-weight
+                # gather + dot product. neg_logits is sliced (view) to match.
+                mtp_loss_sum = out.new_tensor(0.0)
+                mtp_terms = 0
+                # Reshape primary neg_logits back to (B, T, K_neg) so we can slice positions
+                neg_logits_bt = neg_logits.view(B, T, K_neg)
+                for k in range(2, self._mtp_k + 1):
+                    shift = k - 1
+                    if T <= shift:
+                        continue
+                    n_k = B * (T - shift)
+                    h_k_flat = x[:, :T - shift, :].reshape(n_k, -1)  # (n_k, d)
+                    t_k = targets[:, shift:].reshape(-1)             # (n_k,)
+                    mask_k = (t_k >= 0)
+                    t_k_safe = torch.where(mask_k, t_k, torch.zeros_like(t_k))
+                    tgt_w_k = self.lm_head.weight[t_k_safe]          # (n_k, d)
+                    tgt_logit_k = (h_k_flat * tgt_w_k).sum(-1)       # (n_k,)
+                    if not _softcap_clamp:
+                        tgt_logit_k = softcap * torch.tanh(tgt_logit_k / softcap)
+                    # REUSE primary neg_logits — slice positions [:T-shift]
+                    neg_logits_k = neg_logits_bt[:, :T - shift, :].reshape(n_k, K_neg)
+                    all_logits_k = torch.cat([
+                        tgt_logit_k.unsqueeze(-1),
+                        neg_logits_k + log_correction,
+                    ], dim=-1).float()
+                    ce_targets_k = torch.zeros(n_k, dtype=torch.long, device=x.device)
+                    per_tok_ce_k = F.cross_entropy(
+                        all_logits_k, ce_targets_k, reduction='none',
+                        label_smoothing=smoothing,
+                    )
+                    per_tok_ce_k = torch.where(mask_k, per_tok_ce_k, torch.zeros_like(per_tok_ce_k))
+                    n_valid_k = mask_k.sum().clamp(min=1)
+                    mtp_loss_sum = mtp_loss_sum + per_tok_ce_k.sum() / n_valid_k
+                    mtp_terms += 1
+                if mtp_terms > 0:
+                    out = (out + mtp_loss_sum) / float(mtp_terms + 1)
+
+            # -----------------------------------------------------------
+            # Learnability #6: output entropy penalty.
+            # L += -lambda * H(softmax(logits)). Negative entropy penalizes
+            # peaked distributions; encourages diverse predictions and
+            # breaks repetition loops. Computed on a small subset of
+            # positions to keep V-sized logits cost bounded.
+            # -----------------------------------------------------------
+            if reduction == 'mean' and self._entropy_penalty > 0.0 and self.training:
+                # Sample up to 64 random positions. V-sized logits on 64
+                # positions = 64 * V * 4 bytes (~50 MB at V=200k) — fits
+                # on the 3060 and adds ~2 ms.
+                h_flat = x.reshape(-1, x.shape[-1])
+                n_pos = h_flat.shape[0]
+                n_sample = min(64, n_pos)
+                idx_sample = torch.randint(0, n_pos, (n_sample,), device=x.device)
+                h_sample = h_flat[idx_sample]
+                logits_s = F.linear(h_sample, self.lm_head.weight).float()
+                if _softcap_clamp:
+                    logits_s = torch.clamp(logits_s, -softcap, softcap)
+                else:
+                    logits_s = softcap * torch.tanh(logits_s / softcap)
+                log_probs = F.log_softmax(logits_s, dim=-1)
+                probs = log_probs.exp()
+                entropy = -(probs * log_probs).sum(-1).mean()   # scalar, nats
+                out = out - self._entropy_penalty * entropy
+
+            if _profile:
+                _t_end = _ev()
+                torch.cuda.synchronize()
+                def _ms(a, b): return a.elapsed_time(b)
+                print(
+                    f"[PROFILE B={B} T={T}] "
+                    f"htm_launch={_ms(_t0, _t_htm_async):.2f} "
+                    f"wte={_ms(_t_htm_async, _t_wte):.2f} "
+                    f"htm_await={_ms(_t_wte, _t_htm_await):.2f} "
+                    f"htm_proj={_ms(_t_htm_await, _t_htm_proj):.2f} "
+                    f"mamba_mhc_engram={_ms(_t_htm_proj, _t_blocks):.2f} "
+                    f"merge={_ms(_t_blocks, _t_merge):.2f} "
+                    f"lm_head_loss={_ms(_t_merge, _t_end):.2f} "
+                    f"total={_ms(_t0, _t_end):.2f} ms",
+                    flush=True,
+                )
+            return out
+
+        logits = self.lm_head(x).float()
+        if _softcap_clamp:
+            logits = torch.clamp(logits, -softcap, softcap)
+        else:
+            logits = softcap * torch.tanh(logits / softcap)
+        return logits

overlay/hydra/optimizer.py

@@ -1,252 +1,252 @@
-"""MuonAdamW optimizer — combined Muon (2D matrices) + AdamW (everything else).
-
-Extracted verbatim from train.py (W1 modularization). Semantics unchanged.
-
-F1-F15 state preserved:
-- F7 REVERTED: `stacked_params_buf` persistent across steps was REMOVED — each
-  step calls `torch.stack([p.grad for p in params])` / `torch.stack(params)`
-  fresh. Persistent copies of param storage would be mutated between forward
-  passes (via lerp_/sub_ on stacked tensors that share storage with params),
-  triggering "modified in-place" errors on grad_accum=2 backwards.
-- F11/F15: `@torch.compile` on `adamw_step_fused` / `muon_step_fused` intact.
-- F15 compile is default-ON (HYDRA_MUON_COMPILE=1), configured with
-  dynamic=True + mode="default" to avoid the step-17→18 cudagraphs
-  stream-capture deadlock. See .omc/muon_compile_bug.md for the full
-  investigation.
-"""
-
-from __future__ import annotations
-
-import os
-
-import torch
-
-# HYDRA_FUSED_ADAMW=1 (default) -> vectorized torch._fused_adamw_ kernel.
-_HYDRA_FUSED_ADAMW = os.environ.get("HYDRA_FUSED_ADAMW", "1") == "1"
-_HAS_FUSED_ADAMW = hasattr(torch, "_fused_adamw_")
-
-
-polar_express_coeffs = [
-    (8.156554524902461, -22.48329292557795, 15.878769915207462),
-    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
-    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
-    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
-    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
-]
-
-
-def adamw_step_fused(p, grad, exp_avg, exp_avg_sq, step_t, lr_t, beta1_t, beta2_t, eps_t, wd_t):
-    # Per-param AdamW fallback. Fast path is torch._fused_adamw_ (1 CUDA launch
-    # for the whole group) driven from MuonAdamW._step_adamw below.
-    grad = grad.to(p.dtype)  # handle mixed bf16/fp32 from autocast
-    p.mul_(1 - lr_t * wd_t)
-    exp_avg.lerp_(grad, 1 - beta1_t)
-    exp_avg_sq.lerp_(grad.square(), 1 - beta2_t)
-    bias1 = 1 - beta1_t ** step_t
-    bias2 = 1 - beta2_t ** step_t
-    denom = (exp_avg_sq / bias2).sqrt() + eps_t
-    step_size = lr_t / bias1
-    p.add_(exp_avg / denom, alpha=-step_size)
-
-
-# ---------------------------------------------------------------------------
-# F15 muon_step_fused compile strategy.
-#
-# HYDRA_MUON_COMPILE env gate:
-#   "1" (default ON) — wrap with torch.compile(dynamic=True, mode="default").
-#       Dynamic=True collapses the per-shape specialization cache so that N
-#       Muon param-groups with N distinct shapes trigger 1 compile, not N.
-#       mode="default" keeps the inductor codegen but disables cudagraphs,
-#       which is what caused the step-17→18 silent deadlock observed under
-#       the original dynamic=False configuration: cudagraph stream capture
-#       can deadlock against HTM's CUDA kernels running on the default
-#       stream, and the failure mode at capture-time is a silent hang
-#       (100% GPU util, no log output, process state R).
-#   "0" — fall back to eager Python (slower, ~43k tps vs ~63k compiled).
-#       Keeps an escape hatch in case a future torch/inductor regression
-#       reintroduces a deadlock.
-#
-# Defensive .clone() on stacked_grads before in-place lerp_ eliminates the
-# alias-analysis edge case where inductor sees `g is stacked_grads` and
-# subsequent `stacked_grads.square()` operating on the post-lerp storage.
-# ---------------------------------------------------------------------------
-_MUON_COMPILE = os.environ.get("HYDRA_MUON_COMPILE", "1") == "1"
-
-def _maybe_compile(fn):
-    if _MUON_COMPILE:
-        # mode="default" explicitly opts OUT of cudagraphs (which reduce-overhead
-        # would enable) to avoid stream-capture deadlocks against HTM's CUDA
-        # kernels. dynamic=True minimizes recompile count across param-group
-        # shapes.
-        return torch.compile(fn, fullgraph=False, dynamic=True, mode="default")
-    return fn
-
-@_maybe_compile
-def muon_step_fused(stacked_grads, stacked_params, momentum_buffer, second_momentum_buffer,
-                    momentum_t, lr_t, wd_t, beta2_t, ns_steps, red_dim):
-    # Cast grads to param dtype AND clone defensively to break any alias
-    # between the (freshly-stacked) input and the in-place lerp_ below.
-    # Without this, inductor's alias analysis can emit code that reads from
-    # post-mutation storage when computing `v_mean = g.square().mean(...)`.
-    stacked_grads = stacked_grads.to(momentum_buffer.dtype).clone()
-    # Nesterov momentum
-    momentum = momentum_t.to(device=momentum_buffer.device, dtype=stacked_grads.dtype)
-    momentum_buffer.lerp_(stacked_grads, 1 - momentum)
-    g = stacked_grads.lerp_(momentum_buffer, momentum)
-    # Polar express orthogonalization
-    X = g.bfloat16()
-    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
-    if g.size(-2) > g.size(-1):
-        for a, b, c in polar_express_coeffs[:ns_steps]:
-            A = X.mT @ X
-            B = b * A + c * (A @ A)
-            X = a * X + X @ B
-    else:
-        for a, b, c in polar_express_coeffs[:ns_steps]:
-            A = X @ X.mT
-            B = b * A + c * (A @ A)
-            X = a * X + B @ X
-    g = X
-    # NorMuon variance reduction
-    # Keep beta2 in the state-buffer dtype, not g.dtype, so lerp_ on the
-    # float32 second_momentum_buffer doesn't hit a dtype mismatch on h200.
-    beta2 = beta2_t.to(device=second_momentum_buffer.device, dtype=second_momentum_buffer.dtype)
-    v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
-    red_dim_size = g.size(red_dim)
-    v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * red_dim_size
-    v_norm = v_norm_sq.sqrt()
-    second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2)
-    step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()
-    scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square()
-    v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt()
-    final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10))
-    g = g * final_scale.to(g.dtype)
-    # Cautious weight decay + parameter update
-    lr = lr_t.to(device=stacked_params.device, dtype=g.dtype)
-    wd = wd_t.to(device=stacked_params.device, dtype=g.dtype)
-    mask = (g * stacked_params) >= 0
-    stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)
-
-
-class MuonAdamW(torch.optim.Optimizer):
-    """Combined optimizer: Muon for 2D matrix params, AdamW for others."""
-
-    def __init__(self, param_groups):
-        super().__init__(param_groups, defaults={})
-        # 0-D CPU tensors to avoid torch.compile recompilation when values change
-        self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-        self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
-
-    def _step_adamw(self, group):
-        params, grads, exp_avgs, exp_avg_sqs, state_steps = [], [], [], [], []
-        for p in group['params']:
-            if p.grad is None:
-                continue
-            state = self.state[p]
-            if not state:
-                state['step'] = 0
-                state['exp_avg'] = torch.zeros_like(p)
-                state['exp_avg_sq'] = torch.zeros_like(p)
-            if 'step_t' not in state:
-                # _fused_adamw_ wants a per-param float step tensor on-device.
-                state['step_t'] = torch.tensor(
-                    float(state['step']), dtype=torch.float32, device=p.device
-                )
-            state['step'] += 1
-            params.append(p)
-            grads.append(p.grad.to(p.dtype) if p.grad.dtype != p.dtype else p.grad)
-            exp_avgs.append(state['exp_avg'])
-            exp_avg_sqs.append(state['exp_avg_sq'])
-            state_steps.append(state['step_t'])
-
-        if not params:
-            return
-
-        if _HYDRA_FUSED_ADAMW and _HAS_FUSED_ADAMW and params[0].is_cuda:
-            # _fused_adamw_ needs uniform (device, dtype) within a call, so
-            # group by (device, dtype) — same pattern as PyTorch's own
-            # AdamW(fused=True) path (_group_tensors_by_device_and_dtype).
-            buckets = {}
-            for p, g, ea, es, st in zip(params, grads, exp_avgs, exp_avg_sqs, state_steps):
-                key = (p.device, p.dtype)
-                buckets.setdefault(key, ([], [], [], [], []))
-                b_p, b_g, b_ea, b_es, b_st = buckets[key]
-                b_p.append(p); b_g.append(g); b_ea.append(ea); b_es.append(es); b_st.append(st)
-
-            lr_f = float(group['lr'])
-            b1_f = float(group['betas'][0])
-            b2_f = float(group['betas'][1])
-            wd_f = float(group['weight_decay'])
-            eps_f = float(group['eps'])
-            for (_dev, _dt), (b_p, b_g, b_ea, b_es, b_st) in buckets.items():
-                torch._foreach_add_(b_st, 1.0)
-                torch._fused_adamw_(
-                    b_p, b_g, b_ea, b_es,
-                    [],  # max_exp_avg_sqs unused (amsgrad=False)
-                    b_st,
-                    amsgrad=False,
-                    lr=lr_f, beta1=b1_f, beta2=b2_f,
-                    weight_decay=wd_f, eps=eps_f,
-                    maximize=False,
-                    grad_scale=None, found_inf=None,
-                )
-            return
-
-        # Fallback per-param path.
-        self._adamw_lr_t.fill_(group['lr'])
-        self._adamw_beta1_t.fill_(group['betas'][0])
-        self._adamw_beta2_t.fill_(group['betas'][1])
-        self._adamw_eps_t.fill_(group['eps'])
-        self._adamw_wd_t.fill_(group['weight_decay'])
-        for p, grad, exp_avg, exp_avg_sq in zip(params, grads, exp_avgs, exp_avg_sqs):
-            self._adamw_step_t.fill_(self.state[p]['step'])
-            adamw_step_fused(p, grad, exp_avg, exp_avg_sq,
-                             self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t,
-                             self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t)
-
-    def _step_muon(self, group):
-        params = [p for p in group['params'] if p.grad is not None]
-        if not params:
-            return
-        p = params[0]
-        state = self.state[p]
-        num_params = len(params)
-        shape, device, dtype = p.shape, p.device, p.dtype
-        if "momentum_buffer" not in state:
-            state["momentum_buffer"] = torch.zeros(num_params, *shape, dtype=dtype, device=device)
-        red_dim = -1 if shape[-2] >= shape[-1] else -2
-        if "second_momentum_buffer" not in state:
-            # Shape must match v_mean = stacked_grads.square().mean(dim=red_dim, keepdim=True)
-            full_shape = (num_params, *shape)
-            state_shape = list(full_shape)
-            state_shape[len(state_shape) + red_dim] = 1  # red_dim is negative
-            state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
-        # F7 REVERT: fresh stacks each step (no persistent stacked_params_buf).
-        # This was the autograd-safety fix that unblocks grad_accum>=2.
-        stacked_grads = torch.stack([p.grad for p in params])
-        stacked_params = torch.stack(params)
-        self._muon_momentum_t.fill_(group["momentum"])
-        self._muon_beta2_t.fill_(group["beta2"] if group["beta2"] is not None else 0.0)
-        self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1]) ** 0.5)
-        self._muon_wd_t.fill_(group["weight_decay"])
-        muon_step_fused(stacked_grads, stacked_params,
-                        state["momentum_buffer"], state["second_momentum_buffer"],
-                        self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t,
-                        self._muon_beta2_t, group["ns_steps"], red_dim)
-        torch._foreach_copy_(params, list(stacked_params.unbind(0)))
-
-    @torch.no_grad()
-    def step(self):
-        for group in self.param_groups:
-            if group['kind'] == 'adamw':
-                self._step_adamw(group)
-            elif group['kind'] == 'muon':
-                self._step_muon(group)
+"""MuonAdamW optimizer — combined Muon (2D matrices) + AdamW (everything else).
+
+Extracted verbatim from train.py (W1 modularization). Semantics unchanged.
+
+F1-F15 state preserved:
+- F7 REVERTED: `stacked_params_buf` persistent across steps was REMOVED — each
+  step calls `torch.stack([p.grad for p in params])` / `torch.stack(params)`
+  fresh. Persistent copies of param storage would be mutated between forward
+  passes (via lerp_/sub_ on stacked tensors that share storage with params),
+  triggering "modified in-place" errors on grad_accum=2 backwards.
+- F11/F15: `@torch.compile` on `adamw_step_fused` / `muon_step_fused` intact.
+- F15 compile is default-ON (HYDRA_MUON_COMPILE=1), configured with
+  dynamic=True + mode="default" to avoid the step-17→18 cudagraphs
+  stream-capture deadlock. See .omc/muon_compile_bug.md for the full
+  investigation.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+# HYDRA_FUSED_ADAMW=1 (default) -> vectorized torch._fused_adamw_ kernel.
+_HYDRA_FUSED_ADAMW = os.environ.get("HYDRA_FUSED_ADAMW", "1") == "1"
+_HAS_FUSED_ADAMW = hasattr(torch, "_fused_adamw_")
+
+
+polar_express_coeffs = [
+    (8.156554524902461, -22.48329292557795, 15.878769915207462),
+    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
+    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
+    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
+    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
+]
+
+
+def adamw_step_fused(p, grad, exp_avg, exp_avg_sq, step_t, lr_t, beta1_t, beta2_t, eps_t, wd_t):
+    # Per-param AdamW fallback. Fast path is torch._fused_adamw_ (1 CUDA launch
+    # for the whole group) driven from MuonAdamW._step_adamw below.
+    grad = grad.to(p.dtype)  # handle mixed bf16/fp32 from autocast
+    p.mul_(1 - lr_t * wd_t)
+    exp_avg.lerp_(grad, 1 - beta1_t)
+    exp_avg_sq.lerp_(grad.square(), 1 - beta2_t)
+    bias1 = 1 - beta1_t ** step_t
+    bias2 = 1 - beta2_t ** step_t
+    denom = (exp_avg_sq / bias2).sqrt() + eps_t
+    step_size = lr_t / bias1
+    p.add_(exp_avg / denom, alpha=-step_size)
+
+
+# ---------------------------------------------------------------------------
+# F15 muon_step_fused compile strategy.
+#
+# HYDRA_MUON_COMPILE env gate:
+#   "1" (default ON) — wrap with torch.compile(dynamic=True, mode="default").
+#       Dynamic=True collapses the per-shape specialization cache so that N
+#       Muon param-groups with N distinct shapes trigger 1 compile, not N.
+#       mode="default" keeps the inductor codegen but disables cudagraphs,
+#       which is what caused the step-17→18 silent deadlock observed under
+#       the original dynamic=False configuration: cudagraph stream capture
+#       can deadlock against HTM's CUDA kernels running on the default
+#       stream, and the failure mode at capture-time is a silent hang
+#       (100% GPU util, no log output, process state R).
+#   "0" — fall back to eager Python (slower, ~43k tps vs ~63k compiled).
+#       Keeps an escape hatch in case a future torch/inductor regression
+#       reintroduces a deadlock.
+#
+# Defensive .clone() on stacked_grads before in-place lerp_ eliminates the
+# alias-analysis edge case where inductor sees `g is stacked_grads` and
+# subsequent `stacked_grads.square()` operating on the post-lerp storage.
+# ---------------------------------------------------------------------------
+_MUON_COMPILE = os.environ.get("HYDRA_MUON_COMPILE", "1") == "1"
+
+def _maybe_compile(fn):
+    if _MUON_COMPILE:
+        # mode="default" explicitly opts OUT of cudagraphs (which reduce-overhead
+        # would enable) to avoid stream-capture deadlocks against HTM's CUDA
+        # kernels. dynamic=True minimizes recompile count across param-group
+        # shapes.
+        return torch.compile(fn, fullgraph=False, dynamic=True, mode="default")
+    return fn
+
+@_maybe_compile
+def muon_step_fused(stacked_grads, stacked_params, momentum_buffer, second_momentum_buffer,
+                    momentum_t, lr_t, wd_t, beta2_t, ns_steps, red_dim):
+    # Cast grads to param dtype AND clone defensively to break any alias
+    # between the (freshly-stacked) input and the in-place lerp_ below.
+    # Without this, inductor's alias analysis can emit code that reads from
+    # post-mutation storage when computing `v_mean = g.square().mean(...)`.
+    stacked_grads = stacked_grads.to(momentum_buffer.dtype).clone()
+    # Nesterov momentum
+    momentum = momentum_t.to(device=momentum_buffer.device, dtype=stacked_grads.dtype)
+    momentum_buffer.lerp_(stacked_grads, 1 - momentum)
+    g = stacked_grads.lerp_(momentum_buffer, momentum)
+    # Polar express orthogonalization
+    X = g.bfloat16()
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
+    if g.size(-2) > g.size(-1):
+        for a, b, c in polar_express_coeffs[:ns_steps]:
+            A = X.mT @ X
+            B = b * A + c * (A @ A)
+            X = a * X + X @ B
+    else:
+        for a, b, c in polar_express_coeffs[:ns_steps]:
+            A = X @ X.mT
+            B = b * A + c * (A @ A)
+            X = a * X + B @ X
+    g = X
+    # NorMuon variance reduction
+    # Keep beta2 in the state-buffer dtype, not g.dtype, so lerp_ on the
+    # float32 second_momentum_buffer doesn't hit a dtype mismatch on h200.
+    beta2 = beta2_t.to(device=second_momentum_buffer.device, dtype=second_momentum_buffer.dtype)
+    v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
+    red_dim_size = g.size(red_dim)
+    v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * red_dim_size
+    v_norm = v_norm_sq.sqrt()
+    second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2)
+    step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()
+    scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square()
+    v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt()
+    final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10))
+    g = g * final_scale.to(g.dtype)
+    # Cautious weight decay + parameter update
+    lr = lr_t.to(device=stacked_params.device, dtype=g.dtype)
+    wd = wd_t.to(device=stacked_params.device, dtype=g.dtype)
+    mask = (g * stacked_params) >= 0
+    stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)
+
+
+class MuonAdamW(torch.optim.Optimizer):
+    """Combined optimizer: Muon for 2D matrix params, AdamW for others."""
+
+    def __init__(self, param_groups):
+        super().__init__(param_groups, defaults={})
+        # 0-D CPU tensors to avoid torch.compile recompilation when values change
+        self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+
+    def _step_adamw(self, group):
+        params, grads, exp_avgs, exp_avg_sqs, state_steps = [], [], [], [], []
+        for p in group['params']:
+            if p.grad is None:
+                continue
+            state = self.state[p]
+            if not state:
+                state['step'] = 0
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
+            if 'step_t' not in state:
+                # _fused_adamw_ wants a per-param float step tensor on-device.
+                state['step_t'] = torch.tensor(
+                    float(state['step']), dtype=torch.float32, device=p.device
+                )
+            state['step'] += 1
+            params.append(p)
+            grads.append(p.grad.to(p.dtype) if p.grad.dtype != p.dtype else p.grad)
+            exp_avgs.append(state['exp_avg'])
+            exp_avg_sqs.append(state['exp_avg_sq'])
+            state_steps.append(state['step_t'])
+
+        if not params:
+            return
+
+        if _HYDRA_FUSED_ADAMW and _HAS_FUSED_ADAMW and params[0].is_cuda:
+            # _fused_adamw_ needs uniform (device, dtype) within a call, so
+            # group by (device, dtype) — same pattern as PyTorch's own
+            # AdamW(fused=True) path (_group_tensors_by_device_and_dtype).
+            buckets = {}
+            for p, g, ea, es, st in zip(params, grads, exp_avgs, exp_avg_sqs, state_steps):
+                key = (p.device, p.dtype)
+                buckets.setdefault(key, ([], [], [], [], []))
+                b_p, b_g, b_ea, b_es, b_st = buckets[key]
+                b_p.append(p); b_g.append(g); b_ea.append(ea); b_es.append(es); b_st.append(st)
+
+            lr_f = float(group['lr'])
+            b1_f = float(group['betas'][0])
+            b2_f = float(group['betas'][1])
+            wd_f = float(group['weight_decay'])
+            eps_f = float(group['eps'])
+            for (_dev, _dt), (b_p, b_g, b_ea, b_es, b_st) in buckets.items():
+                torch._foreach_add_(b_st, 1.0)
+                torch._fused_adamw_(
+                    b_p, b_g, b_ea, b_es,
+                    [],  # max_exp_avg_sqs unused (amsgrad=False)
+                    b_st,
+                    amsgrad=False,
+                    lr=lr_f, beta1=b1_f, beta2=b2_f,
+                    weight_decay=wd_f, eps=eps_f,
+                    maximize=False,
+                    grad_scale=None, found_inf=None,
+                )
+            return
+
+        # Fallback per-param path.
+        self._adamw_lr_t.fill_(group['lr'])
+        self._adamw_beta1_t.fill_(group['betas'][0])
+        self._adamw_beta2_t.fill_(group['betas'][1])
+        self._adamw_eps_t.fill_(group['eps'])
+        self._adamw_wd_t.fill_(group['weight_decay'])
+        for p, grad, exp_avg, exp_avg_sq in zip(params, grads, exp_avgs, exp_avg_sqs):
+            self._adamw_step_t.fill_(self.state[p]['step'])
+            adamw_step_fused(p, grad, exp_avg, exp_avg_sq,
+                             self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t,
+                             self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t)
+
+    def _step_muon(self, group):
+        params = [p for p in group['params'] if p.grad is not None]
+        if not params:
+            return
+        p = params[0]
+        state = self.state[p]
+        num_params = len(params)
+        shape, device, dtype = p.shape, p.device, p.dtype
+        if "momentum_buffer" not in state:
+            state["momentum_buffer"] = torch.zeros(num_params, *shape, dtype=dtype, device=device)
+        red_dim = -1 if shape[-2] >= shape[-1] else -2
+        if "second_momentum_buffer" not in state:
+            # Shape must match v_mean = stacked_grads.square().mean(dim=red_dim, keepdim=True)
+            full_shape = (num_params, *shape)
+            state_shape = list(full_shape)
+            state_shape[len(state_shape) + red_dim] = 1  # red_dim is negative
+            state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
+        # F7 REVERT: fresh stacks each step (no persistent stacked_params_buf).
+        # This was the autograd-safety fix that unblocks grad_accum>=2.
+        stacked_grads = torch.stack([p.grad for p in params])
+        stacked_params = torch.stack(params)
+        self._muon_momentum_t.fill_(group["momentum"])
+        self._muon_beta2_t.fill_(group["beta2"] if group["beta2"] is not None else 0.0)
+        self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1]) ** 0.5)
+        self._muon_wd_t.fill_(group["weight_decay"])
+        muon_step_fused(stacked_grads, stacked_params,
+                        state["momentum_buffer"], state["second_momentum_buffer"],
+                        self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t,
+                        self._muon_beta2_t, group["ns_steps"], red_dim)
+        torch._foreach_copy_(params, list(stacked_params.unbind(0)))
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            if group['kind'] == 'adamw':
+                self._step_adamw(group)
+            elif group['kind'] == 'muon':
+                self._step_muon(group)

overlay/hydra/training.py

@@ -1,961 +1,961 @@
-"""HYDRA training entry: setup, train loop, eval, summary.
-
-Extracted from the monolithic train.py (W1 modularization). Semantics
-preserved. Public entrypoint: `main()`.
-"""
-
-from __future__ import annotations
-
-import gc
-import json
-import math
-import os
-import sys
-import threading
-import time
-from dataclasses import asdict
-from pathlib import Path
-
-import torch
-
-# Line-buffered stdout so `python -u train.py | tee run.log | grep step` is
-# live (no \r overwrite, no 4k block-buffered pipe stalls). Safe on Python
-# 3.7+ where io.TextIOWrapper.reconfigure exists.
-try:
-    sys.stdout.reconfigure(line_buffering=True)  # type: ignore[attr-defined]
-except Exception:
-    pass
-
-from hydra.config import (
-    ADAM_BETAS, CURRICULUM_SHORT_SEQ_LEN, CURRICULUM_SHORT_STEPS,
-    D_MODEL, D_STATE, DEVICE_BATCH_SIZE, EMA_DECAY, EMBEDDING_LR,
-    ENGRAM_KEY_DIM, ENGRAM_LAYER_IDX, ENGRAM_N_COLUMNS, EXPAND,
-    FINAL_LR_FRAC, GPU_BF16_PEAK_FLOPS, HEADDIM, MATRIX_LR, N_HEADS,
-    N_LAYER, PostSemClawConfig, SCALAR_LR, SEED, TOTAL_BATCH_SIZE,
-    UNEMBEDDING_LR, USE_EMA, WARMUP_RATIO, WEIGHT_DECAY,
-)
-from hydra.diffusion_loss import mdlm_masked_forward_process, mdlm_rb_loss
-from hydra.eval import run_factual_english, run_factual_probes
-from hydra.model import PostSemClawModel
-
-import prepare as _prepare_mod
-from prepare import MAX_SEQ_LEN, TIME_BUDGET as _TIME_BUDGET, Tokenizer, evaluate_bpb as _evaluate_bpb_shards, get_token_bytes, make_dataloader as _make_dataloader_shards
-
-# Streaming Nemotron path (Super3 recipe). Opt-in via HYDRA_USE_NEMOTRON=1.
-if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
-    import prepare_nemotron as _p_nemo
-    make_dataloader = _p_nemo.make_dataloader
-    evaluate_bpb = _p_nemo.evaluate_bpb
-else:
-    make_dataloader = _make_dataloader_shards
-    evaluate_bpb = _evaluate_bpb_shards
-
-TIME_BUDGET = int(os.environ.get("HYDRA_TIME_BUDGET", str(_TIME_BUDGET)))
-_prepare_mod.TIME_BUDGET = TIME_BUDGET  # sync for evaluate_bpb
-
-CACHE_DIR = Path.home() / ".cache" / "autoresearch"
-LATEST_CKPT = CACHE_DIR / "latest.pt"
-PRETRAIN_FINAL_CKPT = CACHE_DIR / "pretrain_final.pt"
-FAILED_CKPT = CACHE_DIR / "latest_failed.pt"          # crash/FAIL path — never overwrites good
-BEST_CKPT = CACHE_DIR / "best_bpb.pt"                 # lowest val_bpb seen
-CKPT_INTERVAL = int(os.environ.get("HYDRA_CKPT_INTERVAL", "250"))
-CKPT_ROTATIONS = int(os.environ.get("HYDRA_CKPT_ROTATIONS", "3"))  # how many .N backups to keep
-RESUME_CKPT = os.environ.get("HYDRA_RESUME_CKPT", str(LATEST_CKPT))
-
-# MDLM (Masked Diffusion LM) Rao-Blackwellized ELBO loss path.
-#   HYDRA_USE_MDLM=1         : switch training loss from AR sampled-softmax CE
-#                              to MDLM RB weighted CE (arXiv:2406.07524).
-#   HYDRA_MDLM_MASK_ID=N     : token id used for the MASK sentinel (default:
-#                              last valid id, vocab_size - 1). Ensure this id
-#                              never appears in training targets — typical
-#                              practice is to reserve it.
-#   HYDRA_MDLM_SCHEDULE=loglinear|linear  : noise schedule (default loglinear).
-# When enabled, the per-step flow is:
-#   1. mdlm_masked_forward_process(y)  ->  (x_noised, mask_positions, weights)
-#   2. logits = model(x_noised)                          (no targets -> full V logits)
-#   3. loss = mdlm_rb_loss(logits, y, mask_positions, weights)
-# Sampled-softmax is bypassed in this path because the RB ELBO needs
-# full-vocab logits on masked positions.
-USE_MDLM = os.environ.get("HYDRA_USE_MDLM", "0") == "1"
-MDLM_MASK_ID = int(os.environ.get("HYDRA_MDLM_MASK_ID", "-1"))  # -1 => default to vocab_size-1 at runtime
-MDLM_SCHEDULE = os.environ.get("HYDRA_MDLM_SCHEDULE", "loglinear")
-
-
-# ---------------------------------------------------------------------------
-# Schedules
-# ---------------------------------------------------------------------------
-
-def get_lr_multiplier(progress: float) -> float:
-    if progress < WARMUP_RATIO:
-        return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
-    decay_progress = (progress - WARMUP_RATIO) / (1.0 - WARMUP_RATIO)
-    return FINAL_LR_FRAC + 0.5 * (1.0 - FINAL_LR_FRAC) * (1 + math.cos(math.pi * decay_progress))
-
-
-def get_muon_momentum(step: int) -> float:
-    frac = min(step / 300, 1)
-    return (1 - frac) * 0.85 + frac * 0.95
-
-
-def get_weight_decay(progress: float) -> float:
-    return WEIGHT_DECAY * (1 - progress)
-
-
-_CKPT_WORKER_THREAD: threading.Thread | None = None
-
-
-def _ckpt_snapshot_state_dicts(
-    model: PostSemClawModel,
-    optimizer: torch.optim.Optimizer,
-) -> tuple[dict, dict]:
-    """Detach + CPU-clone every tensor so a bg thread can serialize safely
-    while the main loop keeps mutating live weights/optimizer state."""
-    msd = {k: (v.detach().to("cpu", copy=True) if torch.is_tensor(v) else v)
-           for k, v in model.state_dict().items()}
-    # optimizer.state_dict() is a nested dict; walk it.
-    osd_raw = optimizer.state_dict()
-
-    def _to_cpu(obj):
-        if torch.is_tensor(obj):
-            return obj.detach().to("cpu", copy=True)
-        if isinstance(obj, dict):
-            return {k: _to_cpu(v) for k, v in obj.items()}
-        if isinstance(obj, list):
-            return [_to_cpu(v) for v in obj]
-        if isinstance(obj, tuple):
-            return tuple(_to_cpu(v) for v in obj)
-        return obj
-
-    osd = _to_cpu(osd_raw)
-    return msd, osd
-
-
-def save_ckpt(
-    model: PostSemClawModel,
-    optimizer: torch.optim.Optimizer,
-    config: PostSemClawConfig,
-    step: int,
-    total_training_time: float,
-    smooth_train_loss: float,
-    bpt_ema: float,
-    epoch: int,
-    path: Path,
-    *,
-    val_bpb: float | None = None,
-    blocking: bool = False,
-) -> None:
-    """Save a training checkpoint.
-
-    Default behavior is async: the GPU→CPU state_dict clone runs on the main
-    thread (unavoidable; needs to happen before the next optimizer.step that
-    mutates live weights), then `torch.save` is dispatched to a daemon
-    worker thread. The next call joins any still-running prior save so only
-    one disk write is in flight.
-
-    `blocking=True` restores the original synchronous behavior — used for
-    end-of-training saves where correctness on process exit matters.
-    """
-    global _CKPT_WORKER_THREAD
-    try:
-        CACHE_DIR.mkdir(parents=True, exist_ok=True)
-        msd, osd = _ckpt_snapshot_state_dicts(model, optimizer)
-        # asdict() recursively converts dataclass fields to a dict and
-        # renders tuples as lists. hyena_layers therefore round-trips as a
-        # JSON-safe list; config_from_dict normalizes it back to a tuple.
-        payload = {
-            "model_state_dict": msd,
-            "optimizer_state_dict": osd,
-            "config": asdict(config),
-            "step": step,
-            "epoch": epoch,
-            "train_seconds": total_training_time,
-            "smoothed_loss": smooth_train_loss,
-            "bpt_ema": bpt_ema,
-            "val_bpb": val_bpb,
-        }
-        path_str = str(path)
-
-        def _rotate(p: str) -> None:
-            """Keep up to CKPT_ROTATIONS previous versions as p.1, p.2, ..."""
-            if CKPT_ROTATIONS <= 0:
-                return
-            try:
-                # Walk from oldest to newest so we don't clobber newer with older.
-                for i in range(CKPT_ROTATIONS, 0, -1):
-                    src = f"{p}.{i-1}" if i > 1 else p
-                    dst = f"{p}.{i}"
-                    if os.path.exists(src):
-                        os.replace(src, dst)
-            except Exception as e:
-                # Rotation is best-effort; never block a save on it.
-                print(f"[ckpt] rotate warn {p}: {type(e).__name__}: {e}", flush=True)
-
-        def _write():
-            try:
-                _rotate(path_str)
-                tmp = path_str + ".tmp"
-                torch.save(payload, tmp)
-                os.replace(tmp, path_str)
-                print(f"[ckpt] saved {path_str} (step={step})", flush=True)
-            except Exception as e:
-                print(f"[ckpt] SAVE FAILED {path_str}: {type(e).__name__}: {e}", flush=True)
-
-        if blocking:
-            _write()
-            return
-
-        # Join previous writer so at most one torch.save runs at a time.
-        if _CKPT_WORKER_THREAD is not None and _CKPT_WORKER_THREAD.is_alive():
-            _CKPT_WORKER_THREAD.join()
-        _CKPT_WORKER_THREAD = threading.Thread(
-            target=_write, daemon=True, name=f"ckpt-save-{step}"
-        )
-        _CKPT_WORKER_THREAD.start()
-    except Exception as e:
-        print(f"[ckpt] SNAPSHOT FAILED {path}: {type(e).__name__}: {e}", flush=True)
-
-
-def config_from_dict(cfg_dict: dict) -> PostSemClawConfig:
-    """Reconstruct a PostSemClawConfig from a checkpoint's asdict() payload.
-
-    Newly-added fields (e.g. `hyena_layers`) are defaulted when absent in
-    older checkpoints, and list-ified tuples are coerced back to tuples so
-    the dataclass keeps its declared types.
-
-    This is the ckpt-safe inverse of `asdict(config)` used by save_ckpt and
-    guarantees that a resume path can rebuild the exact same model topology
-    (Mamba3 vs HyenaBlock per layer) regardless of env-var state at resume.
-    """
-    # Only keep keys that are actually declared on PostSemClawConfig — extra
-    # keys in older/newer checkpoints must not crash construction.
-    field_names = {f.name for f in PostSemClawConfig.__dataclass_fields__.values()}
-    filtered = {k: v for k, v in cfg_dict.items() if k in field_names}
-    # asdict renders tuple[int,...] as list[int]; coerce back so the model
-    # builder sees the declared type.
-    if "hyena_layers" in filtered and filtered["hyena_layers"] is not None:
-        filtered["hyena_layers"] = tuple(sorted(int(x) for x in filtered["hyena_layers"]))
-    return PostSemClawConfig(**filtered)
-
-
-def _try_load_ckpt(path: Path, model, optimizer, device):
-    """Attempt to load a single ckpt. Returns the tuple on success, None on any failure."""
-    if not path.exists():
-        return None
-    ckpt = torch.load(str(path), map_location=device, weights_only=False)
-    state = ckpt.get("model_state_dict", ckpt)
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing:
-        print(f"[ckpt] {path.name} missing={len(missing)}", flush=True)
-    if unexpected:
-        print(f"[ckpt] {path.name} unexpected={len(unexpected)}", flush=True)
-    optimizer_state = ckpt.get("optimizer_state_dict")
-    if optimizer_state is not None:
-        try:
-            optimizer.load_state_dict(optimizer_state)
-        except Exception as e:
-            print(f"[ckpt] optimizer restore failed from {path.name}: {type(e).__name__}: {e}", flush=True)
-    step = int(ckpt.get("step", 0))
-    total_training_time = float(ckpt.get("train_seconds", 0.0))
-    smooth_train_loss = float(ckpt.get("smoothed_loss", 0.0))
-    bpt_ema = float(ckpt.get("bpt_ema", 0.0))
-    epoch = int(ckpt.get("epoch", 0))
-    print(
-        f"[ckpt] resumed {path} step={step} train_seconds={total_training_time:.1f}",
-        flush=True,
-    )
-    # Warn if resuming a schedule-exhausted ckpt — user is probably warm-starting.
-    budget = float(os.environ.get("HYDRA_TIME_BUDGET", "0") or 0)
-    if budget and total_training_time >= 0.99 * budget:
-        print(
-            f"[ckpt] WARNING: resumed ckpt used {total_training_time:.0f}s of {budget:.0f}s "
-            f"budget. LR schedule is essentially exhausted. "
-            f"Set HYDRA_WARMSTART=1 to reset optimizer + scheduler and keep only weights.",
-            flush=True,
-        )
-    return step, total_training_time, smooth_train_loss, bpt_ema, epoch
-
-
-def maybe_resume_ckpt(
-    model: PostSemClawModel,
-    optimizer: torch.optim.Optimizer,
-    device: torch.device,
-) -> tuple[int, float, float, float, int]:
-    if not RESUME_CKPT or RESUME_CKPT.lower() == "none":
-        print("[ckpt] resume disabled; starting fresh", flush=True)
-        return 0, 0.0, 0.0, 0.0, 0
-
-    resume_path = Path(os.path.expanduser(RESUME_CKPT))
-    # Try the primary path, then rotated backups. This is crucial because a
-    # partial / killed torch.save on the primary path would leave a corrupt
-    # file. If that fails we fall back to latest.pt.1, .2, .3 automatically.
-    candidates: list[Path] = [resume_path]
-    for i in range(1, CKPT_ROTATIONS + 1):
-        candidates.append(Path(str(resume_path) + f".{i}"))
-
-    for cand in candidates:
-        if not cand.exists():
-            continue
-        try:
-            result = _try_load_ckpt(cand, model, optimizer, device)
-            if result is not None:
-                if cand != resume_path:
-                    print(f"[ckpt] fell back to rotation {cand.name}", flush=True)
-                return result
-        except Exception as e:
-            print(f"[ckpt] {cand.name} load failed: {type(e).__name__}: {e}", flush=True)
-            continue
-
-    print(f"[ckpt] no usable checkpoint in {resume_path} + rotations; starting fresh", flush=True)
-    return 0, 0.0, 0.0, 0.0, 0
-
-
-# ---------------------------------------------------------------------------
-# Main entry
-# ---------------------------------------------------------------------------
-
-def main() -> None:
-    t_start = time.time()
-    torch.manual_seed(SEED)
-    torch.cuda.manual_seed(SEED)
-    # Precision / kernel-selection knobs for peak throughput on Ampere.
-    # - high : matmul uses TF32 (Ampere's 10-bit mantissa accum) for fp32 ops
-    # - allow_tf32 : explicit for both matmul + cudnn paths
-    # - cudnn.benchmark : env-gated (HYDRA_CUDNN_BENCHMARK, default OFF).
-    #   TRUE can lock in a locally-better-but-globally-slower algorithm
-    #   after the autotune phase ends, causing tps to degrade 15-20%
-    #   over the first ~100 steps. Observed 2026-04-22 and confirmed by
-    #   differential profiling. Default is now FALSE; set =1 only if you
-    #   see a specific workload where benchmark helps sustained tps.
-    torch.set_float32_matmul_precision("high")
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    torch.backends.cudnn.benchmark = os.environ.get("HYDRA_CUDNN_BENCHMARK", "0") == "1"
-    device = torch.device("cuda")
-    autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
-
-    # Streaming path skips prepare.py (which normally trains the tokenizer
-    # and builds the retina), so we must materialize both before model init.
-    if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
-        _p_nemo.ensure_tokenizer()
-        # Retina: HF Hub cache hit for this (vocab, n_bits, target_active) combo
-        # returns in seconds; otherwise build_retina streams Nemotron docs to
-        # compute cooccurrence + train SOM, then uploads back to the cache.
-        import subsystems.sdr_retina as _sdr_retina
-        _sdr_retina.build_retina()
-    tokenizer = Tokenizer.from_directory()
-    vocab_size = tokenizer.get_vocab_size()
-    print(f"Vocab size: {vocab_size:,}")
-
-    config = PostSemClawConfig(
-        sequence_len=MAX_SEQ_LEN,
-        vocab_size=vocab_size,
-        n_layer=N_LAYER,
-        d_model=D_MODEL,
-        d_state=D_STATE,
-        headdim=HEADDIM,
-        n_heads=N_HEADS,
-        expand=EXPAND,
-        engram_n_columns=ENGRAM_N_COLUMNS,
-        engram_key_dim=ENGRAM_KEY_DIM,
-        engram_layer_idx=ENGRAM_LAYER_IDX,
-    )
-    print(f"Model config: {asdict(config)}")
-
-    with torch.device("meta"):
-        model = PostSemClawModel(config)
-    model.to_empty(device=device)
-    model.init_weights()
-
-    param_counts = model.num_scaling_params()
-    print("Parameter counts:")
-    for key, value in param_counts.items():
-        print(f"  {key:24s}: {value:,}")
-    num_params = param_counts['total']
-    num_flops_per_token = model.estimate_flops()
-    print(f"Estimated FLOPs per token: {num_flops_per_token:e}")
-
-    tokens_per_fwdbwd = DEVICE_BATCH_SIZE * MAX_SEQ_LEN
-    assert TOTAL_BATCH_SIZE % tokens_per_fwdbwd == 0
-    grad_accum_steps = TOTAL_BATCH_SIZE // tokens_per_fwdbwd
-
-    optimizer = model.setup_optimizer(
-        unembedding_lr=UNEMBEDDING_LR,
-        embedding_lr=EMBEDDING_LR,
-        scalar_lr=SCALAR_LR,
-        adam_betas=ADAM_BETAS,
-        matrix_lr=MATRIX_LR,
-        weight_decay=WEIGHT_DECAY,
-    )
-
-    step, total_training_time, smooth_train_loss, bpt_ema, resume_epoch = maybe_resume_ckpt(
-        model, optimizer, device,
-    )
-
-    # Learnability #4: inform the model of the BOS token id so it can mask
-    # doc-separator positions in packed sequences. Always set (the mask only
-    # fires when HYDRA_DOC_SEP_MASK=1 is also on).
-    if hasattr(model, 'set_bos_token_id'):
-        model.set_bos_token_id(tokenizer.get_bos_token_id())
-
-    # Learnability #2: EMA shadow copy of weights. AveragedModel clones every
-    # parameter; we update it after every optimizer step and save it at the
-    # end alongside the raw checkpoint. Defaults OFF.
-    ema_model = None
-    if USE_EMA:
-        try:
-            from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn
-            # decay=EMA_DECAY; avg_fn uses get_ema_multi_avg_fn for numerical
-            # stability across bf16/fp32 mixed parameter groups.
-            ema_model = AveragedModel(
-                model,
-                multi_avg_fn=get_ema_multi_avg_fn(EMA_DECAY),
-            )
-            print(f"[EMA] enabled with decay={EMA_DECAY}")
-        except Exception as _e:
-            print(f"[EMA] disabled — AveragedModel init failed: {_e}")
-            ema_model = None
-
-    print("torch.compile: Muon step compiled; AdamW uses torch._fused_adamw_ (model blocks use native CUDA kernels)")
-
-    # Learnability #7: curriculum short-then-long. If enabled, build the
-    # initial dataloader at the short seq_len; we swap to full MAX_SEQ_LEN
-    # after CURRICULUM_SHORT_STEPS optimizer steps (see loop below).
-    _curriculum_active = CURRICULUM_SHORT_STEPS > 0 and CURRICULUM_SHORT_SEQ_LEN < MAX_SEQ_LEN
-    _current_seq_len = CURRICULUM_SHORT_SEQ_LEN if _curriculum_active else MAX_SEQ_LEN
-    if _curriculum_active:
-        print(
-            f"[CURRICULUM] starting at T={_current_seq_len} for "
-            f"{CURRICULUM_SHORT_STEPS} steps, then switching to T={MAX_SEQ_LEN}"
-        )
-    train_loader = make_dataloader(tokenizer, DEVICE_BATCH_SIZE, _current_seq_len, "train")
-    x, y, epoch = next(train_loader)  # prefetch first batch
-    if resume_epoch > 0:
-        epoch = max(epoch, resume_epoch)
-
-    print(f"Time budget: {TIME_BUDGET}s")
-    print(f"Gradient accumulation steps: {grad_accum_steps}")
-
-    # Token→byte LUT for bits-per-byte computation. evaluate_bpb in prepare.py
-    # uses total_nats / (ln(2) * total_bytes); our live metric needs to match.
-    # Without this, `bpb = loss/ln(2)` is actually bits-per-TOKEN, which at
-    # vocab=8192 scales by ~4 and makes live train bpb non-comparable with
-    # val_bpb (champion 1.279 bpb vs train printing "8.04").
-    token_bytes = get_token_bytes(device=device)
-
-    # -----------------------------------------------------------------------
-    # Training loop
-    # -----------------------------------------------------------------------
-
-    t_start_training = time.time()
-
-    # Async postprocessing — run SOM + Hestia on background threads so
-    # the GPU doesn't idle during their CPU-bound work.
-    _ASYNC_POSTPROCESS = os.environ.get("HYDRA_ASYNC_POSTPROCESS", "1") == "1"
-    _som_thread: threading.Thread | None = None
-    _hestia_thread: threading.Thread | None = None
-    _hestia_stream: torch.cuda.Stream | None = (
-        torch.cuda.Stream() if _ASYNC_POSTPROCESS else None
-    )
-
-    # HYDRA_PROFILE_STEPS=N prints a per-phase cpu/gpu time breakdown for the
-    # first N steps (and every 100th step thereafter if N<0). Zero overhead
-    # when disabled. Used to find what's eating CPU budget when GPU should
-    # be the bottleneck.
-    _profile_steps = int(os.environ.get("HYDRA_PROFILE_STEPS", "0"))
-
-    while True:
-        torch.cuda.synchronize()
-        t0 = time.time()
-        _prof = _profile_steps and (step < _profile_steps or (_profile_steps < 0 and step % 100 == 0))
-        _gpu_ms = 0.0
-        _data_ms = 0.0
-        for micro_step in range(grad_accum_steps):
-            if _prof:
-                torch.cuda.synchronize(); _t_micro = time.time()
-            if USE_MDLM:
-                # MDLM path: corrupt y -> x_noised, run model to get full-V logits,
-                # compute RB weighted CE on masked positions. x (original input) is
-                # unused in this path — the model only sees the noised version of y.
-                _mask_id = MDLM_MASK_ID if MDLM_MASK_ID >= 0 else (vocab_size - 1)
-                x_noised, mask_positions, loss_weights = mdlm_masked_forward_process(
-                    y, mask_token_id=_mask_id, alpha_schedule=MDLM_SCHEDULE,
-                )
-                with autocast_ctx:
-                    logits = model(x_noised)  # targets=None -> (B, T, V) logits
-                loss = mdlm_rb_loss(logits, y, mask_positions, loss_weights)
-            else:
-                with autocast_ctx:
-                    loss = model(x, y)
-            train_loss = loss.detach()
-            loss = loss / grad_accum_steps
-            loss.backward()
-            if _prof:
-                torch.cuda.synchronize()
-                _gpu_ms += (time.time() - _t_micro) * 1000
-                _t_data = time.time()
-            x, y, epoch = next(train_loader)
-            if _prof:
-                _data_ms += (time.time() - _t_data) * 1000
-        if _prof:
-            torch.cuda.synchronize(); _t_fb = time.time()
-
-        # Progress and schedules
-        progress = min(total_training_time / TIME_BUDGET, 1.0)
-        lrm = get_lr_multiplier(progress)
-        muon_momentum = get_muon_momentum(step)
-        muon_weight_decay = get_weight_decay(progress)
-        for group in optimizer.param_groups:
-            group["lr"] = group["initial_lr"] * lrm
-            if group['kind'] == 'muon':
-                group["momentum"] = muon_momentum
-                group["weight_decay"] = muon_weight_decay
-        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-        optimizer.step()
-        if _prof:
-            torch.cuda.synchronize(); _t_opt = time.time()
-
-        # Learnability #2: EMA update after every optimizer step.
-        if ema_model is not None:
-            try:
-                ema_model.update_parameters(model)
-            except Exception as _e:
-                print(f"[EMA] update failed at step {step}: {_e}", flush=True)
-
-        # Learnability #7: curriculum transition. After
-        # CURRICULUM_SHORT_STEPS optimizer steps, rebuild the dataloader at
-        # MAX_SEQ_LEN. Done once, then the flag flips off.
-        if _curriculum_active and step + 1 >= CURRICULUM_SHORT_STEPS:
-            print(
-                f"[CURRICULUM] step={step+1} — switching from T={_current_seq_len} "
-                f"to T={MAX_SEQ_LEN}",
-                flush=True,
-            )
-            _current_seq_len = MAX_SEQ_LEN
-            _curriculum_active = False
-            train_loader = make_dataloader(tokenizer, DEVICE_BATCH_SIZE, _current_seq_len, "train")
-            # Prefetch the next batch at the new seq_len so the following
-            # loop iteration consumes fresh data.
-            x, y, epoch = next(train_loader)
-
-        # Online SOM update — retina is now a plain Python attribute (not a
-        # registered buffer) so mutations do not invalidate torch.compile guards.
-        # Runs fully on CPU; safe to overlap with GPU forward pass.
-        _last_sdr = getattr(model, "_last_sdr", None)
-        if _last_sdr is not None:
-            if _ASYNC_POSTPROCESS:
-                if _som_thread is not None:
-                    _som_thread.join()
-                # Clone tensors before next step overwrites them
-                _som_x = x.clone()
-                _som_sdr = _last_sdr.clone()
-                _som_thread = threading.Thread(
-                    target=model.sdr_semantic.maybe_som_update,
-                    args=(_som_x, _som_sdr),
-                    daemon=True,
-                )
-                _som_thread.start()
-            else:
-                model.sdr_semantic.maybe_som_update(x, _last_sdr)
-
-        # Hestia QAT — anneal temperature every step, snap every N steps.
-        # apply_to walks all Linear modules (CPU) then does .data.copy_ (GPU).
-        # Background thread + separate CUDA stream lets this overlap with
-        # the next forward pass on the default stream.
-        _hestia_progress = (time.time() - t_start_training) / max(TIME_BUDGET, 1)
-        _hestia_interval = int(os.environ.get("HYDRA_HESTIA_INTERVAL", "100"))
-        if step % _hestia_interval == 0:
-            if _ASYNC_POSTPROCESS:
-                if _hestia_thread is not None:
-                    _hestia_thread.join()
-
-                def _hestia_bg(mdl: torch.nn.Module, prog: float) -> None:
-                    assert _hestia_stream is not None
-                    with torch.cuda.stream(_hestia_stream):
-                        mdl.hestia.anneal_temperature(prog)
-                        mdl.hestia.apply_to(mdl)
-
-                _hestia_thread = threading.Thread(
-                    target=_hestia_bg,
-                    args=(model, _hestia_progress),
-                    daemon=True,
-                )
-                _hestia_thread.start()
-            else:
-                model.hestia.anneal_temperature(_hestia_progress)
-                model.hestia.apply_to(model)
-        else:
-            # anneal_temperature is cheap (~1 us), keep inline
-            model.hestia.anneal_temperature(_hestia_progress)
-
-        model.zero_grad(set_to_none=True)
-
-        train_loss_f = train_loss.item()
-        if math.isnan(train_loss_f) or train_loss_f > 100:
-            print("FAIL")
-            # Save to a DIFFERENT file — never clobber a good latest.pt with
-            # a NaN/diverged state. The good ckpt from the last periodic save
-            # is the right place to resume from.
-            save_ckpt(
-                model,
-                optimizer,
-                config,
-                step,
-                total_training_time,
-                smooth_train_loss,
-                bpt_ema,
-                epoch,
-                FAILED_CKPT,
-                blocking=True,
-            )
-            raise SystemExit(1)
-
-        torch.cuda.synchronize()
-        t1 = time.time()
-        dt = t1 - t0
-
-        if _prof:
-            fb = (_t_fb - t0) * 1000
-            opt = (_t_opt - _t_fb) * 1000
-            rest = (t1 - _t_opt) * 1000
-            print(
-                f"[PROF step={step:05d}] gpu={_gpu_ms:.0f}ms data_fetch={_data_ms:.0f}ms "
-                f"(sum_fb={fb:.0f}) opt={opt:.0f}ms rest={rest:.0f}ms total={dt*1000:.0f}ms",
-                flush=True,
-            )
-
-        if step > 10:
-            total_training_time += dt
-
-        ema_beta = 0.9
-        smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss_f
-        debiased_smooth_loss = smooth_train_loss / (1 - ema_beta ** (step + 1))
-        pct_done = 100 * progress
-        tok_per_sec = int(TOTAL_BATCH_SIZE / dt)
-        mfu = 100 * num_flops_per_token * TOTAL_BATCH_SIZE / dt / GPU_BF16_PEAK_FLOPS
-        remaining = max(0, TIME_BUDGET - total_training_time)
-
-        # Bytes-per-token for the CURRENT batch. evaluate_bpb in prepare.py
-        # computes bits-per-BYTE (total_nats / (ln2 * total_bytes)); to match
-        # that semantics live, we EMA-smooth the per-batch bytes/token and
-        # divide. Without this, the old `bpb = loss/ln2` was actually
-        # bits-per-token — ~4× larger than val_bpb at vocab=8192 and
-        # therefore not comparable to the champion 1.279 bpb metric.
-        with torch.no_grad():
-            y_flat = y.view(-1)
-            nbytes_batch = token_bytes[y_flat]
-            mask = nbytes_batch > 0
-            mask_count = mask.sum().clamp(min=1).float()
-            avg_bytes_per_tok = (nbytes_batch.float() * mask.float()).sum() / mask_count
-            bpt_batch = float(avg_bytes_per_tok.item())
-        if step == 0 or bpt_ema <= 0.0:
-            bpt_ema = bpt_batch
-        else:
-            bpt_ema = 0.98 * bpt_ema + 0.02 * bpt_batch
-
-        # Dual metric: bpb (byte-normalized, comparable with val_bpb) AND
-        # bpt (bits per token, the raw loss in bits). bpt_div exposes the
-        # current avg bytes-per-token so the conversion is transparent.
-        bpt = debiased_smooth_loss / math.log(2)
-        bpb = bpt / max(bpt_ema, 1e-6)
-        vram_mib = torch.cuda.memory_allocated() / 1024 / 1024
-        current_lr = optimizer.param_groups[0]["lr"]
-
-        # Per-step line-buffered log. NOT \r-overwritten so tee/grep see it.
-        # Keep key=value pairs grep-friendly.
-        ppl = 2.0 ** bpb  # perplexity (byte-level)
-        print(
-            f"step={step:05d} loss={debiased_smooth_loss:.4f} bpb={bpb:.4f} ppl={ppl:.3f} "
-            f"bpt={bpt:.3f} bpt_div={bpt_ema:.2f} "
-            f"tps={tok_per_sec} dt_ms={dt*1000:.0f} mfu={mfu:.1f} "
-            f"lr={current_lr:.2e} vram={vram_mib:.0f}MiB "
-            f"pct={pct_done:.1f} epoch={epoch} remaining={remaining:.0f}s",
-            flush=True,
-        )
-
-        if step == 0:
-            gc.collect()
-            gc.freeze()
-            gc.disable()
-        # No periodic gc.collect() — we disabled+froze at step 0 on purpose,
-        # so a manual collect every 5k steps just re-scans frozen objects
-        # (burned ~900 ms/event in production) for no live-garbage reason.
-
-        if CKPT_INTERVAL > 0 and step > 0 and step % CKPT_INTERVAL == 0:
-            save_ckpt(
-                model,
-                optimizer,
-                config,
-                step,
-                total_training_time,
-                smooth_train_loss,
-                bpt_ema,
-                epoch,
-                LATEST_CKPT,
-            )
-
-        # Periodic mid-training validation so we can see the model learning
-        # English in real time (not just at the end). Small val batch so it
-        # doesn't eat significant training time.
-        mid_val_interval = int(os.environ.get("HYDRA_MID_VAL_INTERVAL", "500"))
-        if mid_val_interval > 0 and step > 0 and step % mid_val_interval == 0:
-            model.eval()
-            try:
-                # Defrag GPU memory before eval allocates fresh chunks —
-                # without this the eval path can OOM on 6GB cards even
-                # though total usage fits, because the allocator's free
-                # blocks are fragmented.
-                torch.cuda.empty_cache()
-                _orig_mid = _prepare_mod.EVAL_TOKENS
-                _prepare_mod.EVAL_TOKENS = 262144  # ~260K tokens, fast
-                with torch.no_grad():
-                    with autocast_ctx:
-                        mid_bpb = evaluate_bpb(model, tokenizer, DEVICE_BATCH_SIZE)
-                _prepare_mod.EVAL_TOKENS = _orig_mid
-                mid_ppl = 2.0 ** mid_bpb
-                print(f"[MID_VAL] step={step} val_bpb={mid_bpb:.4f} val_ppl={mid_ppl:.3f}", flush=True)
-
-                # Per-layer diagnostic panel. Only printed when HYDRA_LAYER_DIAGNOSTICS=1
-                # is set (otherwise the layer_* keys are absent from _metrics).
-                _diag_metrics = model.get_secondary_metrics()
-                _layer_keys = sorted([k for k in _diag_metrics.keys() if k.startswith('layer_')])
-                if _layer_keys:
-                    # Condense: one row per layer showing the four core signals.
-                    n_layers = len(model.blocks)
-                    print(f"[LAYER_DIAG] step={step}", flush=True)
-                    for li in range(n_layers):
-                        d_ratio = _diag_metrics.get(f'layer_{li}_delta_ratio', float('nan'))
-                        out_n   = _diag_metrics.get(f'layer_{li}_out_norm',    float('nan'))
-                        g_norm  = _diag_metrics.get(f'layer_{li}_grad_norm',   float('nan'))
-                        eff_r   = _diag_metrics.get(f'layer_{li}_eff_rank',    float('nan'))
-                        f_std   = _diag_metrics.get(f'layer_{li}_feat_std',    float('nan'))
-                        print(
-                            f"[LAYER_DIAG]   L{li:02d}  delta_ratio={d_ratio:.4f}  "
-                            f"out_norm={out_n:.4f}  grad_norm={g_norm:.3e}  "
-                            f"eff_rank={eff_r:.1f}  feat_std={f_std:.4f}",
-                            flush=True,
-                        )
-                    htm_proj_g = _diag_metrics.get('htm_proj_grad_norm', None)
-                    if htm_proj_g is not None:
-                        print(f"[LAYER_DIAG]   htm_proj grad_norm={htm_proj_g:.3e}", flush=True)
-            except Exception as e:
-                print(f"[MID_VAL] failed: {e}", flush=True)
-            model.train()
-
-        step += 1
-
-        if step > 10 and total_training_time >= TIME_BUDGET:
-            break
-
-    # Drain async postprocessing threads before eval
-    if _som_thread is not None:
-        _som_thread.join()
-    if _hestia_thread is not None:
-        _hestia_thread.join()
-    if _hestia_stream is not None:
-        _hestia_stream.synchronize()
-
-    total_tokens = step * TOTAL_BATCH_SIZE
-
-    # ----------------------------------------------------------------------
-    # SAVE ORDER (critical):
-    #   1. Save PRETRAIN_FINAL_CKPT with val_bpb=None  (hedge against eval OOM)
-    #   2. Save LATEST_CKPT with val_bpb=None          (hedge against eval OOM)
-    #   3. Run eval (may OOM on small GPUs; we survive it)
-    #   4. Re-save both ckpts with val_bpb filled in
-    # This way we NEVER lose the final trained weights to an eval crash.
-    # Previous ordering put eval first, so an eval-time OOM destroyed the
-    # only record of a 6h training run (2026-04-22 incident).
-    # ----------------------------------------------------------------------
-
-    save_ckpt(
-        model, optimizer, config, step, total_training_time,
-        smooth_train_loss, bpt_ema, epoch, PRETRAIN_FINAL_CKPT,
-        val_bpb=None, blocking=True,
-    )
-    save_ckpt(
-        model, optimizer, config, step, total_training_time,
-        smooth_train_loss, bpt_ema, epoch, LATEST_CKPT,
-        val_bpb=None, blocking=True,
-    )
-
-    # Now it's safe to eval — ckpts are on disk regardless of what happens here.
-    # HYDRA_EVAL_BATCH overrides DEVICE_BATCH_SIZE (env-tunable; default halves
-    # the training batch because eval holds activations for full sequence and
-    # does not benefit from overlap with backward). HYDRA_EVAL_TOKENS controls
-    # how many val tokens to sweep (default 2 M, short enough for autoresearch
-    # 5-min budgets).
-    val_bpb: float | None = None
-    # Eval batch: default to 4 on cloud GPUs (enough freed VRAM after optimizer
-    # clear), fall back to DEVICE_BATCH_SIZE//2 on tiny cards. Env-overridable.
-    _eval_B = int(os.environ.get("HYDRA_EVAL_BATCH",
-        str(max(1, DEVICE_BATCH_SIZE // 2) if DEVICE_BATCH_SIZE <= 8 else 4)))
-    # Eval tokens: default 1M (1,048,576) — gives statistically meaningful BPB
-    # (256 forward passes at B=4, seq=1024). Env-overridable for fast/slow sweeps.
-    _eval_tokens = int(os.environ.get("HYDRA_EVAL_TOKENS", str(1048576)))
-    try:
-        # Aggressive VRAM reclaim for 6GB cards. Peak training VRAM = 5.1GB
-        # which leaves < 1GB for the eval forward — the driver can't satisfy
-        # the allocation. Free EVERY tensor we don't strictly need:
-        #   - optimizer grads (set_to_none releases tensor)
-        #   - optimizer.state (fp32 Muon NS workspace, AdamW moments — ~size-of-params each)
-        #   - model internal caches (HTM subsample cache, SDR stash)
-        # After this, VRAM should be ~params only (bf16 ≈ 120MB at 60M params).
-        optimizer.zero_grad(set_to_none=True)
-        if hasattr(optimizer, 'state') and optimizer.state:
-            for p, st in list(optimizer.state.items()):
-                st.clear()
-            optimizer.state.clear()
-        for p in model.parameters():
-            if p.grad is not None:
-                p.grad = None
-        if hasattr(model, '_htm_cache'):
-            model._htm_cache = None
-        if hasattr(model, '_last_sdr'):
-            model._last_sdr = None
-        import gc as _gc
-        _gc.collect()
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-        try:
-            _free_mb = torch.cuda.mem_get_info()[0] / 1024 / 1024
-            print(f"[VAL] free_vram_mb={_free_mb:.0f} (cleared optimizer state)", flush=True)
-        except Exception:
-            pass
-        print(f"[VAL] running eval on {_eval_tokens} tokens at B={_eval_B}...", flush=True)
-        model.eval()
-        _orig = _prepare_mod.EVAL_TOKENS
-        _prepare_mod.EVAL_TOKENS = _eval_tokens
-        # Nemotron path reads HYDRA_STREAM_EVAL_TOKENS env var directly,
-        # not _prepare_mod.EVAL_TOKENS. Sync both so eval budget is
-        # respected regardless of which dataloader path is active.
-        _orig_stream = os.environ.get("HYDRA_STREAM_EVAL_TOKENS")
-        os.environ["HYDRA_STREAM_EVAL_TOKENS"] = str(_eval_tokens)
-        with autocast_ctx:
-            val_bpb = evaluate_bpb(model, tokenizer, _eval_B)
-        _prepare_mod.EVAL_TOKENS = _orig
-        if _orig_stream is not None:
-            os.environ["HYDRA_STREAM_EVAL_TOKENS"] = _orig_stream
-        else:
-            os.environ.pop("HYDRA_STREAM_EVAL_TOKENS", None)
-        val_ppl = 2 ** val_bpb
-        print(f"[VAL] step={step} val_bpb={val_bpb:.4f} val_ppl={val_ppl:.3f}", flush=True)
-    except torch.cuda.OutOfMemoryError as e:
-        print(f"[VAL] SKIPPED (OOM): {e}", flush=True)
-        torch.cuda.empty_cache()
-    except Exception as e:
-        import traceback as _tb
-        print(f"[VAL] SKIPPED ({type(e).__name__}): {e}", flush=True)
-        _tb.print_exc()
-        try:
-            _free = torch.cuda.mem_get_info()[0] / 1024 / 1024
-            print(f"[VAL] post-crash free_vram_mb={_free:.0f}", flush=True)
-        except Exception:
-            pass
-
-    # Final ckpts with val_bpb filled in (if eval succeeded).
-    save_ckpt(
-        model, optimizer, config, step, total_training_time,
-        smooth_train_loss, bpt_ema, epoch, LATEST_CKPT,
-        val_bpb=val_bpb, blocking=True,
-    )
-    save_ckpt(
-        model, optimizer, config, step, total_training_time,
-        smooth_train_loss, bpt_ema, epoch, PRETRAIN_FINAL_CKPT,
-        val_bpb=val_bpb, blocking=True,
-    )
-
-    # Learnability #2: persist EMA weights alongside the raw checkpoint.
-    # latest_ema.pt contains ema_model.module (the Averaged params) so it
-    # can be loaded by evaluation / inference code that expects the same
-    # state_dict shape as the raw model.
-    if ema_model is not None:
-        try:
-            ema_ckpt_path = CACHE_DIR / "latest_ema.pt"
-            CACHE_DIR.mkdir(parents=True, exist_ok=True)
-            torch.save({
-                "model_state_dict": ema_model.module.state_dict(),
-                "config": asdict(config),
-                "step": step,
-                "epoch": epoch,
-                "train_seconds": total_training_time,
-                "val_bpb": val_bpb,
-                "ema_decay": EMA_DECAY,
-            }, str(ema_ckpt_path))
-            print(f"[EMA] saved {ema_ckpt_path} (step={step})", flush=True)
-        except Exception as _e:
-            print(f"[EMA] save failed: {_e}", flush=True)
-
-    run_factual_probes(model, tokenizer, device, autocast_ctx)
-
-    t_end = time.time()
-    startup_time = t_start_training - t_start
-    steady_state_mfu = (
-        100 * num_flops_per_token * TOTAL_BATCH_SIZE * (step - 10)
-        / total_training_time / GPU_BF16_PEAK_FLOPS
-        if total_training_time > 0 else 0
-    )
-    peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
-    metrics = model.get_secondary_metrics()
-
-    print("---")
-    print(f"val_bpb:          {val_bpb:.6f}" if val_bpb is not None else "val_bpb:          SKIPPED")
-    print(f"training_seconds: {total_training_time:.1f}")
-    print(f"total_seconds:    {t_end - t_start:.1f}")
-    print(f"peak_vram_mb:     {peak_vram_mb:.1f}")
-    print(f"mfu_percent:      {steady_state_mfu:.2f}")
-    print(f"total_tokens_M:   {total_tokens / 1e6:.1f}")
-    print(f"num_steps:        {step}")
-    print(f"num_params_M:     {num_params / 1e6:.1f}")
-    print(f"n_layer:          {N_LAYER}")
-    print(f"d_model:          {D_MODEL}")
-    print(f"engram_hit_rate:   {metrics.get('engram_hit_rate', 0.0):.4f}")
-    print(f"sdr_active_bits:  {metrics.get('sdr_active_bits', 0):.1f}")
-    print(f"htm_anomaly:      {metrics.get('htm_anomaly', 0):.4f}")
-
-    # Per-layer summary panel — only printed when diagnostics were active.
-    _layer_keys = sorted([k for k in metrics.keys() if k.startswith('layer_')])
-    if _layer_keys:
-        n_layers = len(model.blocks)
-        print("--- per-layer diagnostic panel ---")
-        for li in range(n_layers):
-            d_ratio = metrics.get(f'layer_{li}_delta_ratio', float('nan'))
-            out_n   = metrics.get(f'layer_{li}_out_norm',    float('nan'))
-            g_norm  = metrics.get(f'layer_{li}_grad_norm',   float('nan'))
-            eff_r   = metrics.get(f'layer_{li}_eff_rank',    float('nan'))
-            f_std   = metrics.get(f'layer_{li}_feat_std',    float('nan'))
-            print(
-                f"L{li:02d}  delta_ratio={d_ratio:.4f}  out_norm={out_n:.4f}  "
-                f"grad_norm={g_norm:.3e}  eff_rank={eff_r:.1f}  feat_std={f_std:.4f}"
-            )
-
-    # Emit full metrics dictionary as JSON for sweep aggregation. Path from
-    # HYDRA_METRICS_OUT env var; default=/tmp/hydra_run_metrics.json. Always
-    # written (even without diagnostics) so the aggregator can compare runs.
-    _metrics_out = os.environ.get("HYDRA_METRICS_OUT", "/tmp/hydra_run_metrics.json")
-    try:
-        _dump = dict(metrics)
-        _dump.update({
-            'val_bpb': (float(val_bpb) if val_bpb is not None else None),
-            'val_ppl': (float(val_ppl) if val_ppl is not None else None),
-            'n_layer': int(N_LAYER),
-            'd_model': int(D_MODEL),
-            'num_params_M': float(num_params / 1e6),
-            'num_steps': int(step),
-            'total_tokens_M': float(total_tokens / 1e6),
-            'peak_vram_mb': float(peak_vram_mb),
-            'training_seconds': float(total_training_time),
-            'sdr_target_active': int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327")),
-        })
-        Path(_metrics_out).parent.mkdir(parents=True, exist_ok=True)
-        with open(_metrics_out, 'w') as _f:
-            json.dump(_dump, _f, indent=2, sort_keys=True)
-        print(f"[METRICS] wrote {_metrics_out}", flush=True)
-        # Also emit a single-line JSON to stdout so the sweep aggregator can
-        # scrape it from HF Jobs logs without pulling files out of the container.
-        print("[METRICS_JSON] " + json.dumps(_dump, sort_keys=True), flush=True)
-    except Exception as _e:
-        print(f"[METRICS] write failed: {_e}", flush=True)
-
-    run_factual_english(model, tokenizer, MAX_SEQ_LEN)
-    # startup_time is informative but not printed (preserve historical output)
-    _ = startup_time
+"""HYDRA training entry: setup, train loop, eval, summary.
+
+Extracted from the monolithic train.py (W1 modularization). Semantics
+preserved. Public entrypoint: `main()`.
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import math
+import os
+import sys
+import threading
+import time
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+
+# Line-buffered stdout so `python -u train.py | tee run.log | grep step` is
+# live (no \r overwrite, no 4k block-buffered pipe stalls). Safe on Python
+# 3.7+ where io.TextIOWrapper.reconfigure exists.
+try:
+    sys.stdout.reconfigure(line_buffering=True)  # type: ignore[attr-defined]
+except Exception:
+    pass
+
+from hydra.config import (
+    ADAM_BETAS, CURRICULUM_SHORT_SEQ_LEN, CURRICULUM_SHORT_STEPS,
+    D_MODEL, D_STATE, DEVICE_BATCH_SIZE, EMA_DECAY, EMBEDDING_LR,
+    ENGRAM_KEY_DIM, ENGRAM_LAYER_IDX, ENGRAM_N_COLUMNS, EXPAND,
+    FINAL_LR_FRAC, GPU_BF16_PEAK_FLOPS, HEADDIM, MATRIX_LR, N_HEADS,
+    N_LAYER, PostSemClawConfig, SCALAR_LR, SEED, TOTAL_BATCH_SIZE,
+    UNEMBEDDING_LR, USE_EMA, WARMUP_RATIO, WEIGHT_DECAY,
+)
+from hydra.diffusion_loss import mdlm_masked_forward_process, mdlm_rb_loss
+from hydra.eval import run_factual_english, run_factual_probes
+from hydra.model import PostSemClawModel
+
+import prepare as _prepare_mod
+from prepare import MAX_SEQ_LEN, TIME_BUDGET as _TIME_BUDGET, Tokenizer, evaluate_bpb as _evaluate_bpb_shards, get_token_bytes, make_dataloader as _make_dataloader_shards
+
+# Streaming Nemotron path (Super3 recipe). Opt-in via HYDRA_USE_NEMOTRON=1.
+if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
+    import prepare_nemotron as _p_nemo
+    make_dataloader = _p_nemo.make_dataloader
+    evaluate_bpb = _p_nemo.evaluate_bpb
+else:
+    make_dataloader = _make_dataloader_shards
+    evaluate_bpb = _evaluate_bpb_shards
+
+TIME_BUDGET = int(os.environ.get("HYDRA_TIME_BUDGET", str(_TIME_BUDGET)))
+_prepare_mod.TIME_BUDGET = TIME_BUDGET  # sync for evaluate_bpb
+
+CACHE_DIR = Path.home() / ".cache" / "autoresearch"
+LATEST_CKPT = CACHE_DIR / "latest.pt"
+PRETRAIN_FINAL_CKPT = CACHE_DIR / "pretrain_final.pt"
+FAILED_CKPT = CACHE_DIR / "latest_failed.pt"          # crash/FAIL path — never overwrites good
+BEST_CKPT = CACHE_DIR / "best_bpb.pt"                 # lowest val_bpb seen
+CKPT_INTERVAL = int(os.environ.get("HYDRA_CKPT_INTERVAL", "250"))
+CKPT_ROTATIONS = int(os.environ.get("HYDRA_CKPT_ROTATIONS", "3"))  # how many .N backups to keep
+RESUME_CKPT = os.environ.get("HYDRA_RESUME_CKPT", str(LATEST_CKPT))
+
+# MDLM (Masked Diffusion LM) Rao-Blackwellized ELBO loss path.
+#   HYDRA_USE_MDLM=1         : switch training loss from AR sampled-softmax CE
+#                              to MDLM RB weighted CE (arXiv:2406.07524).
+#   HYDRA_MDLM_MASK_ID=N     : token id used for the MASK sentinel (default:
+#                              last valid id, vocab_size - 1). Ensure this id
+#                              never appears in training targets — typical
+#                              practice is to reserve it.
+#   HYDRA_MDLM_SCHEDULE=loglinear|linear  : noise schedule (default loglinear).
+# When enabled, the per-step flow is:
+#   1. mdlm_masked_forward_process(y)  ->  (x_noised, mask_positions, weights)
+#   2. logits = model(x_noised)                          (no targets -> full V logits)
+#   3. loss = mdlm_rb_loss(logits, y, mask_positions, weights)
+# Sampled-softmax is bypassed in this path because the RB ELBO needs
+# full-vocab logits on masked positions.
+USE_MDLM = os.environ.get("HYDRA_USE_MDLM", "0") == "1"
+MDLM_MASK_ID = int(os.environ.get("HYDRA_MDLM_MASK_ID", "-1"))  # -1 => default to vocab_size-1 at runtime
+MDLM_SCHEDULE = os.environ.get("HYDRA_MDLM_SCHEDULE", "loglinear")
+
+
+# ---------------------------------------------------------------------------
+# Schedules
+# ---------------------------------------------------------------------------
+
+def get_lr_multiplier(progress: float) -> float:
+    if progress < WARMUP_RATIO:
+        return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
+    decay_progress = (progress - WARMUP_RATIO) / (1.0 - WARMUP_RATIO)
+    return FINAL_LR_FRAC + 0.5 * (1.0 - FINAL_LR_FRAC) * (1 + math.cos(math.pi * decay_progress))
+
+
+def get_muon_momentum(step: int) -> float:
+    frac = min(step / 300, 1)
+    return (1 - frac) * 0.85 + frac * 0.95
+
+
+def get_weight_decay(progress: float) -> float:
+    return WEIGHT_DECAY * (1 - progress)
+
+
+_CKPT_WORKER_THREAD: threading.Thread | None = None
+
+
+def _ckpt_snapshot_state_dicts(
+    model: PostSemClawModel,
+    optimizer: torch.optim.Optimizer,
+) -> tuple[dict, dict]:
+    """Detach + CPU-clone every tensor so a bg thread can serialize safely
+    while the main loop keeps mutating live weights/optimizer state."""
+    msd = {k: (v.detach().to("cpu", copy=True) if torch.is_tensor(v) else v)
+           for k, v in model.state_dict().items()}
+    # optimizer.state_dict() is a nested dict; walk it.
+    osd_raw = optimizer.state_dict()
+
+    def _to_cpu(obj):
+        if torch.is_tensor(obj):
+            return obj.detach().to("cpu", copy=True)
+        if isinstance(obj, dict):
+            return {k: _to_cpu(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [_to_cpu(v) for v in obj]
+        if isinstance(obj, tuple):
+            return tuple(_to_cpu(v) for v in obj)
+        return obj
+
+    osd = _to_cpu(osd_raw)
+    return msd, osd
+
+
+def save_ckpt(
+    model: PostSemClawModel,
+    optimizer: torch.optim.Optimizer,
+    config: PostSemClawConfig,
+    step: int,
+    total_training_time: float,
+    smooth_train_loss: float,
+    bpt_ema: float,
+    epoch: int,
+    path: Path,
+    *,
+    val_bpb: float | None = None,
+    blocking: bool = False,
+) -> None:
+    """Save a training checkpoint.
+
+    Default behavior is async: the GPU→CPU state_dict clone runs on the main
+    thread (unavoidable; needs to happen before the next optimizer.step that
+    mutates live weights), then `torch.save` is dispatched to a daemon
+    worker thread. The next call joins any still-running prior save so only
+    one disk write is in flight.
+
+    `blocking=True` restores the original synchronous behavior — used for
+    end-of-training saves where correctness on process exit matters.
+    """
+    global _CKPT_WORKER_THREAD
+    try:
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        msd, osd = _ckpt_snapshot_state_dicts(model, optimizer)
+        # asdict() recursively converts dataclass fields to a dict and
+        # renders tuples as lists. hyena_layers therefore round-trips as a
+        # JSON-safe list; config_from_dict normalizes it back to a tuple.
+        payload = {
+            "model_state_dict": msd,
+            "optimizer_state_dict": osd,
+            "config": asdict(config),
+            "step": step,
+            "epoch": epoch,
+            "train_seconds": total_training_time,
+            "smoothed_loss": smooth_train_loss,
+            "bpt_ema": bpt_ema,
+            "val_bpb": val_bpb,
+        }
+        path_str = str(path)
+
+        def _rotate(p: str) -> None:
+            """Keep up to CKPT_ROTATIONS previous versions as p.1, p.2, ..."""
+            if CKPT_ROTATIONS <= 0:
+                return
+            try:
+                # Walk from oldest to newest so we don't clobber newer with older.
+                for i in range(CKPT_ROTATIONS, 0, -1):
+                    src = f"{p}.{i-1}" if i > 1 else p
+                    dst = f"{p}.{i}"
+                    if os.path.exists(src):
+                        os.replace(src, dst)
+            except Exception as e:
+                # Rotation is best-effort; never block a save on it.
+                print(f"[ckpt] rotate warn {p}: {type(e).__name__}: {e}", flush=True)
+
+        def _write():
+            try:
+                _rotate(path_str)
+                tmp = path_str + ".tmp"
+                torch.save(payload, tmp)
+                os.replace(tmp, path_str)
+                print(f"[ckpt] saved {path_str} (step={step})", flush=True)
+            except Exception as e:
+                print(f"[ckpt] SAVE FAILED {path_str}: {type(e).__name__}: {e}", flush=True)
+
+        if blocking:
+            _write()
+            return
+
+        # Join previous writer so at most one torch.save runs at a time.
+        if _CKPT_WORKER_THREAD is not None and _CKPT_WORKER_THREAD.is_alive():
+            _CKPT_WORKER_THREAD.join()
+        _CKPT_WORKER_THREAD = threading.Thread(
+            target=_write, daemon=True, name=f"ckpt-save-{step}"
+        )
+        _CKPT_WORKER_THREAD.start()
+    except Exception as e:
+        print(f"[ckpt] SNAPSHOT FAILED {path}: {type(e).__name__}: {e}", flush=True)
+
+
+def config_from_dict(cfg_dict: dict) -> PostSemClawConfig:
+    """Reconstruct a PostSemClawConfig from a checkpoint's asdict() payload.
+
+    Newly-added fields (e.g. `hyena_layers`) are defaulted when absent in
+    older checkpoints, and list-ified tuples are coerced back to tuples so
+    the dataclass keeps its declared types.
+
+    This is the ckpt-safe inverse of `asdict(config)` used by save_ckpt and
+    guarantees that a resume path can rebuild the exact same model topology
+    (Mamba3 vs HyenaBlock per layer) regardless of env-var state at resume.
+    """
+    # Only keep keys that are actually declared on PostSemClawConfig — extra
+    # keys in older/newer checkpoints must not crash construction.
+    field_names = {f.name for f in PostSemClawConfig.__dataclass_fields__.values()}
+    filtered = {k: v for k, v in cfg_dict.items() if k in field_names}
+    # asdict renders tuple[int,...] as list[int]; coerce back so the model
+    # builder sees the declared type.
+    if "hyena_layers" in filtered and filtered["hyena_layers"] is not None:
+        filtered["hyena_layers"] = tuple(sorted(int(x) for x in filtered["hyena_layers"]))
+    return PostSemClawConfig(**filtered)
+
+
+def _try_load_ckpt(path: Path, model, optimizer, device):
+    """Attempt to load a single ckpt. Returns the tuple on success, None on any failure."""
+    if not path.exists():
+        return None
+    ckpt = torch.load(str(path), map_location=device, weights_only=False)
+    state = ckpt.get("model_state_dict", ckpt)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing:
+        print(f"[ckpt] {path.name} missing={len(missing)}", flush=True)
+    if unexpected:
+        print(f"[ckpt] {path.name} unexpected={len(unexpected)}", flush=True)
+    optimizer_state = ckpt.get("optimizer_state_dict")
+    if optimizer_state is not None:
+        try:
+            optimizer.load_state_dict(optimizer_state)
+        except Exception as e:
+            print(f"[ckpt] optimizer restore failed from {path.name}: {type(e).__name__}: {e}", flush=True)
+    step = int(ckpt.get("step", 0))
+    total_training_time = float(ckpt.get("train_seconds", 0.0))
+    smooth_train_loss = float(ckpt.get("smoothed_loss", 0.0))
+    bpt_ema = float(ckpt.get("bpt_ema", 0.0))
+    epoch = int(ckpt.get("epoch", 0))
+    print(
+        f"[ckpt] resumed {path} step={step} train_seconds={total_training_time:.1f}",
+        flush=True,
+    )
+    # Warn if resuming a schedule-exhausted ckpt — user is probably warm-starting.
+    budget = float(os.environ.get("HYDRA_TIME_BUDGET", "0") or 0)
+    if budget and total_training_time >= 0.99 * budget:
+        print(
+            f"[ckpt] WARNING: resumed ckpt used {total_training_time:.0f}s of {budget:.0f}s "
+            f"budget. LR schedule is essentially exhausted. "
+            f"Set HYDRA_WARMSTART=1 to reset optimizer + scheduler and keep only weights.",
+            flush=True,
+        )
+    return step, total_training_time, smooth_train_loss, bpt_ema, epoch
+
+
+def maybe_resume_ckpt(
+    model: PostSemClawModel,
+    optimizer: torch.optim.Optimizer,
+    device: torch.device,
+) -> tuple[int, float, float, float, int]:
+    if not RESUME_CKPT or RESUME_CKPT.lower() == "none":
+        print("[ckpt] resume disabled; starting fresh", flush=True)
+        return 0, 0.0, 0.0, 0.0, 0
+
+    resume_path = Path(os.path.expanduser(RESUME_CKPT))
+    # Try the primary path, then rotated backups. This is crucial because a
+    # partial / killed torch.save on the primary path would leave a corrupt
+    # file. If that fails we fall back to latest.pt.1, .2, .3 automatically.
+    candidates: list[Path] = [resume_path]
+    for i in range(1, CKPT_ROTATIONS + 1):
+        candidates.append(Path(str(resume_path) + f".{i}"))
+
+    for cand in candidates:
+        if not cand.exists():
+            continue
+        try:
+            result = _try_load_ckpt(cand, model, optimizer, device)
+            if result is not None:
+                if cand != resume_path:
+                    print(f"[ckpt] fell back to rotation {cand.name}", flush=True)
+                return result
+        except Exception as e:
+            print(f"[ckpt] {cand.name} load failed: {type(e).__name__}: {e}", flush=True)
+            continue
+
+    print(f"[ckpt] no usable checkpoint in {resume_path} + rotations; starting fresh", flush=True)
+    return 0, 0.0, 0.0, 0.0, 0
+
+
+# ---------------------------------------------------------------------------
+# Main entry
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    t_start = time.time()
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    # Precision / kernel-selection knobs for peak throughput on Ampere.
+    # - high : matmul uses TF32 (Ampere's 10-bit mantissa accum) for fp32 ops
+    # - allow_tf32 : explicit for both matmul + cudnn paths
+    # - cudnn.benchmark : env-gated (HYDRA_CUDNN_BENCHMARK, default OFF).
+    #   TRUE can lock in a locally-better-but-globally-slower algorithm
+    #   after the autotune phase ends, causing tps to degrade 15-20%
+    #   over the first ~100 steps. Observed 2026-04-22 and confirmed by
+    #   differential profiling. Default is now FALSE; set =1 only if you
+    #   see a specific workload where benchmark helps sustained tps.
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = os.environ.get("HYDRA_CUDNN_BENCHMARK", "0") == "1"
+    device = torch.device("cuda")
+    autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+
+    # Streaming path skips prepare.py (which normally trains the tokenizer
+    # and builds the retina), so we must materialize both before model init.
+    if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
+        _p_nemo.ensure_tokenizer()
+        # Retina: HF Hub cache hit for this (vocab, n_bits, target_active) combo
+        # returns in seconds; otherwise build_retina streams Nemotron docs to
+        # compute cooccurrence + train SOM, then uploads back to the cache.
+        import subsystems.sdr_retina as _sdr_retina
+        _sdr_retina.build_retina()
+    tokenizer = Tokenizer.from_directory()
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Vocab size: {vocab_size:,}")
+
+    config = PostSemClawConfig(
+        sequence_len=MAX_SEQ_LEN,
+        vocab_size=vocab_size,
+        n_layer=N_LAYER,
+        d_model=D_MODEL,
+        d_state=D_STATE,
+        headdim=HEADDIM,
+        n_heads=N_HEADS,
+        expand=EXPAND,
+        engram_n_columns=ENGRAM_N_COLUMNS,
+        engram_key_dim=ENGRAM_KEY_DIM,
+        engram_layer_idx=ENGRAM_LAYER_IDX,
+    )
+    print(f"Model config: {asdict(config)}")
+
+    with torch.device("meta"):
+        model = PostSemClawModel(config)
+    model.to_empty(device=device)
+    model.init_weights()
+
+    param_counts = model.num_scaling_params()
+    print("Parameter counts:")
+    for key, value in param_counts.items():
+        print(f"  {key:24s}: {value:,}")
+    num_params = param_counts['total']
+    num_flops_per_token = model.estimate_flops()
+    print(f"Estimated FLOPs per token: {num_flops_per_token:e}")
+
+    tokens_per_fwdbwd = DEVICE_BATCH_SIZE * MAX_SEQ_LEN
+    assert TOTAL_BATCH_SIZE % tokens_per_fwdbwd == 0
+    grad_accum_steps = TOTAL_BATCH_SIZE // tokens_per_fwdbwd
+
+    optimizer = model.setup_optimizer(
+        unembedding_lr=UNEMBEDDING_LR,
+        embedding_lr=EMBEDDING_LR,
+        scalar_lr=SCALAR_LR,
+        adam_betas=ADAM_BETAS,
+        matrix_lr=MATRIX_LR,
+        weight_decay=WEIGHT_DECAY,
+    )
+
+    step, total_training_time, smooth_train_loss, bpt_ema, resume_epoch = maybe_resume_ckpt(
+        model, optimizer, device,
+    )
+
+    # Learnability #4: inform the model of the BOS token id so it can mask
+    # doc-separator positions in packed sequences. Always set (the mask only
+    # fires when HYDRA_DOC_SEP_MASK=1 is also on).
+    if hasattr(model, 'set_bos_token_id'):
+        model.set_bos_token_id(tokenizer.get_bos_token_id())
+
+    # Learnability #2: EMA shadow copy of weights. AveragedModel clones every
+    # parameter; we update it after every optimizer step and save it at the
+    # end alongside the raw checkpoint. Defaults OFF.
+    ema_model = None
+    if USE_EMA:
+        try:
+            from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn
+            # decay=EMA_DECAY; avg_fn uses get_ema_multi_avg_fn for numerical
+            # stability across bf16/fp32 mixed parameter groups.
+            ema_model = AveragedModel(
+                model,
+                multi_avg_fn=get_ema_multi_avg_fn(EMA_DECAY),
+            )
+            print(f"[EMA] enabled with decay={EMA_DECAY}")
+        except Exception as _e:
+            print(f"[EMA] disabled — AveragedModel init failed: {_e}")
+            ema_model = None
+
+    print("torch.compile: Muon step compiled; AdamW uses torch._fused_adamw_ (model blocks use native CUDA kernels)")
+
+    # Learnability #7: curriculum short-then-long. If enabled, build the
+    # initial dataloader at the short seq_len; we swap to full MAX_SEQ_LEN
+    # after CURRICULUM_SHORT_STEPS optimizer steps (see loop below).
+    _curriculum_active = CURRICULUM_SHORT_STEPS > 0 and CURRICULUM_SHORT_SEQ_LEN < MAX_SEQ_LEN
+    _current_seq_len = CURRICULUM_SHORT_SEQ_LEN if _curriculum_active else MAX_SEQ_LEN
+    if _curriculum_active:
+        print(
+            f"[CURRICULUM] starting at T={_current_seq_len} for "
+            f"{CURRICULUM_SHORT_STEPS} steps, then switching to T={MAX_SEQ_LEN}"
+        )
+    train_loader = make_dataloader(tokenizer, DEVICE_BATCH_SIZE, _current_seq_len, "train")
+    x, y, epoch = next(train_loader)  # prefetch first batch
+    if resume_epoch > 0:
+        epoch = max(epoch, resume_epoch)
+
+    print(f"Time budget: {TIME_BUDGET}s")
+    print(f"Gradient accumulation steps: {grad_accum_steps}")
+
+    # Token→byte LUT for bits-per-byte computation. evaluate_bpb in prepare.py
+    # uses total_nats / (ln(2) * total_bytes); our live metric needs to match.
+    # Without this, `bpb = loss/ln(2)` is actually bits-per-TOKEN, which at
+    # vocab=8192 scales by ~4 and makes live train bpb non-comparable with
+    # val_bpb (champion 1.279 bpb vs train printing "8.04").
+    token_bytes = get_token_bytes(device=device)
+
+    # -----------------------------------------------------------------------
+    # Training loop
+    # -----------------------------------------------------------------------
+
+    t_start_training = time.time()
+
+    # Async postprocessing — run SOM + Hestia on background threads so
+    # the GPU doesn't idle during their CPU-bound work.
+    _ASYNC_POSTPROCESS = os.environ.get("HYDRA_ASYNC_POSTPROCESS", "1") == "1"
+    _som_thread: threading.Thread | None = None
+    _hestia_thread: threading.Thread | None = None
+    _hestia_stream: torch.cuda.Stream | None = (
+        torch.cuda.Stream() if _ASYNC_POSTPROCESS else None
+    )
+
+    # HYDRA_PROFILE_STEPS=N prints a per-phase cpu/gpu time breakdown for the
+    # first N steps (and every 100th step thereafter if N<0). Zero overhead
+    # when disabled. Used to find what's eating CPU budget when GPU should
+    # be the bottleneck.
+    _profile_steps = int(os.environ.get("HYDRA_PROFILE_STEPS", "0"))
+
+    while True:
+        torch.cuda.synchronize()
+        t0 = time.time()
+        _prof = _profile_steps and (step < _profile_steps or (_profile_steps < 0 and step % 100 == 0))
+        _gpu_ms = 0.0
+        _data_ms = 0.0
+        for micro_step in range(grad_accum_steps):
+            if _prof:
+                torch.cuda.synchronize(); _t_micro = time.time()
+            if USE_MDLM:
+                # MDLM path: corrupt y -> x_noised, run model to get full-V logits,
+                # compute RB weighted CE on masked positions. x (original input) is
+                # unused in this path — the model only sees the noised version of y.
+                _mask_id = MDLM_MASK_ID if MDLM_MASK_ID >= 0 else (vocab_size - 1)
+                x_noised, mask_positions, loss_weights = mdlm_masked_forward_process(
+                    y, mask_token_id=_mask_id, alpha_schedule=MDLM_SCHEDULE,
+                )
+                with autocast_ctx:
+                    logits = model(x_noised)  # targets=None -> (B, T, V) logits
+                loss = mdlm_rb_loss(logits, y, mask_positions, loss_weights)
+            else:
+                with autocast_ctx:
+                    loss = model(x, y)
+            train_loss = loss.detach()
+            loss = loss / grad_accum_steps
+            loss.backward()
+            if _prof:
+                torch.cuda.synchronize()
+                _gpu_ms += (time.time() - _t_micro) * 1000
+                _t_data = time.time()
+            x, y, epoch = next(train_loader)
+            if _prof:
+                _data_ms += (time.time() - _t_data) * 1000
+        if _prof:
+            torch.cuda.synchronize(); _t_fb = time.time()
+
+        # Progress and schedules
+        progress = min(total_training_time / TIME_BUDGET, 1.0)
+        lrm = get_lr_multiplier(progress)
+        muon_momentum = get_muon_momentum(step)
+        muon_weight_decay = get_weight_decay(progress)
+        for group in optimizer.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+            if group['kind'] == 'muon':
+                group["momentum"] = muon_momentum
+                group["weight_decay"] = muon_weight_decay
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        optimizer.step()
+        if _prof:
+            torch.cuda.synchronize(); _t_opt = time.time()
+
+        # Learnability #2: EMA update after every optimizer step.
+        if ema_model is not None:
+            try:
+                ema_model.update_parameters(model)
+            except Exception as _e:
+                print(f"[EMA] update failed at step {step}: {_e}", flush=True)
+
+        # Learnability #7: curriculum transition. After
+        # CURRICULUM_SHORT_STEPS optimizer steps, rebuild the dataloader at
+        # MAX_SEQ_LEN. Done once, then the flag flips off.
+        if _curriculum_active and step + 1 >= CURRICULUM_SHORT_STEPS:
+            print(
+                f"[CURRICULUM] step={step+1} — switching from T={_current_seq_len} "
+                f"to T={MAX_SEQ_LEN}",
+                flush=True,
+            )
+            _current_seq_len = MAX_SEQ_LEN
+            _curriculum_active = False
+            train_loader = make_dataloader(tokenizer, DEVICE_BATCH_SIZE, _current_seq_len, "train")
+            # Prefetch the next batch at the new seq_len so the following
+            # loop iteration consumes fresh data.
+            x, y, epoch = next(train_loader)
+
+        # Online SOM update — retina is now a plain Python attribute (not a
+        # registered buffer) so mutations do not invalidate torch.compile guards.
+        # Runs fully on CPU; safe to overlap with GPU forward pass.
+        _last_sdr = getattr(model, "_last_sdr", None)
+        if _last_sdr is not None:
+            if _ASYNC_POSTPROCESS:
+                if _som_thread is not None:
+                    _som_thread.join()
+                # Clone tensors before next step overwrites them
+                _som_x = x.clone()
+                _som_sdr = _last_sdr.clone()
+                _som_thread = threading.Thread(
+                    target=model.sdr_semantic.maybe_som_update,
+                    args=(_som_x, _som_sdr),
+                    daemon=True,
+                )
+                _som_thread.start()
+            else:
+                model.sdr_semantic.maybe_som_update(x, _last_sdr)
+
+        # Hestia QAT — anneal temperature every step, snap every N steps.
+        # apply_to walks all Linear modules (CPU) then does .data.copy_ (GPU).
+        # Background thread + separate CUDA stream lets this overlap with
+        # the next forward pass on the default stream.
+        _hestia_progress = (time.time() - t_start_training) / max(TIME_BUDGET, 1)
+        _hestia_interval = int(os.environ.get("HYDRA_HESTIA_INTERVAL", "100"))
+        if step % _hestia_interval == 0:
+            if _ASYNC_POSTPROCESS:
+                if _hestia_thread is not None:
+                    _hestia_thread.join()
+
+                def _hestia_bg(mdl: torch.nn.Module, prog: float) -> None:
+                    assert _hestia_stream is not None
+                    with torch.cuda.stream(_hestia_stream):
+                        mdl.hestia.anneal_temperature(prog)
+                        mdl.hestia.apply_to(mdl)
+
+                _hestia_thread = threading.Thread(
+                    target=_hestia_bg,
+                    args=(model, _hestia_progress),
+                    daemon=True,
+                )
+                _hestia_thread.start()
+            else:
+                model.hestia.anneal_temperature(_hestia_progress)
+                model.hestia.apply_to(model)
+        else:
+            # anneal_temperature is cheap (~1 us), keep inline
+            model.hestia.anneal_temperature(_hestia_progress)
+
+        model.zero_grad(set_to_none=True)
+
+        train_loss_f = train_loss.item()
+        if math.isnan(train_loss_f) or train_loss_f > 100:
+            print("FAIL")
+            # Save to a DIFFERENT file — never clobber a good latest.pt with
+            # a NaN/diverged state. The good ckpt from the last periodic save
+            # is the right place to resume from.
+            save_ckpt(
+                model,
+                optimizer,
+                config,
+                step,
+                total_training_time,
+                smooth_train_loss,
+                bpt_ema,
+                epoch,
+                FAILED_CKPT,
+                blocking=True,
+            )
+            raise SystemExit(1)
+
+        torch.cuda.synchronize()
+        t1 = time.time()
+        dt = t1 - t0
+
+        if _prof:
+            fb = (_t_fb - t0) * 1000
+            opt = (_t_opt - _t_fb) * 1000
+            rest = (t1 - _t_opt) * 1000
+            print(
+                f"[PROF step={step:05d}] gpu={_gpu_ms:.0f}ms data_fetch={_data_ms:.0f}ms "
+                f"(sum_fb={fb:.0f}) opt={opt:.0f}ms rest={rest:.0f}ms total={dt*1000:.0f}ms",
+                flush=True,
+            )
+
+        if step > 10:
+            total_training_time += dt
+
+        ema_beta = 0.9
+        smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss_f
+        debiased_smooth_loss = smooth_train_loss / (1 - ema_beta ** (step + 1))
+        pct_done = 100 * progress
+        tok_per_sec = int(TOTAL_BATCH_SIZE / dt)
+        mfu = 100 * num_flops_per_token * TOTAL_BATCH_SIZE / dt / GPU_BF16_PEAK_FLOPS
+        remaining = max(0, TIME_BUDGET - total_training_time)
+
+        # Bytes-per-token for the CURRENT batch. evaluate_bpb in prepare.py
+        # computes bits-per-BYTE (total_nats / (ln2 * total_bytes)); to match
+        # that semantics live, we EMA-smooth the per-batch bytes/token and
+        # divide. Without this, the old `bpb = loss/ln2` was actually
+        # bits-per-token — ~4× larger than val_bpb at vocab=8192 and
+        # therefore not comparable to the champion 1.279 bpb metric.
+        with torch.no_grad():
+            y_flat = y.view(-1)
+            nbytes_batch = token_bytes[y_flat]
+            mask = nbytes_batch > 0
+            mask_count = mask.sum().clamp(min=1).float()
+            avg_bytes_per_tok = (nbytes_batch.float() * mask.float()).sum() / mask_count
+            bpt_batch = float(avg_bytes_per_tok.item())
+        if step == 0 or bpt_ema <= 0.0:
+            bpt_ema = bpt_batch
+        else:
+            bpt_ema = 0.98 * bpt_ema + 0.02 * bpt_batch
+
+        # Dual metric: bpb (byte-normalized, comparable with val_bpb) AND
+        # bpt (bits per token, the raw loss in bits). bpt_div exposes the
+        # current avg bytes-per-token so the conversion is transparent.
+        bpt = debiased_smooth_loss / math.log(2)
+        bpb = bpt / max(bpt_ema, 1e-6)
+        vram_mib = torch.cuda.memory_allocated() / 1024 / 1024
+        current_lr = optimizer.param_groups[0]["lr"]
+
+        # Per-step line-buffered log. NOT \r-overwritten so tee/grep see it.
+        # Keep key=value pairs grep-friendly.
+        ppl = 2.0 ** bpb  # perplexity (byte-level)
+        print(
+            f"step={step:05d} loss={debiased_smooth_loss:.4f} bpb={bpb:.4f} ppl={ppl:.3f} "
+            f"bpt={bpt:.3f} bpt_div={bpt_ema:.2f} "
+            f"tps={tok_per_sec} dt_ms={dt*1000:.0f} mfu={mfu:.1f} "
+            f"lr={current_lr:.2e} vram={vram_mib:.0f}MiB "
+            f"pct={pct_done:.1f} epoch={epoch} remaining={remaining:.0f}s",
+            flush=True,
+        )
+
+        if step == 0:
+            gc.collect()
+            gc.freeze()
+            gc.disable()
+        # No periodic gc.collect() — we disabled+froze at step 0 on purpose,
+        # so a manual collect every 5k steps just re-scans frozen objects
+        # (burned ~900 ms/event in production) for no live-garbage reason.
+
+        if CKPT_INTERVAL > 0 and step > 0 and step % CKPT_INTERVAL == 0:
+            save_ckpt(
+                model,
+                optimizer,
+                config,
+                step,
+                total_training_time,
+                smooth_train_loss,
+                bpt_ema,
+                epoch,
+                LATEST_CKPT,
+            )
+
+        # Periodic mid-training validation so we can see the model learning
+        # English in real time (not just at the end). Small val batch so it
+        # doesn't eat significant training time.
+        mid_val_interval = int(os.environ.get("HYDRA_MID_VAL_INTERVAL", "500"))
+        if mid_val_interval > 0 and step > 0 and step % mid_val_interval == 0:
+            model.eval()
+            try:
+                # Defrag GPU memory before eval allocates fresh chunks —
+                # without this the eval path can OOM on 6GB cards even
+                # though total usage fits, because the allocator's free
+                # blocks are fragmented.
+                torch.cuda.empty_cache()
+                _orig_mid = _prepare_mod.EVAL_TOKENS
+                _prepare_mod.EVAL_TOKENS = 262144  # ~260K tokens, fast
+                with torch.no_grad():
+                    with autocast_ctx:
+                        mid_bpb = evaluate_bpb(model, tokenizer, DEVICE_BATCH_SIZE)
+                _prepare_mod.EVAL_TOKENS = _orig_mid
+                mid_ppl = 2.0 ** mid_bpb
+                print(f"[MID_VAL] step={step} val_bpb={mid_bpb:.4f} val_ppl={mid_ppl:.3f}", flush=True)
+
+                # Per-layer diagnostic panel. Only printed when HYDRA_LAYER_DIAGNOSTICS=1
+                # is set (otherwise the layer_* keys are absent from _metrics).
+                _diag_metrics = model.get_secondary_metrics()
+                _layer_keys = sorted([k for k in _diag_metrics.keys() if k.startswith('layer_')])
+                if _layer_keys:
+                    # Condense: one row per layer showing the four core signals.
+                    n_layers = len(model.blocks)
+                    print(f"[LAYER_DIAG] step={step}", flush=True)
+                    for li in range(n_layers):
+                        d_ratio = _diag_metrics.get(f'layer_{li}_delta_ratio', float('nan'))
+                        out_n   = _diag_metrics.get(f'layer_{li}_out_norm',    float('nan'))
+                        g_norm  = _diag_metrics.get(f'layer_{li}_grad_norm',   float('nan'))
+                        eff_r   = _diag_metrics.get(f'layer_{li}_eff_rank',    float('nan'))
+                        f_std   = _diag_metrics.get(f'layer_{li}_feat_std',    float('nan'))
+                        print(
+                            f"[LAYER_DIAG]   L{li:02d}  delta_ratio={d_ratio:.4f}  "
+                            f"out_norm={out_n:.4f}  grad_norm={g_norm:.3e}  "
+                            f"eff_rank={eff_r:.1f}  feat_std={f_std:.4f}",
+                            flush=True,
+                        )
+                    htm_proj_g = _diag_metrics.get('htm_proj_grad_norm', None)
+                    if htm_proj_g is not None:
+                        print(f"[LAYER_DIAG]   htm_proj grad_norm={htm_proj_g:.3e}", flush=True)
+            except Exception as e:
+                print(f"[MID_VAL] failed: {e}", flush=True)
+            model.train()
+
+        step += 1
+
+        if step > 10 and total_training_time >= TIME_BUDGET:
+            break
+
+    # Drain async postprocessing threads before eval
+    if _som_thread is not None:
+        _som_thread.join()
+    if _hestia_thread is not None:
+        _hestia_thread.join()
+    if _hestia_stream is not None:
+        _hestia_stream.synchronize()
+
+    total_tokens = step * TOTAL_BATCH_SIZE
+
+    # ----------------------------------------------------------------------
+    # SAVE ORDER (critical):
+    #   1. Save PRETRAIN_FINAL_CKPT with val_bpb=None  (hedge against eval OOM)
+    #   2. Save LATEST_CKPT with val_bpb=None          (hedge against eval OOM)
+    #   3. Run eval (may OOM on small GPUs; we survive it)
+    #   4. Re-save both ckpts with val_bpb filled in
+    # This way we NEVER lose the final trained weights to an eval crash.
+    # Previous ordering put eval first, so an eval-time OOM destroyed the
+    # only record of a 6h training run (2026-04-22 incident).
+    # ----------------------------------------------------------------------
+
+    save_ckpt(
+        model, optimizer, config, step, total_training_time,
+        smooth_train_loss, bpt_ema, epoch, PRETRAIN_FINAL_CKPT,
+        val_bpb=None, blocking=True,
+    )
+    save_ckpt(
+        model, optimizer, config, step, total_training_time,
+        smooth_train_loss, bpt_ema, epoch, LATEST_CKPT,
+        val_bpb=None, blocking=True,
+    )
+
+    # Now it's safe to eval — ckpts are on disk regardless of what happens here.
+    # HYDRA_EVAL_BATCH overrides DEVICE_BATCH_SIZE (env-tunable; default halves
+    # the training batch because eval holds activations for full sequence and
+    # does not benefit from overlap with backward). HYDRA_EVAL_TOKENS controls
+    # how many val tokens to sweep (default 2 M, short enough for autoresearch
+    # 5-min budgets).
+    val_bpb: float | None = None
+    # Eval batch: default to 4 on cloud GPUs (enough freed VRAM after optimizer
+    # clear), fall back to DEVICE_BATCH_SIZE//2 on tiny cards. Env-overridable.
+    _eval_B = int(os.environ.get("HYDRA_EVAL_BATCH",
+        str(max(1, DEVICE_BATCH_SIZE // 2) if DEVICE_BATCH_SIZE <= 8 else 4)))
+    # Eval tokens: default 1M (1,048,576) — gives statistically meaningful BPB
+    # (256 forward passes at B=4, seq=1024). Env-overridable for fast/slow sweeps.
+    _eval_tokens = int(os.environ.get("HYDRA_EVAL_TOKENS", str(1048576)))
+    try:
+        # Aggressive VRAM reclaim for 6GB cards. Peak training VRAM = 5.1GB
+        # which leaves < 1GB for the eval forward — the driver can't satisfy
+        # the allocation. Free EVERY tensor we don't strictly need:
+        #   - optimizer grads (set_to_none releases tensor)
+        #   - optimizer.state (fp32 Muon NS workspace, AdamW moments — ~size-of-params each)
+        #   - model internal caches (HTM subsample cache, SDR stash)
+        # After this, VRAM should be ~params only (bf16 ≈ 120MB at 60M params).
+        optimizer.zero_grad(set_to_none=True)
+        if hasattr(optimizer, 'state') and optimizer.state:
+            for p, st in list(optimizer.state.items()):
+                st.clear()
+            optimizer.state.clear()
+        for p in model.parameters():
+            if p.grad is not None:
+                p.grad = None
+        if hasattr(model, '_htm_cache'):
+            model._htm_cache = None
+        if hasattr(model, '_last_sdr'):
+            model._last_sdr = None
+        import gc as _gc
+        _gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        try:
+            _free_mb = torch.cuda.mem_get_info()[0] / 1024 / 1024
+            print(f"[VAL] free_vram_mb={_free_mb:.0f} (cleared optimizer state)", flush=True)
+        except Exception:
+            pass
+        print(f"[VAL] running eval on {_eval_tokens} tokens at B={_eval_B}...", flush=True)
+        model.eval()
+        _orig = _prepare_mod.EVAL_TOKENS
+        _prepare_mod.EVAL_TOKENS = _eval_tokens
+        # Nemotron path reads HYDRA_STREAM_EVAL_TOKENS env var directly,
+        # not _prepare_mod.EVAL_TOKENS. Sync both so eval budget is
+        # respected regardless of which dataloader path is active.
+        _orig_stream = os.environ.get("HYDRA_STREAM_EVAL_TOKENS")
+        os.environ["HYDRA_STREAM_EVAL_TOKENS"] = str(_eval_tokens)
+        with autocast_ctx:
+            val_bpb = evaluate_bpb(model, tokenizer, _eval_B)
+        _prepare_mod.EVAL_TOKENS = _orig
+        if _orig_stream is not None:
+            os.environ["HYDRA_STREAM_EVAL_TOKENS"] = _orig_stream
+        else:
+            os.environ.pop("HYDRA_STREAM_EVAL_TOKENS", None)
+        val_ppl = 2 ** val_bpb
+        print(f"[VAL] step={step} val_bpb={val_bpb:.4f} val_ppl={val_ppl:.3f}", flush=True)
+    except torch.cuda.OutOfMemoryError as e:
+        print(f"[VAL] SKIPPED (OOM): {e}", flush=True)
+        torch.cuda.empty_cache()
+    except Exception as e:
+        import traceback as _tb
+        print(f"[VAL] SKIPPED ({type(e).__name__}): {e}", flush=True)
+        _tb.print_exc()
+        try:
+            _free = torch.cuda.mem_get_info()[0] / 1024 / 1024
+            print(f"[VAL] post-crash free_vram_mb={_free:.0f}", flush=True)
+        except Exception:
+            pass
+
+    # Final ckpts with val_bpb filled in (if eval succeeded).
+    save_ckpt(
+        model, optimizer, config, step, total_training_time,
+        smooth_train_loss, bpt_ema, epoch, LATEST_CKPT,
+        val_bpb=val_bpb, blocking=True,
+    )
+    save_ckpt(
+        model, optimizer, config, step, total_training_time,
+        smooth_train_loss, bpt_ema, epoch, PRETRAIN_FINAL_CKPT,
+        val_bpb=val_bpb, blocking=True,
+    )
+
+    # Learnability #2: persist EMA weights alongside the raw checkpoint.
+    # latest_ema.pt contains ema_model.module (the Averaged params) so it
+    # can be loaded by evaluation / inference code that expects the same
+    # state_dict shape as the raw model.
+    if ema_model is not None:
+        try:
+            ema_ckpt_path = CACHE_DIR / "latest_ema.pt"
+            CACHE_DIR.mkdir(parents=True, exist_ok=True)
+            torch.save({
+                "model_state_dict": ema_model.module.state_dict(),
+                "config": asdict(config),
+                "step": step,
+                "epoch": epoch,
+                "train_seconds": total_training_time,
+                "val_bpb": val_bpb,
+                "ema_decay": EMA_DECAY,
+            }, str(ema_ckpt_path))
+            print(f"[EMA] saved {ema_ckpt_path} (step={step})", flush=True)
+        except Exception as _e:
+            print(f"[EMA] save failed: {_e}", flush=True)
+
+    run_factual_probes(model, tokenizer, device, autocast_ctx)
+
+    t_end = time.time()
+    startup_time = t_start_training - t_start
+    steady_state_mfu = (
+        100 * num_flops_per_token * TOTAL_BATCH_SIZE * (step - 10)
+        / total_training_time / GPU_BF16_PEAK_FLOPS
+        if total_training_time > 0 else 0
+    )
+    peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
+    metrics = model.get_secondary_metrics()
+
+    print("---")
+    print(f"val_bpb:          {val_bpb:.6f}" if val_bpb is not None else "val_bpb:          SKIPPED")
+    print(f"training_seconds: {total_training_time:.1f}")
+    print(f"total_seconds:    {t_end - t_start:.1f}")
+    print(f"peak_vram_mb:     {peak_vram_mb:.1f}")
+    print(f"mfu_percent:      {steady_state_mfu:.2f}")
+    print(f"total_tokens_M:   {total_tokens / 1e6:.1f}")
+    print(f"num_steps:        {step}")
+    print(f"num_params_M:     {num_params / 1e6:.1f}")
+    print(f"n_layer:          {N_LAYER}")
+    print(f"d_model:          {D_MODEL}")
+    print(f"engram_hit_rate:   {metrics.get('engram_hit_rate', 0.0):.4f}")
+    print(f"sdr_active_bits:  {metrics.get('sdr_active_bits', 0):.1f}")
+    print(f"htm_anomaly:      {metrics.get('htm_anomaly', 0):.4f}")
+
+    # Per-layer summary panel — only printed when diagnostics were active.
+    _layer_keys = sorted([k for k in metrics.keys() if k.startswith('layer_')])
+    if _layer_keys:
+        n_layers = len(model.blocks)
+        print("--- per-layer diagnostic panel ---")
+        for li in range(n_layers):
+            d_ratio = metrics.get(f'layer_{li}_delta_ratio', float('nan'))
+            out_n   = metrics.get(f'layer_{li}_out_norm',    float('nan'))
+            g_norm  = metrics.get(f'layer_{li}_grad_norm',   float('nan'))
+            eff_r   = metrics.get(f'layer_{li}_eff_rank',    float('nan'))
+            f_std   = metrics.get(f'layer_{li}_feat_std',    float('nan'))
+            print(
+                f"L{li:02d}  delta_ratio={d_ratio:.4f}  out_norm={out_n:.4f}  "
+                f"grad_norm={g_norm:.3e}  eff_rank={eff_r:.1f}  feat_std={f_std:.4f}"
+            )
+
+    # Emit full metrics dictionary as JSON for sweep aggregation. Path from
+    # HYDRA_METRICS_OUT env var; default=/tmp/hydra_run_metrics.json. Always
+    # written (even without diagnostics) so the aggregator can compare runs.
+    _metrics_out = os.environ.get("HYDRA_METRICS_OUT", "/tmp/hydra_run_metrics.json")
+    try:
+        _dump = dict(metrics)
+        _dump.update({
+            'val_bpb': (float(val_bpb) if val_bpb is not None else None),
+            'val_ppl': (float(val_ppl) if val_ppl is not None else None),
+            'n_layer': int(N_LAYER),
+            'd_model': int(D_MODEL),
+            'num_params_M': float(num_params / 1e6),
+            'num_steps': int(step),
+            'total_tokens_M': float(total_tokens / 1e6),
+            'peak_vram_mb': float(peak_vram_mb),
+            'training_seconds': float(total_training_time),
+            'sdr_target_active': int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327")),
+        })
+        Path(_metrics_out).parent.mkdir(parents=True, exist_ok=True)
+        with open(_metrics_out, 'w') as _f:
+            json.dump(_dump, _f, indent=2, sort_keys=True)
+        print(f"[METRICS] wrote {_metrics_out}", flush=True)
+        # Also emit a single-line JSON to stdout so the sweep aggregator can
+        # scrape it from HF Jobs logs without pulling files out of the container.
+        print("[METRICS_JSON] " + json.dumps(_dump, sort_keys=True), flush=True)
+    except Exception as _e:
+        print(f"[METRICS] write failed: {_e}", flush=True)
+
+    run_factual_english(model, tokenizer, MAX_SEQ_LEN)
+    # startup_time is informative but not printed (preserve historical output)
+    _ = startup_time