Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+benchmark_results/model_comparison.png filter=lfs diff=lfs merge=lfs -text
+benchmark_results/trained_model_comparison.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,170 @@
+# WandB
+/wandb/
+# checkpoints
+/checkpoints/
+# cache
+/cache/
+# data
+/data/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/*.safetensors
+hf_upload/

.gitmodules

@@ -0,0 +1,9 @@
+[submodule "dataset/raw-data/ConceptARC"]
+	path = dataset/raw-data/ConceptARC
+	url = git@github.com:victorvikram/ConceptARC.git
+[submodule "dataset/raw-data/ARC-AGI"]
+	path = dataset/raw-data/ARC-AGI
+	url = git@github.com:fchollet/ARC-AGI.git
+[submodule "dataset/raw-data/ARC-AGI-2"]
+	path = dataset/raw-data/ARC-AGI-2
+	url = git@github.com:arcprize/ARC-AGI-2.git

.vscode/launch.json

@@ -0,0 +1,26 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        },
+        {
+            "name": "Debug: Single GPU",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "pretrain.py",
+            "args": [],
+            "env": {
+                "OMP_NUM_THREADS": "1",
+                "DISABLE_COMPILE": "true"
+            }
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.analysis.typeCheckingMode": "standard"
+}

GUIDE_SPEED_TEST.md

@@ -0,0 +1,62 @@
+# Speed Test Guide: HRM SRAM/DRAM Benchmarking
+
+Use this guide to measure the **hardware performance** (latency, throughput, memory overhead) of the memory-tiered architecture.
+
+> [!NOTE]
+> This guide uses **synthetic data**. No external datasets or prior training are required. You can run this immediately after setup.
+
+---
+
+### 1. Environment Setup
+Create a clean virtual environment and install dependencies:
+
+```bash
+python -m venv venv
+source venv/bin/activate
+pip install --upgrade pip
+pip install -r requirements.txt
+pip install triton matplotlib
+```
+
+---
+
+### 2. Run Hardware Comparison
+This command compares the **Baseline (Original)** model vs. the **Tiered (Optimized)** model across multiple configurations.
+
+```bash
+python run_benchmark.py --mode compare --plot
+```
+
+**What this does:**
+1. Generates random tensors to simulate a reasoning workload.
+2. Compiles Triton kernels for the L-level (SRAM) and H-level (DRAM) paths.
+3. Records precise GPU timing using `torch.cuda.Event`.
+4. Saves charts to `benchmark_results/`.
+
+---
+
+### 3. Deep Dive into Tiered Metrics
+To get a detailed breakdown of memory tier behavior:
+
+```bash
+python run_benchmark.py --mode tiered --iterations 50 --batch-sizes 8 --seq-lens 128
+```
+
+**Key Metrics to watch:**
+- **L_lat(μs):** Speed of the fast-updating L-level module in SRAM.
+- **H_lat(μs):** Speed of the planning H-level module in DRAM.
+- **H/L Ratio:** Shows the latency multiplier between memory tiers.
+- **Memory Efficiency:** Percentage of time spent on math vs. memory transfer.
+
+---
+
+### 4. Custom Hardware Stress Test
+Test the limits of your GPU by increasing batch sizes or sequence lengths:
+
+```bash
+python run_benchmark.py \
+    --mode compare \
+    --batch-sizes 64,128 \
+    --seq-lens 256,512 \
+    --hidden-size 1024
+```

GUIDE_TRAIN_TEST.md

@@ -0,0 +1,67 @@
+# Training Guide: HRM SRAM/DRAM Learning Performance
+
+Use this guide to train the memory-tiered HRM on **real datasets** and verify its reasoning accuracy.
+
+---
+
+### 1. Dataset Generation
+The model needs to learn how to solve puzzles. First, generate a training/test set (e.g., Sudoku).
+
+```bash
+# Activate environment
+source venv/bin/activate
+
+# Build a small Sudoku dataset (1000 examples)
+python dataset/build_sudoku_dataset.py \
+    --output-dir data/sudoku-1k \
+    --subsample-size 1000 \
+    --num-aug 10
+```
+
+---
+
+### 2. Start Training
+Point the trainer to the tiered architecture configuration.
+
+#### Train the Tiered Model (New)
+```bash
+python pretrain.py \
+    arch=hrm_tiered \
+    data_path=data/sudoku-1k \
+    epochs=1000 \
+    global_batch_size=384
+```
+
+#### Train the Baseline Model (Comparison)
+```bash
+python pretrain.py \
+    arch=hrm_v1 \
+    data_path=data/sudoku-1k \
+    epochs=1000 \
+    global_batch_size=384
+```
+
+---
+
+### 3. Verify Reasoning Accuracy
+Once training is complete (or during training), check the accuracy metrics in your W&B dashboard or via the evaluation script:
+
+```bash
+# Replace with the path to your generated checkpoint
+python evaluate.py checkpoint=checkpoints/Sudoku_ACT-torch/YOUR_RUN/step_1000
+```
+
+**What to look for:**
+- **`eval/exact_accuracy`**: Fraction of puzzles solved perfectly.
+- **`eval/steps`**: Average number of reasoning steps taken by the ACT (Adaptive Computation Time) module.
+- **`eval/q_halt_accuracy`**: How well the model learns when to stop thinking.
+
+---
+
+### 4. Training on ARC (Artificial General Intelligence Benchmark)
+To train on the more complex ARC-AGI-2 dataset:
+
+```bash
+python dataset/build_arc_dataset.py --output-dir data/arc-2
+python pretrain.py arch=hrm_tiered data_path=data/arc-2
+```

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,215 @@
----
-license: mit
----
+# Hierarchical Reasoning Model
+
+![](./assets/hrm.png)
+
+Reasoning, the process of devising and executing complex goal-oriented action sequences, remains a critical challenge in AI.
+Current large language models (LLMs) primarily employ Chain-of-Thought (CoT) techniques, which suffer from brittle task decomposition, extensive data requirements, and high latency. Inspired by the hierarchical and multi-timescale processing in the human brain, we propose the Hierarchical Reasoning Model (HRM), a novel recurrent architecture that attains significant computational depth while maintaining both training stability and efficiency.
+HRM executes sequential reasoning tasks in a single forward pass without explicit supervision of the intermediate process, through two interdependent recurrent modules: a high-level module responsible for slow, abstract planning, and a low-level module handling rapid, detailed computations. With only 27 million parameters, HRM achieves exceptional performance on complex reasoning tasks using only 1000 training samples. The model operates without pre-training or CoT data, yet achieves nearly perfect performance on challenging tasks including complex Sudoku puzzles and optimal path finding in large mazes.
+Furthermore, HRM outperforms much larger models with significantly longer context windows on the Abstraction and Reasoning Corpus (ARC), a key benchmark for measuring artificial general intelligence capabilities.
+These results underscore HRM’s potential as a transformative advancement toward universal computation and general-purpose reasoning systems.
+
+**Join our Discord Community: [https://discord.gg/sapient](https://discord.gg/sapient)**
+
+
+## Quick Start Guide 🚀
+
+### Prerequisites ⚙️
+
+Ensure PyTorch and CUDA are installed. The repo needs CUDA extensions to be built. If not present, run the following commands:
+
+```bash
+# Install CUDA 12.6
+CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run
+
+wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
+sudo sh cuda_installer.run --silent --toolkit --override
+
+export CUDA_HOME=/usr/local/cuda-12.6
+
+# Install PyTorch with CUDA 12.6
+PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu126
+
+pip3 install torch torchvision torchaudio --index-url $PYTORCH_INDEX_URL
+
+# Additional packages for building extensions
+pip3 install packaging ninja wheel setuptools setuptools-scm
+```
+
+Then install FlashAttention. For Hopper GPUs, install FlashAttention 3
+
+```bash
+git clone git@github.com:Dao-AILab/flash-attention.git
+cd flash-attention/hopper
+python setup.py install
+```
+
+For Ampere or earlier GPUs, install FlashAttention 2
+
+```bash
+pip3 install flash-attn
+```
+
+## Install Python Dependencies 🐍
+
+```bash
+pip install -r requirements.txt
+```
+
+## W&B Integration 📈
+
+This project uses [Weights & Biases](https://wandb.ai/) for experiment tracking and metric visualization. Ensure you're logged in:
+
+```bash
+wandb login
+```
+
+## Run Experiments
+
+### Quick Demo: Sudoku Solver 💻🗲
+
+Train a master-level Sudoku AI capable of solving extremely difficult puzzles on a modern laptop GPU. 🧩
+
+```bash
+# Download and build Sudoku dataset
+python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000
+
+# Start training (single GPU, smaller batch size)
+OMP_NUM_THREADS=8 python pretrain.py data_path=data/sudoku-extreme-1k-aug-1000 epochs=20000 eval_interval=2000 global_batch_size=384 lr=7e-5 puzzle_emb_lr=7e-5 weight_decay=1.0 puzzle_emb_weight_decay=1.0
+```
+
+Runtime: ~10 hours on a RTX 4070 laptop GPU
+
+## Trained Checkpoints 🚧
+
+ - [ARC-AGI-2](https://huggingface.co/sapientinc/HRM-checkpoint-ARC-2)
+ - [Sudoku 9x9 Extreme (1000 examples)](https://huggingface.co/sapientinc/HRM-checkpoint-sudoku-extreme)
+ - [Maze 30x30 Hard (1000 examples)](https://huggingface.co/sapientinc/HRM-checkpoint-maze-30x30-hard)
+
+To use the checkpoints, see Evaluation section below.
+
+## Full-scale Experiments 🔵
+
+Experiments below assume an 8-GPU setup.
+
+### Dataset Preparation
+
+```bash
+# Initialize submodules
+git submodule update --init --recursive
+
+# ARC-1
+python dataset/build_arc_dataset.py  # ARC offical + ConceptARC, 960 examples
+# ARC-2
+python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000  # ARC-2 official, 1120 examples
+
+# Sudoku-Extreme
+python dataset/build_sudoku_dataset.py  # Full version
+python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000  # 1000 examples
+
+# Maze
+python dataset/build_maze_dataset.py  # 1000 examples
+```
+
+### Dataset Visualization
+
+Explore the puzzles visually:
+
+* Open `puzzle_visualizer.html` in your browser.
+* Upload the generated dataset folder located in `data/...`.
+
+## Launch experiments
+
+### Small-sample (1K)
+
+ARC-1:
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py 
+```
+
+*Runtime:* ~24 hours
+
+ARC-2:
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/arc-2-aug-1000
+```
+
+*Runtime:* ~24 hours (checkpoint after 8 hours is often sufficient)
+
+Sudoku Extreme (1k):
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-extreme-1k-aug-1000 epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0
+```
+
+*Runtime:* ~10 minutes
+
+Maze 30x30 Hard (1k):
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/maze-30x30-hard-1k epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0
+```
+
+*Runtime:* ~1 hour
+
+### Full Sudoku-Hard
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-hard-full epochs=100 eval_interval=10 lr_min_ratio=0.1 global_batch_size=2304 lr=3e-4 puzzle_emb_lr=3e-4 weight_decay=0.1 puzzle_emb_weight_decay=0.1 arch.loss.loss_type=softmax_cross_entropy arch.L_cycles=8 arch.halt_max_steps=8 arch.pos_encodings=learned
+```
+
+*Runtime:* ~2 hours
+
+## Streamlined Training & Benchmarking 🛠️
+For easier monitoring and automated reporting, use these consolidated scripts:
+
+### 1. Unified Training
+Train either the Baseline or Tiered (SRAM/DRAM) model with multi-process support and integrated W&B monitoring.
+```bash
+./train_hrm.py arch=hrm_tiered epochs=1000 data_path=data/sudoku-1k
+```
+
+### 2. Comprehensive Evaluation
+Load a checkpoint and generate a detailed Markdown report (`_report.md`) with accuracy breakdowns.
+```bash
+./eval_hrm.py checkpoint=checkpoints/hrm_run_.../best_model.pt
+```
+
+### 3. Hardware Benchmark
+Compare Baseline vs. Tiered performance across multiple batch sizes and sequence lengths. Generates plots and CSV reports.
+```bash
+./benchmark_hrm.py --batch-sizes 1,8,32 --seq-lens 64,128 --plot
+```
+
+## Evaluation
+
+Evaluate your trained models:
+
+* Check `eval/exact_accuracy` in W&B.
+* For ARC-AGI, follow these additional steps:
+
+```bash
+OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 evaluate.py checkpoint=<CHECKPOINT_PATH>
+```
+
+* Then use the provided `arc_eval.ipynb` notebook to finalize and inspect your results.
+
+## Notes
+
+ - Small-sample learning typically exhibits accuracy variance of around ±2 points.
+ - For Sudoku-Extreme (1,000-example dataset), late-stage overfitting may cause numerical instability during training and Q-learning. It is advisable to use early stopping once the training accuracy approaches 100%.
+
+## Citation 📜
+
+```bibtex
+@misc{wang2025hierarchicalreasoningmodel,
+      title={Hierarchical Reasoning Model}, 
+      author={Guan Wang and Jin Li and Yuhao Sun and Xing Chen and Changling Liu and Yue Wu and Meng Lu and Sen Song and Yasin Abbasi Yadkori},
+      year={2025},
+      eprint={2506.21734},
+      archivePrefix={arXiv},
+      primaryClass={cs.AI},
+      url={https://arxiv.org/abs/2506.21734}, 
+}
+```

README_TIERED.md

@@ -0,0 +1,116 @@
+# HRM SRAM/DRAM Memory Tiering: Training & Inference Guide
+
+This guide explains how to train, benchmark, and evaluate both the **Original (Baseline)** and the **New (Memory-Tiered)** HRM models.
+
+---
+
+## 🚀 Setup
+
+### 1. Prerequisites
+Ensure you have a CUDA-capable GPU. The tiered model uses **Python Triton** for high-speed SRAM kernels.
+
+### 2. Environment Setup (Recommended)
+Create and activate a new virtual environment to keep dependencies isolated:
+
+```bash
+# Create venv
+python -m venv venv
+
+# Activate venv (Mac/Linux)
+source venv/bin/activate
+
+# Install dependencies
+pip install --upgrade pip
+pip install -r requirements.txt
+pip install triton matplotlib
+```
+
+---
+
+## 🏋️ Training the Models
+
+Training is handled by `pretrain.py` using Hydra for configuration.
+
+### A. Train the Original (Baseline) Model
+Uses standard GPU global memory (DRAM) for all layers.
+
+```bash
+# Example for Sudoku dataset
+OMP_NUM_THREADS=8 python pretrain.py \
+    arch=hrm_v1 \
+    data_path=data/sudoku-extreme-1k-aug-1000 \
+    epochs=10000 \
+    global_batch_size=384
+```
+
+### B. Train the New (Tiered) Model
+Optimizes the L-level module by pinning it to simulated SRAM using Triton kernels.
+
+```bash
+OMP_NUM_THREADS=8 python pretrain.py \
+    arch=hrm_tiered \
+    data_path=data/sudoku-extreme-1k-aug-1000 \
+    epochs=10000 \
+    global_batch_size=384
+```
+
+---
+
+## 📊 Testing Inference & Benchmarking
+
+The `run_benchmark.py` script is the primary tool for measuring latencies, throughput, and memory efficiency.
+
+### 1. Comparison Mode (Recommended)
+Automatically runs both models across multiple batch sizes and sequence lengths, then generates performance charts.
+
+```bash
+python run_benchmark.py \
+    --mode compare \
+    --batch-sizes 1,8,32 \
+    --seq-lens 64,128 \
+    --plot \
+    --output-dir benchmark_results/
+```
+*   **Outputs:** `benchmark_results/benchmark_comparison.png` and `results.json`.
+
+### 2. Tiered-Only Benchmark (Deep Dive)
+Get detailed metrics for just the tiered model, including memory transfer overhead and SRAM hit rates.
+
+```bash
+python run_benchmark.py \
+    --mode tiered \
+    --iterations 50 \
+    --batch-sizes 8 \
+    --output tiered_stats.json
+```
+
+---
+
+## ✅ Evaluation
+
+To verify the accuracy of a trained checkpoint on a test set:
+
+### Using the original evaluate script:
+```bash
+python evaluate.py checkpoint=checkpoints/PATH_TO_YOUR_STEP
+```
+
+### Using the tiered model class manually:
+If you are writing a custom script, you can import the models as follows:
+
+```python
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1 # Original
+from models.hrm.hrm_tiered import HRM_Tiered                       # Tiered
+```
+
+---
+
+## 🛠 Model Location Reference
+
+| Folder/File | Contents |
+| :--- | :--- |
+| `models/hrm/hrm_act_v1.py` | Original model code (Baseline) |
+| `models/hrm/hrm_tiered.py` | New Tiered model code |
+| `models/triton_kernels.py` | Triton SRAM/DRAM kernels |
+| `models/memory_tier.py` | Memory Tier Manager and CUDA stream logic |
+| `config/arch/hrm_tiered.yaml` | Configuration for the tiered version |

arc_eval.ipynb

@@ -0,0 +1,252 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "from glob import glob\n",
+    "import hashlib\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.colors as mcolors\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "import numpy as np\n",
+    "from numba import njit\n",
+    "\n",
+    "from dataset.common import inverse_dihedral_transform\n",
+    "\n",
+    "\n",
+    "DATASET_PATH = \"data/arc-aug-1000\"  # ARC-1\n",
+    "# DATASET_PATH = \"data/arc-2-aug-1000\"  # ARC-2\n",
+    "\n",
+    "CHECKPOINT_PATH = \"checkpoints/Arc-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 amphibian-turaco/step_414456\"\n",
+    "\n",
+    "\n",
+    "PAD_PUZZLE_IDENTIFIER = 0\n",
+    "\n",
+    "# Visualization\n",
+    "ARC_COLOR_MAP = mcolors.ListedColormap([\n",
+    "    \"#000000\",  # symbol_0: black\n",
+    "    \"#0074D9\",  # symbol_1: blue\n",
+    "    \"#FF4136\",  # symbol_2: red\n",
+    "    \"#2ECC40\",  # symbol_3: green\n",
+    "    \"#FFDC00\",  # symbol_4: yellow\n",
+    "    \"#AAAAAA\",  # symbol_5: grey\n",
+    "    \"#F012BE\",  # symbol_6: fuschia\n",
+    "    \"#FF851B\",  # symbol_7: orange\n",
+    "    \"#7FDBFF\",  # symbol_8: teal\n",
+    "    \"#870C25\"   # symbol_9: brown\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_identifiers_and_preds(dataset_path: str, checkpoint_path: str):\n",
+    "    # Load puzzle identifiers\n",
+    "    with open(os.path.join(dataset_path, \"identifiers.json\"), \"r\") as f:\n",
+    "        identifier_map = json.load(f)\n",
+    "        \n",
+    "    # Load preds\n",
+    "    all_preds = {}\n",
+    "    for filename in glob(f\"{checkpoint_path}_all_preds.*\"):\n",
+    "        preds = torch.load(filename)\n",
+    "        for k, v in preds.items():\n",
+    "            all_preds.setdefault(k, [])\n",
+    "            all_preds[k].append(v)\n",
+    "            \n",
+    "        del preds\n",
+    "\n",
+    "    all_preds = {k: torch.cat(v, dim=0) for k, v in all_preds.items()}\n",
+    "    \n",
+    "    # Remove paddings\n",
+    "    mask = all_preds[\"puzzle_identifiers\"] != PAD_PUZZLE_IDENTIFIER\n",
+    "    all_preds = {k: v[mask] for k, v in all_preds.items()}\n",
+    "\n",
+    "    return identifier_map, all_preds\n",
+    "\n",
+    "\n",
+    "def inverse_aug(name: str, grid: np.ndarray):\n",
+    "    if \"_\" not in name:\n",
+    "        return grid\n",
+    "\n",
+    "    trans_id, perm = name.split(\"_\")[-2:]\n",
+    "    trans_id = int(trans_id[1:])  # Remove \"t\" letter\n",
+    "    inv_perm = np.argsort(list(perm))\n",
+    "    \n",
+    "    return inv_perm[inverse_dihedral_transform(grid, trans_id)]\n",
+    "\n",
+    "\n",
+    "def grid_hash(grid: np.ndarray):\n",
+    "    return hash((grid.tobytes(), grid.shape))\n",
+    "\n",
+    "\n",
+    "@njit\n",
+    "def crop(grid: np.ndarray):\n",
+    "    # Find maximum-sized rectangle without any EOS token inside.\n",
+    "    grid = grid.reshape(30, 30)\n",
+    "\n",
+    "    max_area = 0\n",
+    "    max_size = (0, 0)\n",
+    "    nr, nc = grid.shape\n",
+    "    \n",
+    "    num_c = nc\n",
+    "    for num_r in range(1, nr + 1):\n",
+    "        # Scan for maximum c\n",
+    "        for c in range(1, num_c + 1):\n",
+    "            x = grid[num_r - 1, c - 1]\n",
+    "            if (x < 2) | (x > 11):\n",
+    "                num_c = c - 1\n",
+    "                break\n",
+    "        \n",
+    "        area = num_r * num_c\n",
+    "        if area > max_area:\n",
+    "            max_area = area\n",
+    "            max_size = (num_r, num_c)\n",
+    "\n",
+    "    return grid[:max_size[0], :max_size[1]] - 2\n",
+    "\n",
+    "\n",
+    "def test(visualize, Ks=[1, 2, 10, 100, 1000]):\n",
+    "    identifier_map, all_preds = load_identifiers_and_preds(DATASET_PATH, CHECKPOINT_PATH)\n",
+    "    \n",
+    "    global_hmap = {}\n",
+    "    \n",
+    "    # Get puzzles and corresponding answers\n",
+    "    puzzle_labels = {}\n",
+    "    for identifier, input, label in zip(all_preds[\"puzzle_identifiers\"], all_preds[\"inputs\"], all_preds[\"labels\"]):\n",
+    "        name = identifier_map[identifier]\n",
+    "        if \"_\" not in name:   # Not-augmented\n",
+    "            puzzle_labels.setdefault(name, {})\n",
+    "            \n",
+    "            input = crop(input.numpy())\n",
+    "            label = crop(label.numpy())\n",
+    "\n",
+    "            input_hash = grid_hash(input)\n",
+    "            label_hash = grid_hash(label)\n",
+    "\n",
+    "            global_hmap[input_hash] = input\n",
+    "            global_hmap[label_hash] = label\n",
+    "\n",
+    "            assert input_hash not in puzzle_labels[name]\n",
+    "            puzzle_labels[name][input_hash] = label_hash\n",
+    "            \n",
+    "    print (\"Number of puzzles\", len(puzzle_labels))\n",
+    "    \n",
+    "    # Argmax prediction\n",
+    "    preds = all_preds[\"logits\"].argmax(-1)\n",
+    "\n",
+    "    # Collate\n",
+    "    pred_answers = {}\n",
+    "    for identifier, input, pred, q in zip(all_preds[\"puzzle_identifiers\"], all_preds[\"inputs\"], preds, all_preds[\"q_halt_logits\"].sigmoid()):\n",
+    "        name = identifier_map[identifier]\n",
+    "        orig_name = name.split(\"_\")[0]\n",
+    "        \n",
+    "        input = input.numpy()\n",
+    "        input_hash = grid_hash(inverse_aug(name, crop(input)))\n",
+    "        assert input_hash in puzzle_labels[orig_name]\n",
+    "        \n",
+    "        pred = inverse_aug(name, crop(pred.numpy()))\n",
+    "        pred_hash = grid_hash(pred)\n",
+    "        global_hmap[pred_hash] = pred\n",
+    "        \n",
+    "        pred_answers.setdefault(orig_name, {})\n",
+    "        pred_answers[orig_name].setdefault(input_hash, [])\n",
+    "        pred_answers[orig_name][input_hash].append((pred_hash, q.item()))\n",
+    "\n",
+    "    # test-1\n",
+    "    if visualize:\n",
+    "        num_figs = sum(len(tests) for name, tests in puzzle_labels.items())\n",
+    "        fig, axes = plt.subplots(num_figs, 4, figsize=(8, num_figs * 4))\n",
+    "        \n",
+    "        fig_id = 0\n",
+    "    \n",
+    "    correct = [0 for _ in range(len(Ks))]\n",
+    "    for name, tests in puzzle_labels.items():\n",
+    "        num_test_correct = [0 for _ in range(len(Ks))]\n",
+    "        for input_hash, label_hash in tests.items():\n",
+    "            p = pred_answers[name][input_hash]\n",
+    "            p_map = {}\n",
+    "            \n",
+    "            for h, q in p:\n",
+    "                p_map.setdefault(h, [0, 0])\n",
+    "                p_map[h][0] += 1\n",
+    "                p_map[h][1] += q\n",
+    "                \n",
+    "            for h, stats in p_map.items():\n",
+    "                stats[1] /= stats[0]\n",
+    "                \n",
+    "            p_map = sorted(p_map.items(), key=lambda kv: kv[1], reverse=True)\n",
+    "\n",
+    "            # 2-vote\n",
+    "            for i, k in enumerate(Ks):\n",
+    "                ok = False\n",
+    "                for h, stats in p_map[:k]:\n",
+    "                    ok |= h == label_hash\n",
+    "                    \n",
+    "                num_test_correct[i] += ok\n",
+    "\n",
+    "            if visualize:\n",
+    "                # Show input and ground truth\n",
+    "                axes[fig_id, 0].imshow(global_hmap[input_hash], cmap=ARC_COLOR_MAP)\n",
+    "                axes[fig_id, 0].set_title(f\"{name}\\nInput\")\n",
+    "                axes[fig_id, 0].axis('off')\n",
+    "                \n",
+    "                axes[fig_id, 1].imshow(global_hmap[label_hash], cmap=ARC_COLOR_MAP)\n",
+    "                axes[fig_id, 1].set_title(f\"{name}\\nAnswer\")\n",
+    "                axes[fig_id, 1].axis('off')\n",
+    "                \n",
+    "                trial_id = 2\n",
+    "                for h, stats in p_map[:2]:\n",
+    "                    ans = global_hmap[h]\n",
+    "                    \n",
+    "                    axes[fig_id, trial_id].imshow(ans, cmap=ARC_COLOR_MAP)\n",
+    "                    axes[fig_id, trial_id].set_title(f\"{name}\\nTrial {trial_id}\")\n",
+    "                    axes[fig_id, trial_id].axis('off')\n",
+    "                    \n",
+    "                    trial_id += 1\n",
+    "                \n",
+    "                fig_id += 1\n",
+    "            \n",
+    "        # Total correctness\n",
+    "        for i in range(len(Ks)):\n",
+    "            correct[i] += num_test_correct[i] == len(tests)\n",
+    "\n",
+    "    for i, k in enumerate(Ks):\n",
+    "        print (f\"{k}-shot: {correct[i] / len(puzzle_labels) * 100:.2f}%\")\n",
+    "\n",
+    "\n",
+    "test(visualize=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

assets/npyjs.js

@@ -0,0 +1,176 @@
+class npyjs {
+
+    constructor(opts) {
+        if (opts && !('convertFloat16' in opts)) {
+            console.warn([
+                "npyjs constructor now accepts {convertFloat16?: boolean}.",
+                "For usage, go to https://github.com/jhuapl-boss/npyjs."
+            ].join(" "));
+        }
+
+        this.convertFloat16 = opts?.convertFloat16 ?? true;
+
+        this.dtypes = {
+            "<u1": {
+                name: "uint8",
+                size: 8,
+                arrayConstructor: Uint8Array,
+            },
+            "|u1": {
+                name: "uint8",
+                size: 8,
+                arrayConstructor: Uint8Array,
+            },
+            "<u2": {
+                name: "uint16",
+                size: 16,
+                arrayConstructor: Uint16Array,
+            },
+            "|i1": {
+                name: "int8",
+                size: 8,
+                arrayConstructor: Int8Array,
+            },
+            "<i2": {
+                name: "int16",
+                size: 16,
+                arrayConstructor: Int16Array,
+            },
+            "<u4": {
+                name: "uint32",
+                size: 32,
+                arrayConstructor: Uint32Array,
+            },
+            "<i4": {
+                name: "int32",
+                size: 32,
+                arrayConstructor: Int32Array,
+            },
+            "<u8": {
+                name: "uint64",
+                size: 64,
+                arrayConstructor: BigUint64Array,
+            },
+            "<i8": {
+                name: "int64",
+                size: 64,
+                arrayConstructor: BigInt64Array,
+            },
+            "<f4": {
+                name: "float32",
+                size: 32,
+                arrayConstructor: Float32Array
+            },
+            "<f8": {
+                name: "float64",
+                size: 64,
+                arrayConstructor: Float64Array
+            },
+            "<f2": {
+                name: "float16",
+                size: 16,
+                arrayConstructor: Uint16Array,
+                converter: this.convertFloat16 ? this.float16ToFloat32Array : undefined
+            },
+        };
+    }
+
+    float16ToFloat32Array(float16Array) {
+        const length = float16Array.length;
+        const float32Array = new Float32Array(length);
+        
+        for (let i = 0; i < length; i++) {
+            float32Array[i] = npyjs.float16ToFloat32(float16Array[i]);
+        }
+        
+        return float32Array;
+    }
+
+    static float16ToFloat32(float16) {
+        // Extract the parts of the float16
+        const sign = (float16 >> 15) & 0x1;
+        const exponent = (float16 >> 10) & 0x1f;
+        const fraction = float16 & 0x3ff;
+
+        // Handle special cases
+        if (exponent === 0) {
+            if (fraction === 0) {
+                // Zero
+                return sign ? -0 : 0;
+            }
+            // Denormalized number
+            return (sign ? -1 : 1) * Math.pow(2, -14) * (fraction / 0x400);
+        } else if (exponent === 0x1f) {
+            if (fraction === 0) {
+                // Infinity
+                return sign ? -Infinity : Infinity;
+            }
+            // NaN
+            return NaN;
+        }
+
+        // Normalized number
+        return (sign ? -1 : 1) * Math.pow(2, exponent - 15) * (1 + fraction / 0x400);
+    }
+
+    parse(arrayBufferContents) {
+        // const version = arrayBufferContents.slice(6, 8); // Uint8-encoded
+        const headerLength = new DataView(arrayBufferContents.slice(8, 10)).getUint8(0);
+        const offsetBytes = 10 + headerLength;
+
+        const hcontents = new TextDecoder("utf-8").decode(
+            new Uint8Array(arrayBufferContents.slice(10, 10 + headerLength))
+        );
+        const header = JSON.parse(
+            hcontents
+                .toLowerCase() // True -> true
+                .replace(/'/g, '"')
+                .replace("(", "[")
+                .replace(/,*\),*/g, "]")
+        );
+        const shape = header.shape;
+        const dtype = this.dtypes[header.descr];
+
+        if (!dtype) {
+            console.error(`Unsupported dtype: ${header.descr}`);
+            return null;
+        }
+
+        const nums = new dtype.arrayConstructor(
+            arrayBufferContents,
+            offsetBytes
+        );
+
+        // Convert float16 to float32 if converter exists
+        const data = dtype.converter ? dtype.converter.call(this, nums) : nums;
+
+        return {
+            dtype: dtype.name,
+            data: data,
+            shape,
+            fortranOrder: header.fortran_order
+        };
+    }
+
+    async load(filename, callback, fetchArgs) {
+        /*
+        Loads an array from a stream of bytes.
+        */
+        fetchArgs = fetchArgs || {};
+        let arrayBuf;
+        // If filename is ArrayBuffer
+        if (filename instanceof ArrayBuffer) {
+            arrayBuf = filename;
+        }
+        // If filename is a file path
+        else {
+            const resp = await fetch(filename, { ...fetchArgs });
+            arrayBuf = await resp.arrayBuffer();
+        }
+        const result = this.parse(arrayBuf);
+        if (callback) {
+            return callback(result);
+        }
+        return result;
+    }
+}

benchmark.py

@@ -0,0 +1,686 @@
+"""
+Comprehensive Benchmarking Suite for HRM SRAM/DRAM Memory Tiering.
+
+Measures and compares:
+  - Per-module latency (L-level SRAM vs H-level DRAM)
+  - End-to-end inference latency
+  - Throughput (samples/sec)
+  - Memory usage per tier (SRAM / DRAM)
+  - Cross-tier transfer overhead
+  - SRAM hit rate
+  - Triton kernel-level profiling
+  - GPU utilization / power draw (when available)
+"""
+
+import json
+import time
+import math
+import os
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, asdict
+
+import torch
+import torch.nn.functional as F
+
+from models.memory_tier import MemoryTierManager
+from models.hrm.hrm_tiered import HRM_Tiered, HRM_Tiered_Inner
+from models.hrm.hrm_act_v1 import (
+    HierarchicalReasoningModel_ACTV1,
+    HierarchicalReasoningModel_ACTV1Config,
+    HierarchicalReasoningModel_ACTV1InnerCarry,
+)
+from models.triton_kernels import triton_memory_latency_probe
+
+
+@dataclass
+class BenchmarkResult:
+    """Results from a single benchmark run."""
+    model_name: str
+    batch_size: int
+    seq_len: int
+    hidden_size: int
+    H_cycles: int
+    L_cycles: int
+    H_layers: int
+    L_layers: int
+    num_iterations: int
+    warmup_iterations: int
+
+    # Latency (μs)
+    l_level_latency_mean_us: float
+    l_level_latency_min_us: float
+    l_level_latency_max_us: float
+    l_level_latency_std_us: float
+
+    h_level_latency_mean_us: float
+    h_level_latency_min_us: float
+    h_level_latency_max_us: float
+    h_level_latency_std_us: float
+
+    total_inference_latency_mean_ms: float
+    total_inference_latency_min_ms: float
+    total_inference_latency_max_ms: float
+    total_inference_latency_std_ms: float
+
+    # Throughput
+    throughput_samples_per_sec: float
+
+    # Memory (MB)
+    sram_peak_mb: float
+    dram_peak_mb: float
+    total_gpu_memory_mb: float
+
+    # Transfer metrics (μs)
+    h_l_transfer_mean_us: float
+    l_h_transfer_mean_us: float
+
+    # SRAM metrics
+    sram_hit_rate: float
+
+    # Triton kernel metrics
+    triton_sram_probe_latency_us: float
+    triton_dram_probe_latency_us: float
+
+    # GPU metrics
+    gpu_utilization_pct: Optional[float]
+    gpu_power_w: Optional[float]
+    gpu_temperature_c: Optional[float]
+
+    # Derived
+    h_over_l_latency_ratio: float
+    memory_efficiency: float  # useful_compute_time / total_time
+
+
+def _std(values: List[float]) -> float:
+    if len(values) < 2:
+        return 0.0
+    mean = sum(values) / len(values)
+    var = sum((v - mean) ** 2 for v in values) / (len(values) - 1)
+    return math.sqrt(var)
+
+
+def _create_dummy_batch(
+    batch_size: int,
+    seq_len: int,
+    vocab_size: int,
+    device: torch.device,
+) -> Dict[str, torch.Tensor]:
+    """Create a synthetic batch for benchmarking."""
+    return {
+        "inputs": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
+        "labels": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
+        "puzzle_identifiers": torch.arange(batch_size, device=device),
+    }
+
+
+def _get_gpu_metrics() -> Dict[str, Optional[float]]:
+    """Try to read GPU utilization, power, temperature via nvidia-smi."""
+    metrics = {'utilization': None, 'power': None, 'temperature': None}
+    try:
+        import subprocess
+        result = subprocess.run(
+            ['nvidia-smi', '--query-gpu=utilization.gpu,power.draw,temperature.gpu',
+             '--format=csv,noheader,nounits'],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            parts = result.stdout.strip().split(',')
+            if len(parts) >= 3:
+                metrics['utilization'] = float(parts[0].strip())
+                metrics['power'] = float(parts[1].strip())
+                metrics['temperature'] = float(parts[2].strip())
+    except Exception:
+        pass
+    return metrics
+
+
+def _run_triton_latency_probe(
+    batch_size: int,
+    hidden_size: int,
+    device: torch.device,
+    num_iters: int = 100,
+) -> Tuple[float, float]:
+    """Measure SRAM vs DRAM effective latency using Triton probe kernels.
+
+    Returns (sram_latency_us, dram_latency_us).
+    """
+    if not torch.cuda.is_available():
+        return 0.0, 0.0
+
+    # Small tensor → fits in SRAM (L1/L2/registers)
+    sram_data = torch.randn(batch_size, hidden_size, device=device, dtype=torch.float32)
+
+    # Large tensor → forces DRAM access (much larger than L2)
+    dram_size = max(hidden_size, 65536)  # Force spill to DRAM
+    dram_data = torch.randn(batch_size * 64, dram_size, device=device, dtype=torch.float32)
+
+    # Make a contiguous slice for DRAM probe
+    dram_probe_data = dram_data[:batch_size, :hidden_size].contiguous()
+
+    # Warmup
+    triton_memory_latency_probe(sram_data, num_iters=10)
+    triton_memory_latency_probe(dram_probe_data, num_iters=10)
+    torch.cuda.synchronize()
+
+    # SRAM probe
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    triton_memory_latency_probe(sram_data, num_iters=num_iters)
+    end.record()
+    torch.cuda.synchronize()
+    sram_us = start.elapsed_time(end) * 1000  # ms → μs
+
+    # DRAM probe (access scattered to prevent caching)
+    start2 = torch.cuda.Event(enable_timing=True)
+    end2 = torch.cuda.Event(enable_timing=True)
+    start2.record()
+    triton_memory_latency_probe(dram_probe_data, num_iters=num_iters)
+    end2.record()
+    torch.cuda.synchronize()
+    dram_us = start2.elapsed_time(end2) * 1000
+
+    # Cleanup
+    del sram_data, dram_data
+    torch.cuda.empty_cache()
+
+    return sram_us, dram_us
+
+
+def benchmark_tiered_model(
+    batch_size: int = 8,
+    seq_len: int = 64,
+    hidden_size: int = 512,
+    num_heads: int = 8,
+    H_cycles: int = 2,
+    L_cycles: int = 2,
+    H_layers: int = 4,
+    L_layers: int = 4,
+    halt_max_steps: int = 1,
+    warmup: int = 5,
+    iterations: int = 20,
+    device: Optional[torch.device] = None,
+) -> BenchmarkResult:
+    """Benchmark the tiered HRM model."""
+
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    vocab_size = 32
+    config_dict = {
+        'batch_size': batch_size,
+        'seq_len': seq_len,
+        'puzzle_emb_ndim': 0,
+        'num_puzzle_identifiers': batch_size,
+        'vocab_size': vocab_size,
+        'H_cycles': H_cycles,
+        'L_cycles': L_cycles,
+        'H_layers': H_layers,
+        'L_layers': L_layers,
+        'hidden_size': hidden_size,
+        'expansion': 4.0,
+        'num_heads': num_heads,
+        'pos_encodings': 'rope',
+        'halt_max_steps': halt_max_steps,
+        'halt_exploration_prob': 0.0,
+    }
+
+    # Memory manager
+    mem_mgr = MemoryTierManager(device=device, enable_tracking=True)
+
+    # Create model
+    model = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
+    model.eval()
+
+    batch = _create_dummy_batch(batch_size, seq_len, vocab_size, device)
+
+    # ---- Warmup ----
+    with torch.no_grad():
+        for _ in range(warmup):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+            model(carry, batch)
+
+    model.reset_timing()
+    mem_mgr.reset_stats()
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+        torch.cuda.synchronize()
+
+    # ---- Benchmark iterations ----
+    total_latencies_ms = []
+    with torch.no_grad():
+        for _ in range(iterations):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+
+            if torch.cuda.is_available():
+                start_event = torch.cuda.Event(enable_timing=True)
+                end_event = torch.cuda.Event(enable_timing=True)
+                start_event.record()
+
+            t0 = time.perf_counter()
+            model(carry, batch)
+
+            if torch.cuda.is_available():
+                end_event.record()
+                torch.cuda.synchronize()
+                total_latencies_ms.append(start_event.elapsed_time(end_event))
+            else:
+                total_latencies_ms.append((time.perf_counter() - t0) * 1000)
+
+    # ---- Collect results ----
+    timing = model.get_timing_stats()
+    mem_stats = mem_mgr.get_stats()
+
+    # GPU memory
+    total_gpu_mb = 0.0
+    if torch.cuda.is_available():
+        total_gpu_mb = torch.cuda.max_memory_allocated(device) / (1024 * 1024)
+
+    # GPU metrics
+    gpu_metrics = _get_gpu_metrics()
+
+    # Triton latency probe
+    sram_probe_us, dram_probe_us = _run_triton_latency_probe(
+        batch_size, hidden_size, device,
+    )
+
+    # Compute derived metrics
+    l_mean = timing['L_forward_us']['mean_us']
+    h_mean = timing['H_forward_us']['mean_us']
+    ratio = h_mean / l_mean if l_mean > 0 else float('inf')
+
+    total_mean_ms = sum(total_latencies_ms) / len(total_latencies_ms)
+    throughput = batch_size / (total_mean_ms / 1000) if total_mean_ms > 0 else 0
+
+    compute_time = timing['L_forward_us']['total_us'] + timing['H_forward_us']['total_us']
+    transfer_time = timing['H_L_transfer_us']['total_us'] + timing['L_H_transfer_us']['total_us']
+    efficiency = compute_time / (compute_time + transfer_time) if (compute_time + transfer_time) > 0 else 0
+
+    # Build L-level stats
+    l_values = [timing['L_forward_us']['min_us'], timing['L_forward_us']['max_us']]
+
+    result = BenchmarkResult(
+        model_name='HRM_Tiered',
+        batch_size=batch_size,
+        seq_len=seq_len,
+        hidden_size=hidden_size,
+        H_cycles=H_cycles,
+        L_cycles=L_cycles,
+        H_layers=H_layers,
+        L_layers=L_layers,
+        num_iterations=iterations,
+        warmup_iterations=warmup,
+
+        l_level_latency_mean_us=l_mean,
+        l_level_latency_min_us=timing['L_forward_us']['min_us'],
+        l_level_latency_max_us=timing['L_forward_us']['max_us'],
+        l_level_latency_std_us=0.0,
+
+        h_level_latency_mean_us=h_mean,
+        h_level_latency_min_us=timing['H_forward_us']['min_us'],
+        h_level_latency_max_us=timing['H_forward_us']['max_us'],
+        h_level_latency_std_us=0.0,
+
+        total_inference_latency_mean_ms=total_mean_ms,
+        total_inference_latency_min_ms=min(total_latencies_ms),
+        total_inference_latency_max_ms=max(total_latencies_ms),
+        total_inference_latency_std_ms=_std(total_latencies_ms),
+
+        throughput_samples_per_sec=throughput,
+
+        sram_peak_mb=mem_stats['sram']['peak_mb'],
+        dram_peak_mb=mem_stats['dram']['peak_mb'],
+        total_gpu_memory_mb=total_gpu_mb,
+
+        h_l_transfer_mean_us=timing['H_L_transfer_us']['mean_us'],
+        l_h_transfer_mean_us=timing['L_H_transfer_us']['mean_us'],
+
+        sram_hit_rate=mem_stats['sram']['hit_rate'],
+
+        triton_sram_probe_latency_us=sram_probe_us,
+        triton_dram_probe_latency_us=dram_probe_us,
+
+        gpu_utilization_pct=gpu_metrics['utilization'],
+        gpu_power_w=gpu_metrics['power'],
+        gpu_temperature_c=gpu_metrics['temperature'],
+
+        h_over_l_latency_ratio=ratio,
+        memory_efficiency=efficiency,
+    )
+
+    # Cleanup
+    del model, batch, carry
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return result
+
+
+def benchmark_baseline_model(
+    batch_size: int = 8,
+    seq_len: int = 64,
+    hidden_size: int = 512,
+    num_heads: int = 8,
+    H_cycles: int = 2,
+    L_cycles: int = 2,
+    H_layers: int = 4,
+    L_layers: int = 4,
+    halt_max_steps: int = 1,
+    warmup: int = 5,
+    iterations: int = 20,
+    device: Optional[torch.device] = None,
+) -> BenchmarkResult:
+    """Benchmark the original (non-tiered) HRM model."""
+
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    vocab_size = 32
+    config_dict = {
+        'batch_size': batch_size,
+        'seq_len': seq_len,
+        'puzzle_emb_ndim': 0,
+        'num_puzzle_identifiers': batch_size,
+        'vocab_size': vocab_size,
+        'H_cycles': H_cycles,
+        'L_cycles': L_cycles,
+        'H_layers': H_layers,
+        'L_layers': L_layers,
+        'hidden_size': hidden_size,
+        'expansion': 4.0,
+        'num_heads': num_heads,
+        'pos_encodings': 'rope',
+        'halt_max_steps': halt_max_steps,
+        'halt_exploration_prob': 0.0,
+    }
+
+    model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
+    model.eval()
+
+    batch = _create_dummy_batch(batch_size, seq_len, vocab_size, device)
+
+    # ---- Warmup ----
+    with torch.no_grad():
+        for _ in range(warmup):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+            model(carry, batch)
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+        torch.cuda.synchronize()
+
+    # ---- Benchmark ----
+    total_latencies_ms = []
+    with torch.no_grad():
+        for _ in range(iterations):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+
+            if torch.cuda.is_available():
+                start_event = torch.cuda.Event(enable_timing=True)
+                end_event = torch.cuda.Event(enable_timing=True)
+                start_event.record()
+
+            t0 = time.perf_counter()
+            model(carry, batch)
+
+            if torch.cuda.is_available():
+                end_event.record()
+                torch.cuda.synchronize()
+                total_latencies_ms.append(start_event.elapsed_time(end_event))
+            else:
+                total_latencies_ms.append((time.perf_counter() - t0) * 1000)
+
+    total_gpu_mb = 0.0
+    if torch.cuda.is_available():
+        total_gpu_mb = torch.cuda.max_memory_allocated(device) / (1024 * 1024)
+
+    total_mean_ms = sum(total_latencies_ms) / len(total_latencies_ms)
+    throughput = batch_size / (total_mean_ms / 1000) if total_mean_ms > 0 else 0
+
+    gpu_metrics = _get_gpu_metrics()
+
+    result = BenchmarkResult(
+        model_name='HRM_Baseline',
+        batch_size=batch_size,
+        seq_len=seq_len,
+        hidden_size=hidden_size,
+        H_cycles=H_cycles,
+        L_cycles=L_cycles,
+        H_layers=H_layers,
+        L_layers=L_layers,
+        num_iterations=iterations,
+        warmup_iterations=warmup,
+
+        l_level_latency_mean_us=0, l_level_latency_min_us=0,
+        l_level_latency_max_us=0, l_level_latency_std_us=0,
+        h_level_latency_mean_us=0, h_level_latency_min_us=0,
+        h_level_latency_max_us=0, h_level_latency_std_us=0,
+
+        total_inference_latency_mean_ms=total_mean_ms,
+        total_inference_latency_min_ms=min(total_latencies_ms),
+        total_inference_latency_max_ms=max(total_latencies_ms),
+        total_inference_latency_std_ms=_std(total_latencies_ms),
+
+        throughput_samples_per_sec=throughput,
+
+        sram_peak_mb=0, dram_peak_mb=0,
+        total_gpu_memory_mb=total_gpu_mb,
+
+        h_l_transfer_mean_us=0, l_h_transfer_mean_us=0,
+        sram_hit_rate=0,
+
+        triton_sram_probe_latency_us=0, triton_dram_probe_latency_us=0,
+
+        gpu_utilization_pct=gpu_metrics['utilization'],
+        gpu_power_w=gpu_metrics['power'],
+        gpu_temperature_c=gpu_metrics['temperature'],
+
+        h_over_l_latency_ratio=0,
+        memory_efficiency=1.0,
+    )
+
+    del model, batch, carry
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return result
+
+
+def compare_models(
+    batch_sizes: List[int] = [1, 8, 32],
+    seq_lens: List[int] = [64, 128],
+    hidden_size: int = 512,
+    warmup: int = 5,
+    iterations: int = 20,
+    device: Optional[torch.device] = None,
+) -> List[Dict]:
+    """Run comparative benchmark between tiered and baseline HRM."""
+    results = []
+
+    for bs in batch_sizes:
+        for sl in seq_lens:
+            print(f"\n{'='*60}")
+            print(f"  Benchmarking: batch_size={bs}, seq_len={sl}")
+            print(f"{'='*60}")
+
+            print("  → Baseline model...")
+            baseline = benchmark_baseline_model(
+                batch_size=bs, seq_len=sl, hidden_size=hidden_size,
+                warmup=warmup, iterations=iterations, device=device,
+            )
+
+            print("  → Tiered model...")
+            tiered = benchmark_tiered_model(
+                batch_size=bs, seq_len=sl, hidden_size=hidden_size,
+                warmup=warmup, iterations=iterations, device=device,
+            )
+
+            comparison = {
+                'batch_size': bs,
+                'seq_len': sl,
+                'baseline': asdict(baseline),
+                'tiered': asdict(tiered),
+                'speedup': baseline.total_inference_latency_mean_ms / tiered.total_inference_latency_mean_ms if tiered.total_inference_latency_mean_ms > 0 else 0,
+                'memory_savings_mb': baseline.total_gpu_memory_mb - tiered.total_gpu_memory_mb,
+                'throughput_improvement': tiered.throughput_samples_per_sec / baseline.throughput_samples_per_sec if baseline.throughput_samples_per_sec > 0 else 0,
+            }
+            results.append(comparison)
+
+            # Print summary
+            print(f"\n  Results:")
+            print(f"    Baseline latency: {baseline.total_inference_latency_mean_ms:.2f} ms")
+            print(f"    Tiered latency:   {tiered.total_inference_latency_mean_ms:.2f} ms")
+            print(f"    Speedup:          {comparison['speedup']:.2f}x")
+            print(f"    H/L ratio:        {tiered.h_over_l_latency_ratio:.2f}x")
+            print(f"    SRAM probe:       {tiered.triton_sram_probe_latency_us:.1f} μs")
+            print(f"    DRAM probe:       {tiered.triton_dram_probe_latency_us:.1f} μs")
+
+    return results
+
+
+def print_results_table(results: List[BenchmarkResult]):
+    """Pretty-print benchmark results as a table."""
+    header = (
+        f"{'Model':<15} {'BS':>4} {'Seq':>5} "
+        f"{'Latency(ms)':>12} {'Throughput':>12} "
+        f"{'L_lat(μs)':>10} {'H_lat(μs)':>10} {'H/L':>6} "
+        f"{'GPU_MB':>8} {'Efficiency':>10}"
+    )
+    print(f"\n{'='*len(header)}")
+    print(header)
+    print(f"{'='*len(header)}")
+
+    for r in results:
+        print(
+            f"{r.model_name:<15} {r.batch_size:>4} {r.seq_len:>5} "
+            f"{r.total_inference_latency_mean_ms:>12.2f} {r.throughput_samples_per_sec:>12.1f} "
+            f"{r.l_level_latency_mean_us:>10.1f} {r.h_level_latency_mean_us:>10.1f} {r.h_over_l_latency_ratio:>6.2f} "
+            f"{r.total_gpu_memory_mb:>8.1f} {r.memory_efficiency:>10.3f}"
+        )
+    print()
+
+
+def generate_plots(results: List[Dict], output_dir: str = "benchmark_results"):
+    """Generate comparison plots using matplotlib."""
+    try:
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        import numpy as np
+    except ImportError:
+        print("matplotlib not available — skipping plot generation.")
+        return
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # ---- Plot 1: Latency comparison ----
+    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
+    fig.suptitle('HRM SRAM/DRAM Tiering — Benchmark Results', fontsize=14, fontweight='bold')
+
+    configs = [f"bs={r['batch_size']}\nseq={r['seq_len']}" for r in results]
+    baseline_lat = [r['baseline']['total_inference_latency_mean_ms'] for r in results]
+    tiered_lat = [r['tiered']['total_inference_latency_mean_ms'] for r in results]
+
+    x = np.arange(len(configs))
+    w = 0.35
+
+    ax = axes[0]
+    ax.bar(x - w/2, baseline_lat, w, label='Baseline', color='#e74c3c', alpha=0.8)
+    ax.bar(x + w/2, tiered_lat, w, label='Tiered (SRAM/DRAM)', color='#2ecc71', alpha=0.8)
+    ax.set_xlabel('Configuration')
+    ax.set_ylabel('Latency (ms)')
+    ax.set_title('Inference Latency')
+    ax.set_xticks(x)
+    ax.set_xticklabels(configs, fontsize=8)
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+
+    # ---- Plot 2: Throughput ----
+    ax = axes[1]
+    baseline_tp = [r['baseline']['throughput_samples_per_sec'] for r in results]
+    tiered_tp = [r['tiered']['throughput_samples_per_sec'] for r in results]
+    ax.bar(x - w/2, baseline_tp, w, label='Baseline', color='#e74c3c', alpha=0.8)
+    ax.bar(x + w/2, tiered_tp, w, label='Tiered', color='#2ecc71', alpha=0.8)
+    ax.set_xlabel('Configuration')
+    ax.set_ylabel('Samples/sec')
+    ax.set_title('Throughput')
+    ax.set_xticks(x)
+    ax.set_xticklabels(configs, fontsize=8)
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+
+    # ---- Plot 3: H vs L latency (tiered only) ----
+    ax = axes[2]
+    l_lat = [r['tiered']['l_level_latency_mean_us'] for r in results]
+    h_lat = [r['tiered']['h_level_latency_mean_us'] for r in results]
+    ax.bar(x - w/2, l_lat, w, label='L-level (SRAM)', color='#3498db', alpha=0.8)
+    ax.bar(x + w/2, h_lat, w, label='H-level (DRAM)', color='#e67e22', alpha=0.8)
+    ax.set_xlabel('Configuration')
+    ax.set_ylabel('Latency (μs)')
+    ax.set_title('Per-Module Latency')
+    ax.set_xticks(x)
+    ax.set_xticklabels(configs, fontsize=8)
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+
+    plt.tight_layout()
+    plot_path = os.path.join(output_dir, 'benchmark_comparison.png')
+    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"  Plot saved: {plot_path}")
+
+    # ---- Plot 4: Memory breakdown ----
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+    fig.suptitle('Memory Analysis', fontsize=14, fontweight='bold')
+
+    ax = axes[0]
+    gpu_mem_baseline = [r['baseline']['total_gpu_memory_mb'] for r in results]
+    gpu_mem_tiered = [r['tiered']['total_gpu_memory_mb'] for r in results]
+    ax.bar(x - w/2, gpu_mem_baseline, w, label='Baseline', color='#e74c3c', alpha=0.8)
+    ax.bar(x + w/2, gpu_mem_tiered, w, label='Tiered', color='#2ecc71', alpha=0.8)
+    ax.set_xlabel('Configuration')
+    ax.set_ylabel('GPU Memory (MB)')
+    ax.set_title('Total GPU Memory')
+    ax.set_xticks(x)
+    ax.set_xticklabels(configs, fontsize=8)
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+
+    ax = axes[1]
+    triton_sram = [r['tiered']['triton_sram_probe_latency_us'] for r in results]
+    triton_dram = [r['tiered']['triton_dram_probe_latency_us'] for r in results]
+    ax.bar(x - w/2, triton_sram, w, label='SRAM Probe', color='#3498db', alpha=0.8)
+    ax.bar(x + w/2, triton_dram, w, label='DRAM Probe', color='#e67e22', alpha=0.8)
+    ax.set_xlabel('Configuration')
+    ax.set_ylabel('Latency (μs)')
+    ax.set_title('Triton Memory Probe Latency')
+    ax.set_xticks(x)
+    ax.set_xticklabels(configs, fontsize=8)
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+
+    plt.tight_layout()
+    plot_path = os.path.join(output_dir, 'memory_analysis.png')
+    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"  Plot saved: {plot_path}")

benchmark_results/comparison_results.json

@@ -0,0 +1,119 @@
+{
+  "baseline": {
+    "latency_ms": 13.310442765553793,
+    "latency_std": 0.3072074317910763,
+    "throughput": 2404.1273880695444,
+    "peak_gpu_mb": 292.287488,
+    "params_m": 27.29677
+  },
+  "tiered": {
+    "latency_ms": 11.866764799753826,
+    "latency_std": 0.02685813995812233,
+    "throughput": 2696.606913508881,
+    "peak_gpu_mb": 313.03936,
+    "params_m": 27.29677
+  },
+  "speedup": 1.1216572494830197,
+  "sweep": {
+    "baseline": [
+      {
+        "latency_ms": 13.55545597076416,
+        "latency_std": 0.17929869581870853,
+        "throughput": 73.77103375620548,
+        "peak_gpu_mb": 140.93056,
+        "params_m": 27.29677,
+        "batch_size": 1
+      },
+      {
+        "latency_ms": 13.13232650756836,
+        "latency_std": 0.41090034980335216,
+        "throughput": 304.5918784988128,
+        "peak_gpu_mb": 162.348544,
+        "params_m": 27.29677,
+        "batch_size": 4
+      },
+      {
+        "latency_ms": 12.639206314086914,
+        "latency_std": 0.34697112347890935,
+        "throughput": 632.9511364241022,
+        "peak_gpu_mb": 189.266432,
+        "params_m": 27.29677,
+        "batch_size": 8
+      },
+      {
+        "latency_ms": 12.491708850860595,
+        "latency_std": 0.013537496997456985,
+        "throughput": 1280.8495771895696,
+        "peak_gpu_mb": 223.744512,
+        "params_m": 27.29677,
+        "batch_size": 16
+      },
+      {
+        "latency_ms": 12.555190467834473,
+        "latency_std": 0.01494085745352898,
+        "throughput": 2548.7466782747565,
+        "peak_gpu_mb": 268.796416,
+        "params_m": 27.29677,
+        "batch_size": 32
+      },
+      {
+        "latency_ms": 13.44706563949585,
+        "latency_std": 0.013219697890277465,
+        "throughput": 4759.402661947552,
+        "peak_gpu_mb": 350.675456,
+        "params_m": 27.29677,
+        "batch_size": 64
+      }
+    ],
+    "tiered": [
+      {
+        "latency_ms": 11.769158458709716,
+        "latency_std": 0.016388604477900706,
+        "throughput": 84.96784230651208,
+        "peak_gpu_mb": 158.135808,
+        "params_m": 27.29677,
+        "batch_size": 1
+      },
+      {
+        "latency_ms": 12.082194995880126,
+        "latency_std": 0.019275075409013288,
+        "throughput": 331.0656715409699,
+        "peak_gpu_mb": 180.051456,
+        "params_m": 27.29677,
+        "batch_size": 4
+      },
+      {
+        "latency_ms": 11.751276779174805,
+        "latency_std": 0.0175825841180981,
+        "throughput": 680.7770891906243,
+        "peak_gpu_mb": 206.969344,
+        "params_m": 27.29677,
+        "batch_size": 8
+      },
+      {
+        "latency_ms": 11.727971076965332,
+        "latency_std": 0.014815964207941113,
+        "throughput": 1364.2598446908924,
+        "peak_gpu_mb": 241.291776,
+        "params_m": 27.29677,
+        "batch_size": 16
+      },
+      {
+        "latency_ms": 11.87640323638916,
+        "latency_std": 0.021528091812880033,
+        "throughput": 2694.418450019647,
+        "peak_gpu_mb": 289.014272,
+        "params_m": 27.29677,
+        "batch_size": 32
+      },
+      {
+        "latency_ms": 12.873033428192139,
+        "latency_std": 0.021238526686589362,
+        "throughput": 4971.633170767585,
+        "peak_gpu_mb": 373.481984,
+        "params_m": 27.29677,
+        "batch_size": 64
+      }
+    ]
+  }
+}

benchmark_results/eval_dummy.json

@@ -0,0 +1,24 @@
+[
+  {
+    "model_name": "Baseline",
+    "batch_size": 32,
+    "seq_len": 81,
+    "hidden_size": 512,
+    "param_count": 27296770,
+    "latency_mean_ms": 12.367569541931152,
+    "latency_std_ms": 0.03661433936324098,
+    "throughput_sps": 2587.4121743570417,
+    "gpu_memory_mb": 174.16943359375
+  },
+  {
+    "model_name": "Tiered",
+    "batch_size": 32,
+    "seq_len": 81,
+    "hidden_size": 512,
+    "param_count": 27296770,
+    "latency_mean_ms": 11.73528323173523,
+    "latency_std_ms": 0.06700267067328709,
+    "throughput_sps": 2726.81957206314,
+    "gpu_memory_mb": 193.54443359375
+  }
+]

benchmark_results/results.json

@@ -0,0 +1,500 @@
+[
+  {
+    "batch_size": 1,
+    "seq_len": 64,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 1,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.233289670944213,
+      "total_inference_latency_min_ms": 12.15283203125,
+      "total_inference_latency_max_ms": 12.321791648864746,
+      "total_inference_latency_std_ms": 0.05471266632900291,
+      "throughput_samples_per_sec": 81.74416096556112,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 116.29345703125,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 23.0,
+      "gpu_power_w": 62.07,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 1,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1864.8816029230754,
+      "l_level_latency_min_us": 1843.9680337905884,
+      "l_level_latency_max_us": 1886.2080574035645,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1861.43679022789,
+      "h_level_latency_min_us": 1850.3680229187012,
+      "h_level_latency_max_us": 1875.9679794311523,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.600086498260499,
+      "total_inference_latency_min_ms": 11.553088188171387,
+      "total_inference_latency_max_ms": 11.630592346191406,
+      "total_inference_latency_std_ms": 0.022018860122281114,
+      "throughput_samples_per_sec": 86.20625373353516,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 132.66845703125,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 9.9823999684304,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 820.2239871025085,
+      "triton_dram_probe_latency_us": 53.247999399900436,
+      "gpu_utilization_pct": 19.0,
+      "gpu_power_w": 90.58,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9981527981777578,
+      "memory_efficiency": 0.9986629635146693
+    },
+    "speedup": 1.0545860733692518,
+    "memory_savings_mb": -16.375,
+    "throughput_improvement": 1.0545860733692518
+  },
+  {
+    "batch_size": 1,
+    "seq_len": 128,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 1,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.75871524810791,
+      "total_inference_latency_min_ms": 12.291104316711426,
+      "total_inference_latency_max_ms": 13.331456184387207,
+      "total_inference_latency_std_ms": 0.491103465039144,
+      "throughput_samples_per_sec": 78.37779749401476,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 133.57763671875,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 22.0,
+      "gpu_power_w": 91.94,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 1,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1880.560537179311,
+      "l_level_latency_min_us": 1853.6959886550903,
+      "l_level_latency_max_us": 1905.6639671325684,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1872.9376077651978,
+      "h_level_latency_min_us": 1849.2159843444824,
+      "h_level_latency_max_us": 1890.3039693832397,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.684232091903686,
+      "total_inference_latency_min_ms": 11.551487922668457,
+      "total_inference_latency_max_ms": 11.770879745483398,
+      "total_inference_latency_std_ms": 0.06819341259194193,
+      "throughput_samples_per_sec": 85.58542762026497,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 150.07763671875,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 10.091200051829219,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 47.10400104522705,
+      "triton_dram_probe_latency_us": 38.94399851560593,
+      "gpu_utilization_pct": 14.0,
+      "gpu_power_w": 95.82,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9959464589076472,
+      "memory_efficiency": 0.9986589251294146
+    },
+    "speedup": 1.0919601003945105,
+    "memory_savings_mb": -16.5,
+    "throughput_improvement": 1.0919601003945105
+  },
+  {
+    "batch_size": 8,
+    "seq_len": 64,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 8,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.226089572906494,
+      "total_inference_latency_min_ms": 12.181280136108398,
+      "total_inference_latency_max_ms": 12.269696235656738,
+      "total_inference_latency_std_ms": 0.02198156350662233,
+      "throughput_samples_per_sec": 654.3384090468568,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 158.31396484375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 25.0,
+      "gpu_power_w": 100.16,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 8,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1842.8634683291118,
+      "l_level_latency_min_us": 1817.6000118255615,
+      "l_level_latency_max_us": 1919.103980064392,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1826.6352117061615,
+      "h_level_latency_min_us": 1815.551996231079,
+      "h_level_latency_max_us": 1835.8080387115479,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.43253436088562,
+      "total_inference_latency_min_ms": 11.369471549987793,
+      "total_inference_latency_max_ms": 11.523072242736816,
+      "total_inference_latency_std_ms": 0.04659291036998755,
+      "throughput_samples_per_sec": 699.7573545346668,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 175.56396484375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 9.712000098079443,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 46.08000069856644,
+      "triton_dram_probe_latency_us": 40.95999896526337,
+      "gpu_utilization_pct": 26.0,
+      "gpu_power_w": 104.38,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9911939995003188,
+      "memory_efficiency": 0.9986813194349485
+    },
+    "speedup": 1.0694120119801154,
+    "memory_savings_mb": -17.25,
+    "throughput_improvement": 1.0694120119801154
+  },
+  {
+    "batch_size": 8,
+    "seq_len": 128,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 8,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.09585280418396,
+      "total_inference_latency_min_ms": 12.034048080444336,
+      "total_inference_latency_max_ms": 12.1693115234375,
+      "total_inference_latency_std_ms": 0.03280009762608858,
+      "throughput_samples_per_sec": 661.3837097317187,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 186.99365234375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 37.0,
+      "gpu_power_w": 112.76,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 8,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1827.0586788654327,
+      "l_level_latency_min_us": 1786.8800163269043,
+      "l_level_latency_max_us": 1932.2880506515503,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1811.8864059448242,
+      "h_level_latency_min_us": 1794.7200536727905,
+      "h_level_latency_max_us": 1830.7520151138306,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.355580854415894,
+      "total_inference_latency_min_ms": 11.258879661560059,
+      "total_inference_latency_max_ms": 11.465632438659668,
+      "total_inference_latency_std_ms": 0.05902572924729382,
+      "throughput_samples_per_sec": 704.4994089306321,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 204.24365234375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 9.579200064763427,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 47.16800153255463,
+      "triton_dram_probe_latency_us": 39.93599861860275,
+      "gpu_utilization_pct": 33.0,
+      "gpu_power_w": 118.67,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9916957933009408,
+      "memory_efficiency": 0.9986882554925227
+    },
+    "speedup": 1.0651901438824412,
+    "memory_savings_mb": -17.25,
+    "throughput_improvement": 1.0651901438824412
+  },
+  {
+    "batch_size": 32,
+    "seq_len": 64,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 32,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.013897466659547,
+      "total_inference_latency_min_ms": 11.935744285583496,
+      "total_inference_latency_max_ms": 12.108991622924805,
+      "total_inference_latency_std_ms": 0.04052232974770088,
+      "throughput_samples_per_sec": 2663.581913263787,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 228.25927734375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 50.0,
+      "gpu_power_w": 137.35,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 32,
+      "seq_len": 64,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1848.9973326524098,
+      "l_level_latency_min_us": 1829.6960592269897,
+      "l_level_latency_max_us": 1885.4399919509888,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1829.4048011302948,
+      "h_level_latency_min_us": 1820.9600448608398,
+      "h_level_latency_max_us": 1836.0320329666138,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.482071924209595,
+      "total_inference_latency_min_ms": 11.442208290100098,
+      "total_inference_latency_max_ms": 11.52511978149414,
+      "total_inference_latency_std_ms": 0.022921052595205035,
+      "throughput_samples_per_sec": 2786.953453281283,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 247.38427734375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 9.636800037696958,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 46.08000069856644,
+      "triton_dram_probe_latency_us": 40.95999896526337,
+      "gpu_utilization_pct": 45.0,
+      "gpu_power_w": 162.51,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9894036994125842,
+      "memory_efficiency": 0.9986952672353347
+    },
+    "speedup": 1.046317907252315,
+    "memory_savings_mb": -19.125,
+    "throughput_improvement": 1.046317907252315
+  },
+  {
+    "batch_size": 32,
+    "seq_len": 128,
+    "baseline": {
+      "model_name": "HRM_Baseline",
+      "batch_size": 32,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 0,
+      "l_level_latency_min_us": 0,
+      "l_level_latency_max_us": 0,
+      "l_level_latency_std_us": 0,
+      "h_level_latency_mean_us": 0,
+      "h_level_latency_min_us": 0,
+      "h_level_latency_max_us": 0,
+      "h_level_latency_std_us": 0,
+      "total_inference_latency_mean_ms": 12.027006340026855,
+      "total_inference_latency_min_ms": 11.971232414245605,
+      "total_inference_latency_max_ms": 12.089983940124512,
+      "total_inference_latency_std_ms": 0.02667638763107568,
+      "throughput_samples_per_sec": 2660.678733784433,
+      "sram_peak_mb": 0,
+      "dram_peak_mb": 0,
+      "total_gpu_memory_mb": 292.13427734375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 0,
+      "sram_hit_rate": 0,
+      "triton_sram_probe_latency_us": 0,
+      "triton_dram_probe_latency_us": 0,
+      "gpu_utilization_pct": 84.0,
+      "gpu_power_w": 210.32,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0,
+      "memory_efficiency": 1.0
+    },
+    "tiered": {
+      "model_name": "HRM_Tiered",
+      "batch_size": 32,
+      "seq_len": 128,
+      "hidden_size": 512,
+      "H_cycles": 2,
+      "L_cycles": 2,
+      "H_layers": 4,
+      "L_layers": 4,
+      "num_iterations": 20,
+      "warmup_iterations": 5,
+      "l_level_latency_mean_us": 1828.7717362244923,
+      "l_level_latency_min_us": 1806.175947189331,
+      "l_level_latency_max_us": 1862.6240491867065,
+      "l_level_latency_std_us": 0.0,
+      "h_level_latency_mean_us": 1824.127995967865,
+      "h_level_latency_min_us": 1806.3360452651978,
+      "h_level_latency_max_us": 1875.7760524749756,
+      "h_level_latency_std_us": 0.0,
+      "total_inference_latency_mean_ms": 11.574174451828004,
+      "total_inference_latency_min_ms": 11.528096199035645,
+      "total_inference_latency_max_ms": 11.638784408569336,
+      "total_inference_latency_std_ms": 0.03262226881630495,
+      "throughput_samples_per_sec": 2764.7760220985765,
+      "sram_peak_mb": 0.0,
+      "dram_peak_mb": 0.0,
+      "total_gpu_memory_mb": 312.38427734375,
+      "h_l_transfer_mean_us": 0,
+      "l_h_transfer_mean_us": 9.51200001873076,
+      "sram_hit_rate": 0.0,
+      "triton_sram_probe_latency_us": 49.15200173854828,
+      "triton_dram_probe_latency_us": 40.70400074124336,
+      "gpu_utilization_pct": 78.0,
+      "gpu_power_w": 245.59,
+      "gpu_temperature_c": null,
+      "h_over_l_latency_ratio": 0.9974607327067432,
+      "memory_efficiency": 0.9987005384933674
+    },
+    "speedup": 1.0391243358291815,
+    "memory_savings_mb": -20.25,
+    "throughput_improvement": 1.0391243358291815
+  }
+]

benchmark_results/run_nsys_profiler.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""
+NVIDIA Nsight Systems Profiler for HRM Memory Tiering.
+
+This script runs a few iterations of both the baseline and tiered models
+and is designed to be executed via `nsys profile`.
+
+Usage:
+  nsys profile -t cuda,nvtx --stats=true --force-overwrite=true -o hrm_profile python run_nsys_profiler.py
+"""
+
+import torch
+import torch.cuda.nvtx as nvtx
+import argparse
+
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
+from models.hrm.hrm_tiered import HRM_Tiered
+from models.memory_tier import MemoryTierManager
+from run_training_comparison import DummyLossModel, create_dummy_batch
+
+def profile_model(model_name, model, batch, iterations, device):
+    print(f"Profiling {model_name}...")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    model.train()
+
+    # Warmup
+    for _ in range(2):
+        optimizer.zero_grad()
+        carry = model.initial_carry(batch)
+        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+        carry.steps = carry.steps.to(device)
+        carry.halted = carry.halted.to(device)
+        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+        
+        _, loss, _, _, _ = model(carry, batch, return_keys=[])
+        loss.backward()
+        optimizer.step()
+
+    torch.cuda.synchronize()
+
+    # Profiling Phase
+    with torch.autograd.profiler.emit_nvtx():
+        nvtx.range_push(f"{model_name}_Training_Loop")
+        for i in range(iterations):
+            nvtx.range_push(f"Iteration_{i}")
+            optimizer.zero_grad()
+            
+            nvtx.range_push("Forward_Pass")
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+
+            _, loss, _, _, _ = model(carry, batch, return_keys=[])
+            nvtx.range_pop() # End Forward
+            
+            nvtx.range_push("Backward_Pass")
+            loss.backward()
+            optimizer.step()
+            nvtx.range_pop() # End Backward
+            
+            nvtx.range_pop() # End Iteration
+        nvtx.range_pop() # End Loop
+
+    torch.cuda.synchronize()
+    print(f"Finished {model_name}.\n")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch-size', type=int, default=16)
+    parser.add_argument('--seq-len', type=int, default=128)
+    parser.add_argument('--hidden-size', type=int, default=1024)
+    parser.add_argument('--iterations', type=int, default=5)
+    args = parser.parse_args()
+
+    device = torch.device('cuda')
+    vocab_size = 32
+
+    config_dict = {
+        'batch_size': args.batch_size,
+        'seq_len': args.seq_len,
+        'puzzle_emb_ndim': 0,
+        'num_puzzle_identifiers': args.batch_size,
+        'vocab_size': vocab_size,
+        'H_cycles': 2,
+        'L_cycles': 2,
+        'H_layers': 4,
+        'L_layers': 4,
+        'hidden_size': args.hidden_size,
+        'expansion': 4.0,
+        'num_heads': 8,
+        'pos_encodings': 'rope',
+        'halt_max_steps': 1,
+        'halt_exploration_prob': 0.0,
+    }
+
+    batch = create_dummy_batch(args.batch_size, args.seq_len, vocab_size, device)
+
+    # 1. Baseline
+    baseline = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
+    baseline_wrapped = DummyLossModel(baseline)
+    profile_model("HRM_Baseline", baseline_wrapped, batch, args.iterations, device)
+    del baseline_wrapped, baseline
+    torch.cuda.empty_cache()
+
+    # 2. Tiered
+    mem_mgr = MemoryTierManager(device=device, enable_tracking=False)
+    tiered = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
+    tiered_wrapped = DummyLossModel(tiered)
+    profile_model("HRM_Tiered", tiered_wrapped, batch, args.iterations, device)
+
+if __name__ == "__main__":
+    main()

benchmark_results/trained_model_results.json

@@ -0,0 +1,18 @@
+{
+  "baseline": {
+    "accuracy": {},
+    "latency": {
+      "latency_ms": 117.27158050537109,
+      "latency_std": 0.10221427702082414,
+      "throughput": 3274.4506243131314
+    }
+  },
+  "tiered": {
+    "accuracy": {},
+    "latency": {
+      "latency_ms": 95.14221649169922,
+      "latency_std": 0.010479617132043349,
+      "throughput": 4036.0632131531484
+    }
+  }
+}

cleanup.sh

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+# cleanup.sh — Purge venv, downloaded data, checkpoints, caches, and wandb logs
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+echo "=========================================="
+echo "  HRM_optimised — Cleanup"
+echo "=========================================="
+
+# Deactivate venv if currently active
+if [ -n "${VIRTUAL_ENV:-}" ]; then
+    echo "  Deactivating current virtual environment..."
+    deactivate 2>/dev/null || true
+fi
+
+# 1. Remove virtual environment
+if [ -d "venv" ]; then
+    echo "[1/6] Removing virtual environment (venv/)..."
+    rm -rf venv
+else
+    echo "[1/6] No venv found, skipping."
+fi
+
+# 2. Remove downloaded/generated datasets
+if [ -d "data" ]; then
+    echo "[2/6] Removing generated datasets (data/)..."
+    rm -rf data
+else
+    echo "[2/6] No data/ directory found, skipping."
+fi
+
+# 3. Remove raw downloaded data from HuggingFace
+if [ -d "dataset/raw-data" ]; then
+    echo "[3/6] Removing raw dataset downloads (dataset/raw-data/)..."
+    rm -rf dataset/raw-data
+else
+    echo "[3/6] No dataset/raw-data/ found, skipping."
+fi
+
+# 4. Remove training checkpoints
+if [ -d "checkpoints" ]; then
+    echo "[4/6] Removing training checkpoints (checkpoints/)..."
+    rm -rf checkpoints
+else
+    echo "[4/6] No checkpoints/ directory found, skipping."
+fi
+
+# 5. Remove wandb logs
+if [ -d "wandb" ]; then
+    echo "[5/6] Removing W&B logs (wandb/)..."
+    rm -rf wandb
+else
+    echo "[5/6] No wandb/ directory found, skipping."
+fi
+
+# 6. Remove Python caches
+echo "[6/6] Removing __pycache__ directories..."
+find . -type d -name "__pycache__" -not -path "./.git/*" -exec rm -rf {} + 2>/dev/null || true
+find . -type d -name ".hydra" -exec rm -rf {} + 2>/dev/null || true
+
+echo ""
+echo "=========================================="
+echo "  ✅ Cleanup complete!"
+echo "  Run ./startup.sh to set up a fresh environment."
+echo "=========================================="

compare_models.py

@@ -0,0 +1,299 @@
+#!/usr/bin/env python3
+"""
+Compare Baseline vs Tiered HRM — Multi-GPU Benchmark + Comparison Plots.
+
+Usage:
+    source venv/bin/activate
+    python compare_models.py                              # quick compare
+    python compare_models.py --sweep                      # batch-size sweep + plots
+    python compare_models.py --iterations 50 --sweep      # thorough
+"""
+
+import argparse
+import json
+import os
+import sys
+from dataclasses import dataclass, asdict
+
+import torch
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from models.memory_tier import MemoryTierManager
+from models.hrm.hrm_tiered import HRM_Tiered
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
+
+
+# ═══════════════════════════════════════════════════════════
+#  Helpers
+# ═══════════════════════════════════════════════════════════
+
+def make_config(batch_size, seq_len, hidden_size, num_heads):
+    return {
+        "batch_size": batch_size, "seq_len": seq_len,
+        "puzzle_emb_ndim": 0, "num_puzzle_identifiers": batch_size,
+        "vocab_size": 32,
+        "H_cycles": 2, "L_cycles": 2, "H_layers": 4, "L_layers": 4,
+        "hidden_size": hidden_size, "expansion": 4.0,
+        "num_heads": num_heads, "pos_encodings": "rope",
+        "halt_max_steps": 1, "halt_exploration_prob": 0.0,
+    }
+
+
+def make_batch(batch_size, seq_len, device):
+    return {
+        "inputs": torch.randint(0, 31, (batch_size, seq_len), device=device),
+        "labels": torch.randint(0, 31, (batch_size, seq_len), device=device),
+        "puzzle_identifiers": torch.arange(batch_size, device=device),
+    }
+
+
+def build_model(arch, config_dict, device):
+    """Build using the wrapper classes (same pattern as eval_dummy.py)."""
+    if arch == "tiered":
+        mm = MemoryTierManager(device=device, enable_tracking=True)
+        model = HRM_Tiered(config_dict, memory_manager=mm).to(device)
+    else:
+        model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def benchmark(model, batch, device, warmup=5, iterations=20):
+    """Time forward pass, return dict of metrics."""
+    bs = batch["inputs"].shape[0]
+
+    # Warmup
+    for _ in range(warmup):
+        carry = model.initial_carry(batch)
+        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+        carry.steps = carry.steps.to(device)
+        carry.halted = carry.halted.to(device)
+        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+        model(carry, batch)
+
+    torch.cuda.reset_peak_memory_stats(device)
+    torch.cuda.synchronize()
+
+    latencies = []
+    for _ in range(iterations):
+        carry = model.initial_carry(batch)
+        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+        carry.steps = carry.steps.to(device)
+        carry.halted = carry.halted.to(device)
+        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        model(carry, batch)
+        end.record()
+        torch.cuda.synchronize()
+        latencies.append(start.elapsed_time(end))
+
+    lat = np.array(latencies)
+    return {
+        "latency_ms": float(np.mean(lat)),
+        "latency_std": float(np.std(lat)),
+        "throughput": float(bs / (np.mean(lat) / 1000)),
+        "peak_gpu_mb": float(torch.cuda.max_memory_allocated(device) / 1e6),
+        "params_m": sum(p.numel() for p in model.parameters()) / 1e6,
+    }
+
+
+# ═══════════════════════════════════════════════════════════
+#  Plotting
+# ═══════════════════════════════════════════════════════════
+
+def create_plots(base_res, tier_res, sweep_data, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+
+    c_base, c_tier = "#4A90D9", "#E85D75"
+    bg, text, grid = "#1a1a2e", "#e0e0e0", "#333355"
+
+    plt.rcParams.update({
+        "figure.facecolor": bg, "axes.facecolor": "#16213e",
+        "axes.edgecolor": grid, "axes.labelcolor": text,
+        "text.color": text, "xtick.color": text, "ytick.color": text,
+        "grid.color": grid, "grid.alpha": 0.3,
+        "font.family": "sans-serif", "font.size": 11,
+    })
+
+    n_plots = 6 if sweep_data else 5
+    fig = plt.figure(figsize=(16, 10))
+    fig.suptitle("HRM Baseline vs Tiered (SRAM/DRAM) Comparison",
+                 fontsize=18, fontweight="bold", y=0.98)
+    gs = GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.35)
+    labels = ["Baseline", "Tiered"]
+
+    def bar_plot(ax, title, ylabel, vals, fmt=".2f"):
+        bars = ax.bar(labels, vals, color=[c_base, c_tier],
+                      edgecolor="white", linewidth=0.5, width=0.5)
+        ax.set_title(title, fontweight="bold")
+        ax.set_ylabel(ylabel)
+        for b, v in zip(bars, vals):
+            ax.text(b.get_x() + b.get_width()/2, b.get_height() * 1.02,
+                    f"{v:{fmt}}", ha="center", fontsize=10, color=text)
+        ax.grid(axis="y")
+
+    # 1. Latency
+    bar_plot(fig.add_subplot(gs[0, 0]), "Inference Latency", "ms",
+             [base_res["latency_ms"], tier_res["latency_ms"]])
+
+    # 2. Throughput
+    bar_plot(fig.add_subplot(gs[0, 1]), "Throughput", "samples/sec",
+             [base_res["throughput"], tier_res["throughput"]], fmt=".0f")
+
+    # 3. GPU Memory
+    bar_plot(fig.add_subplot(gs[0, 2]), "Peak GPU Memory", "MB",
+             [base_res["peak_gpu_mb"], tier_res["peak_gpu_mb"]], fmt=".0f")
+
+    # 4. Parameters
+    bar_plot(fig.add_subplot(gs[1, 0]), "Model Parameters", "Millions",
+             [base_res["params_m"], tier_res["params_m"]], fmt=".1f")
+
+    # 5. Summary text
+    ax5 = fig.add_subplot(gs[1, 1])
+    speedup = base_res["latency_ms"] / tier_res["latency_ms"]
+    mem_diff = tier_res["peak_gpu_mb"] - base_res["peak_gpu_mb"]
+    tp_gain = (tier_res["throughput"] / base_res["throughput"] - 1) * 100
+    summary = (
+        f"Speedup: {speedup:.2f}x\n"
+        f"Throughput: {tp_gain:+.1f}%\n"
+        f"Memory Δ: {mem_diff:+.0f} MB\n"
+        f"Params: identical"
+    )
+    ax5.text(0.5, 0.5, summary, transform=ax5.transAxes,
+             ha="center", va="center", fontsize=14, fontfamily="monospace",
+             bbox=dict(boxstyle="round,pad=0.5", facecolor="#0f3460", alpha=0.8))
+    ax5.set_title("Summary", fontweight="bold")
+    ax5.axis("off")
+
+    # 6. Sweep plot
+    ax6 = fig.add_subplot(gs[1, 2])
+    if sweep_data:
+        bs_list = [s["batch_size"] for s in sweep_data["baseline"]]
+        ax6.plot(bs_list, [s["latency_ms"] for s in sweep_data["baseline"]],
+                 "o-", color=c_base, label="Baseline", linewidth=2, markersize=6)
+        ax6.plot(bs_list, [s["latency_ms"] for s in sweep_data["tiered"]],
+                 "s-", color=c_tier, label="Tiered", linewidth=2, markersize=6)
+        ax6.set_xlabel("Batch Size")
+        ax6.set_ylabel("Latency (ms)")
+        ax6.set_title("Latency vs Batch Size", fontweight="bold")
+        ax6.legend(facecolor="#16213e", edgecolor=grid)
+        ax6.grid(True)
+    else:
+        ax6.text(0.5, 0.5, "Run with --sweep\nfor batch size\ncomparison",
+                 transform=ax6.transAxes, ha="center", va="center", fontsize=12)
+        ax6.set_title("Latency vs Batch Size", fontweight="bold")
+        ax6.axis("off")
+
+    path = os.path.join(output_dir, "model_comparison.png")
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"  Plot saved → {path}")
+    return path
+
+
+# ═══════════════════════════════════════════════════════════
+#  Main
+# ═══════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser(description="Compare Baseline vs Tiered HRM")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--seq-len", type=int, default=81)
+    parser.add_argument("--hidden-size", type=int, default=512)
+    parser.add_argument("--num-heads", type=int, default=8)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--iterations", type=int, default=20)
+    parser.add_argument("--sweep", action="store_true", help="Batch-size sweep")
+    parser.add_argument("--output-dir", type=str, default="benchmark_results")
+    args = parser.parse_args()
+
+    device = torch.device("cuda")
+    cfg = make_config(args.batch_size, args.seq_len, args.hidden_size, args.num_heads)
+
+    print("=" * 64)
+    print("  HRM Model Comparison: Baseline vs Tiered (SRAM/DRAM)")
+    print(f"  Device: {torch.cuda.get_device_name(0)}")
+    print(f"  Config: bs={args.batch_size}, seq={args.seq_len}, hidden={args.hidden_size}")
+    print("=" * 64)
+
+    # ── Build ──
+    print("\n  Building Baseline...")
+    base_model = build_model("baseline", cfg, device)
+    batch = make_batch(args.batch_size, args.seq_len, device)
+
+    print("  Building Tiered...")
+    tier_model = build_model("tiered", cfg, device)
+
+    # ── Benchmark ──
+    print(f"\n  Benchmarking Baseline ({args.iterations} iters)...")
+    base_res = benchmark(base_model, batch, device, args.warmup, args.iterations)
+    print(f"    → {base_res['latency_ms']:.2f} ms | {base_res['throughput']:.0f} samp/s | {base_res['peak_gpu_mb']:.0f} MB")
+
+    batch_t = make_batch(args.batch_size, args.seq_len, device)
+    print(f"  Benchmarking Tiered ({args.iterations} iters)...")
+    tier_res = benchmark(tier_model, batch_t, device, args.warmup, args.iterations)
+    print(f"    → {tier_res['latency_ms']:.2f} ms | {tier_res['throughput']:.0f} samp/s | {tier_res['peak_gpu_mb']:.0f} MB")
+
+    speedup = base_res["latency_ms"] / tier_res["latency_ms"]
+    print(f"\n  Speedup: {speedup:.2f}x")
+
+    # ── Sweep ──
+    sweep_data = None
+    if args.sweep:
+        print("\n  Running batch-size sweep...")
+        sweep_data = {"baseline": [], "tiered": []}
+        del base_model, tier_model
+        torch.cuda.empty_cache()
+
+        for bs in [1, 4, 8, 16, 32, 64]:
+            print(f"    bs={bs}...", end=" ", flush=True)
+            c = make_config(bs, args.seq_len, args.hidden_size, args.num_heads)
+            b = make_batch(bs, args.seq_len, device)
+
+            bm = build_model("baseline", c, device)
+            br = benchmark(bm, b, device, warmup=3, iterations=10)
+            br["batch_size"] = bs
+            sweep_data["baseline"].append(br)
+            del bm
+
+            tm = build_model("tiered", c, device)
+            tr = benchmark(tm, b, device, warmup=3, iterations=10)
+            tr["batch_size"] = bs
+            sweep_data["tiered"].append(tr)
+            del tm
+            torch.cuda.empty_cache()
+            print(f"base={br['latency_ms']:.2f}ms, tier={tr['latency_ms']:.2f}ms")
+
+    # ── Plots ──
+    print("\n  Generating plots...")
+    create_plots(base_res, tier_res, sweep_data, args.output_dir)
+
+    # ── Save JSON ──
+    results = {"baseline": base_res, "tiered": tier_res, "speedup": speedup}
+    if sweep_data:
+        results["sweep"] = sweep_data
+    json_path = os.path.join(args.output_dir, "comparison_results.json")
+    os.makedirs(args.output_dir, exist_ok=True)
+    with open(json_path, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"  Results saved → {json_path}")
+
+    print("\n" + "=" * 64)
+    print("  Done!")
+    print("=" * 64)
+
+
+if __name__ == "__main__":
+    main()

config/arch/hrm_tiered.yaml

@@ -0,0 +1,27 @@
+# Tiered HRM config — extends hrm_v1 with SRAM/DRAM memory tiering
+name: hrm.hrm_tiered@HRM_Tiered
+loss:
+  name: losses@ACTLossHead
+  loss_type: stablemax_cross_entropy
+
+halt_exploration_prob: 0.1
+halt_max_steps: 16
+
+H_cycles: 2
+L_cycles: 2
+
+H_layers: 4
+L_layers: 4
+
+hidden_size: 512
+num_heads: 8
+expansion: 4
+
+puzzle_emb_ndim: ${.hidden_size}
+
+pos_encodings: rope
+
+# Memory tier config (used by MemoryTierManager)
+memory_tier:
+  sram_capacity_mb: 48         # Simulated SRAM budget (e.g., L2 cache size)
+  enable_tracking: true        # Log all memory operations for benchmarks

config/arch/hrm_v1.yaml

@@ -0,0 +1,21 @@
+name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
+loss:
+  name: losses@ACTLossHead
+  loss_type: stablemax_cross_entropy
+
+halt_exploration_prob: 0.1
+halt_max_steps: 16
+
+H_cycles: 2
+L_cycles: 2
+
+H_layers: 4
+L_layers: 4
+
+hidden_size: 512
+num_heads: 8  # min(2, hidden_size // 64)
+expansion: 4
+
+puzzle_emb_ndim: ${.hidden_size}
+
+pos_encodings: rope

config/cfg_pretrain.yaml

@@ -0,0 +1,31 @@
+# ARC training config
+
+defaults:
+  - arch: hrm_v1
+  - _self_
+
+hydra:
+  output_subdir: null
+
+# Data path
+data_path: data/arc-aug-1000
+
+# Hyperparams - Training
+global_batch_size: 768
+
+epochs: 100000
+eval_interval: 10000
+checkpoint_every_eval: True
+
+lr: 1e-4
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+
+# Standard hyperparameter settings for LM, as used in Llama
+beta1: 0.9
+beta2: 0.95
+weight_decay: 0.1
+puzzle_emb_weight_decay: 0.1
+
+# Hyperparams - Puzzle embeddings training
+puzzle_emb_lr: 1e-2

dataset/build_arc_dataset.py

@@ -0,0 +1,291 @@
+from typing import List, Optional, Tuple, Dict
+from dataclasses import dataclass
+from pathlib import Path
+import os
+import json
+import hashlib
+import numpy as np
+from glob import glob
+
+from argdantic import ArgParser
+from pydantic import BaseModel
+
+from common import PuzzleDatasetMetadata, dihedral_transform
+
+
+cli = ArgParser()
+
+
+class DataProcessConfig(BaseModel):
+    # ARC-1
+    dataset_dirs: List[str] = ["dataset/raw-data/ARC-AGI/data", "dataset/raw-data/ConceptARC/corpus"]
+    output_dir: str = "data/arc-aug-1000"
+    
+    # ARC-2
+    # dataset_dirs: List[str] = ["dataset/raw-data/ARC-AGI-2/data"]
+    # output_dir: str = "data/arc-2-aug-1000"
+
+    seed: int = 42
+    num_aug: int = 1000
+    
+    
+ARCMaxGridSize = 30
+ARCAugmentRetriesFactor = 5
+    
+
+@dataclass
+class ARCPuzzle:
+    id: str
+
+    examples: List[Tuple[np.ndarray, np.ndarray]]
+
+    
+def arc_grid_to_np(grid: List[List[int]]):
+    arr = np.array(grid)
+
+    # Shape check
+    assert arr.ndim == 2
+    assert arr.shape[0] <= ARCMaxGridSize and arr.shape[1] <= ARCMaxGridSize
+    # Element check
+    assert np.all((arr >= 0) & (arr <= 9))
+    return arr.astype(np.uint8)
+
+
+def np_grid_to_seq_translational_augment(inp: np.ndarray, out: np.ndarray, do_translation: bool):
+    # PAD: 0, <eos>: 1, digits: 2 ... 11
+    # Compute random top-left pad
+    if do_translation:
+        pad_r = np.random.randint(0, ARCMaxGridSize - max(inp.shape[0], out.shape[0]) + 1)
+        pad_c = np.random.randint(0, ARCMaxGridSize - max(inp.shape[1], out.shape[1]) + 1)
+    else:
+        pad_r = pad_c = 0
+
+    # Pad grid
+    result = []
+    for grid in [inp, out]:
+        nrow, ncol = grid.shape
+        grid = np.pad(grid + 2, ((pad_r, ARCMaxGridSize - pad_r - nrow), (pad_c, ARCMaxGridSize - pad_c - ncol)), constant_values=0)
+
+        # Add <eos>
+        eos_row, eos_col = pad_r + nrow, pad_c + ncol
+        if eos_row < ARCMaxGridSize:
+            grid[eos_row, pad_c:eos_col] = 1
+        if eos_col < ARCMaxGridSize:
+            grid[pad_r:eos_row, eos_col] = 1
+
+        result.append(grid.flatten())
+
+    return result
+
+
+def puzzle_hash(puzzle: dict):
+    # Hash the puzzle for checking equivalence
+    def _grid_hash(grid: np.ndarray):
+        buffer = [x.to_bytes(1) for x in grid.shape]
+        buffer.append(grid.tobytes())
+        
+        return hashlib.sha256(b"".join(buffer)).hexdigest()
+    
+    hashes = []
+    for example_type, example in puzzle.items():
+        for input, label in example.examples:
+            hashes.append(f"{_grid_hash(input)}|{_grid_hash(label)}")
+            
+    hashes.sort()
+    return hashlib.sha256("|".join(hashes).encode()).hexdigest()
+
+
+def convert_single_arc_puzzle(results: dict, default_name: str, puzzle: dict, aug_count: int, dest_mapping: Dict[str, Tuple[str, str]]):
+    # Remove "name"
+    name = puzzle.pop("name", default_name)
+    
+    # Convert
+    dests = set(dest_mapping.values())
+    converted = {dest: ARCPuzzle(name, []) for dest in dests}
+    for example_type, examples in puzzle.items():
+        dest = dest_mapping[example_type]
+        converted[dest].examples.extend([(arc_grid_to_np(example["input"]), arc_grid_to_np(example["output"])) for example in examples])
+
+    group = [converted]
+    
+    # Augment
+    if aug_count > 0:
+        hashes = {puzzle_hash(converted)}
+
+        for _trial in range(ARCAugmentRetriesFactor * aug_count):
+            # Augment plan
+            trans_id = np.random.randint(0, 8)
+            mapping = np.concatenate([np.arange(0, 1, dtype=np.uint8), np.random.permutation(np.arange(1, 10, dtype=np.uint8))])  # Permute colors, Excluding "0" (black)
+            
+            aug_repr = f"t{trans_id}_{''.join(str(x) for x in mapping)}"
+
+            def _map_grid(grid: np.ndarray):
+                return dihedral_transform(mapping[grid], trans_id)
+            
+            # Check duplicate
+            augmented = {dest: ARCPuzzle(f"{puzzle.id}_{aug_repr}", [(_map_grid(input), _map_grid(label)) for (input, label) in puzzle.examples]) for dest, puzzle in converted.items()}
+            h = puzzle_hash(augmented)
+            if h not in hashes:
+                hashes.add(h)
+                group.append(augmented)
+                
+            if len(group) >= aug_count + 1:
+                break
+            
+        if len(group) < aug_count + 1:
+            print (f"[Puzzle {name}] augmentation not full, only {len(group)}")
+
+    # Append
+    for dest in dests:
+        # Convert the examples
+        dest_split, dest_set = dest
+
+        results.setdefault(dest_split, {})
+        results[dest_split].setdefault(dest_set, [])
+        results[dest_split][dest_set].append([converted[dest] for converted in group])
+
+
+def load_puzzles_arcagi(results: dict, dataset_path: str, config: DataProcessConfig):
+    train_examples_dest = ("train", "all")
+    test_examples_map = {
+        "evaluation": [(1.0, ("test", "all"))],
+        "_default": [(1.0, ("train", "all"))]
+    }
+    
+    total_puzzles = 0
+    for subdir in os.scandir(dataset_path):
+        if subdir.is_dir():
+            # Load all puzzles in this directory
+            puzzles = []
+            for filename in glob(os.path.join(subdir.path, "*.json")):
+                with open(filename, "r") as f:
+                    puzzles.append((Path(filename).stem, json.load(f)))
+                    
+            # Shuffle puzzles
+            np.random.shuffle(puzzles)
+            
+            # Assign by fraction
+            for idx, (default_name, puzzle) in enumerate(puzzles):
+                fraction = idx / len(puzzles)
+                test_examples_dest = None
+                for f, dest in test_examples_map.get(subdir.name, test_examples_map["_default"]):
+                    if fraction < f:
+                        test_examples_dest = dest
+                        break
+                        
+                assert test_examples_dest is not None
+                
+                convert_single_arc_puzzle(results, default_name, puzzle, config.num_aug, {"train": train_examples_dest, "test": test_examples_dest})
+                total_puzzles += 1
+
+    print (f"[{dataset_path}] total puzzles: {total_puzzles}")
+
+
+def convert_dataset(config: DataProcessConfig):
+    np.random.seed(config.seed)
+    
+    # Read dataset
+    data = {}
+    for dataset_dir in config.dataset_dirs:
+        load_puzzles_arcagi(data, dataset_dir, config)
+    
+    # Map global puzzle identifiers
+    num_identifiers = 1  # 0 is blank
+    identifier_map = {}
+    for split_name, split in data.items():
+        for subset_name, subset in split.items():
+            for group in subset:
+                for puzzle in group:
+                    if puzzle.id not in identifier_map:
+                        identifier_map[puzzle.id] = num_identifiers
+                        num_identifiers += 1
+
+    print (f"Total puzzle IDs (including <blank>): {num_identifiers}")
+
+    # Save
+    for split_name, split in data.items():
+        os.makedirs(os.path.join(config.output_dir, split_name), exist_ok=True)
+        
+        # Translational augmentations
+        enable_translational_augment = split_name == "train"
+
+        # Statistics
+        total_examples = 0
+        total_puzzles = 0
+        total_groups = 0
+        
+        for subset_name, subset in split.items():
+            # Construct subset
+            results = {k: [] for k in ["inputs", "labels", "puzzle_identifiers", "puzzle_indices", "group_indices"]}
+            results["puzzle_indices"].append(0)
+            results["group_indices"].append(0)
+            
+            example_id = 0
+            puzzle_id = 0
+            
+            for group in subset:
+                for puzzle in group:
+                    # Push puzzle
+                    no_aug_id = np.random.randint(0, len(puzzle.examples))
+                    for _idx_ex, (inp, out) in enumerate(puzzle.examples):
+                        inp, out = np_grid_to_seq_translational_augment(inp, out, do_translation=enable_translational_augment and _idx_ex != no_aug_id)
+                            
+                        results["inputs"].append(inp)
+                        results["labels"].append(out)
+                        example_id += 1
+                        
+                        total_examples += 1
+
+                    results["puzzle_indices"].append(example_id)
+                    results["puzzle_identifiers"].append(identifier_map[puzzle.id])
+                    
+                    puzzle_id += 1
+                    
+                    total_puzzles += 1
+                    
+                # Push group
+                results["group_indices"].append(puzzle_id)
+                total_groups += 1
+            
+            for k, v in results.items():
+                if k in {"inputs", "labels"}:
+                    v = np.stack(v, 0)
+                else:
+                    v = np.array(v, dtype=np.int32)
+                
+                np.save(os.path.join(config.output_dir, split_name, f"{subset_name}__{k}.npy"), v)
+        
+        # Metadata
+        metadata = PuzzleDatasetMetadata(
+            seq_len=ARCMaxGridSize * ARCMaxGridSize,
+            vocab_size=10 + 2,  # PAD + EOS + "0" ... "9"
+            
+            pad_id=0,
+            ignore_label_id=0,
+            
+            blank_identifier_id=0,
+            num_puzzle_identifiers=num_identifiers,
+            
+            total_groups=total_groups,
+            mean_puzzle_examples=total_examples / total_puzzles,
+            sets=list(split.keys())
+        )
+
+        # Save metadata as JSON.
+        with open(os.path.join(config.output_dir, split_name, "dataset.json"), "w") as f:
+            json.dump(metadata.model_dump(), f)
+            
+    # Save IDs mapping
+    with open(os.path.join(config.output_dir, "identifiers.json"), "w") as f:
+        ids_mapping = {v: k for k, v in identifier_map.items()}
+        
+        json.dump([ids_mapping.get(i, "<blank>") for i in range(num_identifiers)], f)
+
+
+@cli.command(singleton=True)
+def main(config: DataProcessConfig):
+    convert_dataset(config)
+
+
+if __name__ == "__main__":
+    cli()

dataset/build_maze_dataset.py

@@ -0,0 +1,142 @@
+from typing import Optional
+import math
+import os
+import csv
+import json
+import numpy as np
+
+from argdantic import ArgParser
+from pydantic import BaseModel
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+
+from common import PuzzleDatasetMetadata, dihedral_transform
+
+
+CHARSET = "# SGo"
+
+
+cli = ArgParser()
+
+
+class DataProcessConfig(BaseModel):
+    source_repo: str = "sapientinc/maze-30x30-hard-1k"
+    output_dir: str = "data/maze-30x30-hard-1k"
+
+    subsample_size: Optional[int] = None
+    aug: bool = False
+
+
+def convert_subset(set_name: str, config: DataProcessConfig):
+    # Read CSV
+    all_chars = set()
+    grid_size = None
+    inputs = []
+    labels = []
+    
+    with open(hf_hub_download(config.source_repo, f"{set_name}.csv", repo_type="dataset"), newline="") as csvfile:  # type: ignore
+        reader = csv.reader(csvfile)
+        next(reader)  # Skip header
+        for source, q, a, rating in reader:
+            all_chars.update(q)
+            all_chars.update(a)
+
+            if grid_size is None:
+                n = int(len(q) ** 0.5)
+                grid_size = (n, n)
+                
+            inputs.append(np.frombuffer(q.encode(), dtype=np.uint8).reshape(grid_size))
+            labels.append(np.frombuffer(a.encode(), dtype=np.uint8).reshape(grid_size))
+
+    # If subsample_size is specified for the training set,
+    # randomly sample the desired number of examples.
+    if set_name == "train" and config.subsample_size is not None:
+        total_samples = len(inputs)
+        if config.subsample_size < total_samples:
+            indices = np.random.choice(total_samples, size=config.subsample_size, replace=False)
+            inputs = [inputs[i] for i in indices]
+            labels = [labels[i] for i in indices]
+
+    # Generate dataset
+    results = {k: [] for k in ["inputs", "labels", "puzzle_identifiers", "puzzle_indices", "group_indices"]}
+    puzzle_id = 0
+    example_id = 0
+    
+    results["puzzle_indices"].append(0)
+    results["group_indices"].append(0)
+    
+    for inp, out in zip(tqdm(inputs), labels):
+        # Dihedral transformations for augmentation
+        for aug_idx in range(8 if (set_name == "train" and config.aug) else 1):
+            results["inputs"].append(dihedral_transform(inp, aug_idx))
+            results["labels"].append(dihedral_transform(out, aug_idx))
+            example_id += 1
+            puzzle_id += 1
+            
+            results["puzzle_indices"].append(example_id)
+            results["puzzle_identifiers"].append(0)
+            
+        # Push group
+        results["group_indices"].append(puzzle_id)
+            
+    # Char mappings
+    assert len(all_chars - set(CHARSET)) == 0
+    
+    char2id = np.zeros(256, np.uint8)
+    char2id[np.array(list(map(ord, CHARSET)))] = np.arange(len(CHARSET)) + 1
+
+    # To Numpy
+    def _seq_to_numpy(seq):
+        arr = np.vstack([char2id[s.reshape(-1)] for s in seq])
+        
+        return arr
+    
+    results = {
+        "inputs": _seq_to_numpy(results["inputs"]),
+        "labels": _seq_to_numpy(results["labels"]),
+        
+        "group_indices": np.array(results["group_indices"], dtype=np.int32),
+        "puzzle_indices": np.array(results["puzzle_indices"], dtype=np.int32),
+        "puzzle_identifiers": np.array(results["puzzle_identifiers"], dtype=np.int32),
+    }
+
+    # Metadata
+    metadata = PuzzleDatasetMetadata(
+        seq_len=int(math.prod(grid_size)),  # type: ignore
+        vocab_size=len(CHARSET) + 1,  # PAD + Charset
+        
+        pad_id=0,
+        ignore_label_id=0,
+        
+        blank_identifier_id=0,
+        num_puzzle_identifiers=1,
+        
+        total_groups=len(results["group_indices"]) - 1,
+        mean_puzzle_examples=1,
+        sets=["all"]
+    )
+
+    # Save metadata as JSON.
+    save_dir = os.path.join(config.output_dir, set_name)
+    os.makedirs(save_dir, exist_ok=True)
+    
+    with open(os.path.join(save_dir, "dataset.json"), "w") as f:
+        json.dump(metadata.model_dump(), f)
+        
+    # Save data
+    for k, v in results.items():
+        np.save(os.path.join(save_dir, f"all__{k}.npy"), v)
+        
+    # Save IDs mapping (for visualization only)
+    with open(os.path.join(config.output_dir, "identifiers.json"), "w") as f:
+        json.dump(["<blank>"], f)
+
+
+@cli.command(singleton=True)
+def preprocess_data(config: DataProcessConfig):
+    convert_subset("train", config)
+    convert_subset("test", config)
+
+
+if __name__ == "__main__":
+    cli()

dataset/build_sudoku_dataset.py

@@ -0,0 +1,169 @@
+from typing import Optional
+import os
+import csv
+import json
+import numpy as np
+
+from argdantic import ArgParser
+from pydantic import BaseModel
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+
+from common import PuzzleDatasetMetadata
+
+
+cli = ArgParser()
+
+
+class DataProcessConfig(BaseModel):
+    source_repo: str = "sapientinc/sudoku-extreme"
+    output_dir: str = "data/sudoku-extreme-full"
+
+    subsample_size: Optional[int] = None
+    min_difficulty: Optional[int] = None
+    num_aug: int = 0
+
+
+def shuffle_sudoku(board: np.ndarray, solution: np.ndarray):
+    # Create a random digit mapping: a permutation of 1..9, with zero (blank) unchanged
+    digit_map = np.pad(np.random.permutation(np.arange(1, 10)), (1, 0))
+    
+    # Randomly decide whether to transpose.
+    transpose_flag = np.random.rand() < 0.5
+
+    # Generate a valid row permutation:
+    # - Shuffle the 3 bands (each band = 3 rows) and for each band, shuffle its 3 rows.
+    bands = np.random.permutation(3)
+    row_perm = np.concatenate([b * 3 + np.random.permutation(3) for b in bands])
+
+    # Similarly for columns (stacks).
+    stacks = np.random.permutation(3)
+    col_perm = np.concatenate([s * 3 + np.random.permutation(3) for s in stacks])
+
+    # Build an 81->81 mapping. For each new cell at (i, j)
+    # (row index = i // 9, col index = i % 9),
+    # its value comes from old row = row_perm[i//9] and old col = col_perm[i%9].
+    mapping = np.array([row_perm[i // 9] * 9 + col_perm[i % 9] for i in range(81)])
+
+    def apply_transformation(x: np.ndarray) -> np.ndarray:
+        # Apply transpose flag
+        if transpose_flag:
+            x = x.T
+        # Apply the position mapping.
+        new_board = x.flatten()[mapping].reshape(9, 9).copy()
+        # Apply digit mapping
+        return digit_map[new_board]
+
+    return apply_transformation(board), apply_transformation(solution)
+
+
+def convert_subset(set_name: str, config: DataProcessConfig):
+    # Read CSV
+    inputs = []
+    labels = []
+    
+    with open(hf_hub_download(config.source_repo, f"{set_name}.csv", repo_type="dataset"), newline="") as csvfile:
+        reader = csv.reader(csvfile)
+        next(reader)  # Skip header
+        for source, q, a, rating in reader:
+            if (config.min_difficulty is None) or (int(rating) >= config.min_difficulty):
+                assert len(q) == 81 and len(a) == 81
+                
+                inputs.append(np.frombuffer(q.replace('.', '0').encode(), dtype=np.uint8).reshape(9, 9) - ord('0'))
+                labels.append(np.frombuffer(a.encode(), dtype=np.uint8).reshape(9, 9) - ord('0'))
+
+    # If subsample_size is specified for the training set,
+    # randomly sample the desired number of examples.
+    if set_name == "train" and config.subsample_size is not None:
+        total_samples = len(inputs)
+        if config.subsample_size < total_samples:
+            indices = np.random.choice(total_samples, size=config.subsample_size, replace=False)
+            inputs = [inputs[i] for i in indices]
+            labels = [labels[i] for i in indices]
+
+    # Generate dataset
+    num_augments = config.num_aug if set_name == "train" else 0
+
+    results = {k: [] for k in ["inputs", "labels", "puzzle_identifiers", "puzzle_indices", "group_indices"]}
+    puzzle_id = 0
+    example_id = 0
+    
+    results["puzzle_indices"].append(0)
+    results["group_indices"].append(0)
+    
+    for orig_inp, orig_out in zip(tqdm(inputs), labels):
+        for aug_idx in range(1 + num_augments):
+            # First index is not augmented
+            if aug_idx == 0:
+                inp, out = orig_inp, orig_out
+            else:
+                inp, out = shuffle_sudoku(orig_inp, orig_out)
+
+            # Push puzzle (only single example)
+            results["inputs"].append(inp)
+            results["labels"].append(out)
+            example_id += 1
+            puzzle_id += 1
+            
+            results["puzzle_indices"].append(example_id)
+            results["puzzle_identifiers"].append(0)
+            
+        # Push group
+        results["group_indices"].append(puzzle_id)
+        
+    # To Numpy
+    def _seq_to_numpy(seq):
+        arr = np.concatenate(seq).reshape(len(seq), -1)
+        
+        assert np.all((arr >= 0) & (arr <= 9))
+        return arr + 1
+    
+    results = {
+        "inputs": _seq_to_numpy(results["inputs"]),
+        "labels": _seq_to_numpy(results["labels"]),
+        
+        "group_indices": np.array(results["group_indices"], dtype=np.int32),
+        "puzzle_indices": np.array(results["puzzle_indices"], dtype=np.int32),
+        "puzzle_identifiers": np.array(results["puzzle_identifiers"], dtype=np.int32),
+    }
+
+    # Metadata
+    metadata = PuzzleDatasetMetadata(
+        seq_len=81,
+        vocab_size=10 + 1,  # PAD + "0" ... "9"
+        
+        pad_id=0,
+        ignore_label_id=0,
+        
+        blank_identifier_id=0,
+        num_puzzle_identifiers=1,
+        
+        total_groups=len(results["group_indices"]) - 1,
+        mean_puzzle_examples=1,
+        sets=["all"]
+    )
+
+    # Save metadata as JSON.
+    save_dir = os.path.join(config.output_dir, set_name)
+    os.makedirs(save_dir, exist_ok=True)
+    
+    with open(os.path.join(save_dir, "dataset.json"), "w") as f:
+        json.dump(metadata.model_dump(), f)
+        
+    # Save data
+    for k, v in results.items():
+        np.save(os.path.join(save_dir, f"all__{k}.npy"), v)
+        
+    # Save IDs mapping (for visualization only)
+    with open(os.path.join(config.output_dir, "identifiers.json"), "w") as f:
+        json.dump(["<blank>"], f)
+
+
+@cli.command(singleton=True)
+def preprocess_data(config: DataProcessConfig):
+    convert_subset("train", config)
+    convert_subset("test", config)
+
+
+if __name__ == "__main__":
+    cli()

dataset/common.py

@@ -0,0 +1,51 @@
+from typing import List, Optional
+
+import pydantic
+import numpy as np
+
+
+# Global list mapping each dihedral transform id to its inverse.
+# Index corresponds to the original tid, and the value is its inverse.
+DIHEDRAL_INVERSE = [0, 3, 2, 1, 4, 5, 6, 7]
+
+
+class PuzzleDatasetMetadata(pydantic.BaseModel):
+    pad_id: int
+    ignore_label_id: Optional[int]
+    blank_identifier_id: int
+    
+    vocab_size: int
+    seq_len: int
+    num_puzzle_identifiers: int
+    
+    total_groups: int
+    mean_puzzle_examples: float
+
+    sets: List[str]
+
+
+def dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:
+    """8 dihedral symmetries by rotate, flip and mirror"""
+    
+    if tid == 0:
+        return arr  # identity
+    elif tid == 1:
+        return np.rot90(arr, k=1)
+    elif tid == 2:
+        return np.rot90(arr, k=2)
+    elif tid == 3:
+        return np.rot90(arr, k=3)
+    elif tid == 4:
+        return np.fliplr(arr)       # horizontal flip
+    elif tid == 5:
+        return np.flipud(arr)       # vertical flip
+    elif tid == 6:
+        return arr.T                # transpose (reflection along main diagonal)
+    elif tid == 7:
+        return np.fliplr(np.rot90(arr, k=1))  # anti-diagonal reflection
+    else:
+        return arr
+    
+    
+def inverse_dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray:
+    return dihedral_transform(arr, DIHEDRAL_INVERSE[tid])

docs/END_TO_END_EXPLANATION.md

@@ -0,0 +1,1193 @@
+# End-to-End Explanation: PyTorch & Triton Implementation
+## For someone with basic Python knowledge
+
+---
+
+## PART 1: What Are the Available Benchmarks?
+
+Before diving into code, let's understand what benchmarks are available to measure performance.
+
+### Overview of Benchmarking Suite
+
+The benchmarking system measures how fast the model runs and how efficiently it uses memory. It compares two setups:
+- **Tiered Model**: Uses the SRAM/DRAM memory hierarchy (what we want to study)
+- **Baseline Model**: Standard model without memory tiering
+
+### Files Involved: `benchmark.py` and `run_benchmark.py`
+
+#### A. `run_benchmark.py` - The Command-Line Interface
+
+This file lets you **run benchmarks from the terminal**. Think of it as a control panel.
+
+**Key command-line options:**
+
+```bash
+# See available benchmarks and options
+python run_benchmark.py --help
+
+# Compare tiered vs baseline models
+python run_benchmark.py --mode compare --batch-sizes 1,8,32 --seq-lens 64,128
+
+# Benchmark only the tiered (memory-aware) model
+python run_benchmark.py --mode tiered --warmup 5 --iterations 50 --output results.json
+
+# Quick test (for learning)
+python run_benchmark.py --mode tiered --warmup 1 --iterations 3 --batch-sizes 2 --seq-lens 16
+```
+
+**What each option means:**
+
+| Option | Meaning |
+|--------|---------|
+| `--mode` | What to benchmark: `tiered` (new model), `baseline` (standard), or `compare` (both) |
+| `--batch-sizes` | How many samples to process at once (comma-separated: `1,8,32` means test with 1, 8, and 32 samples) |
+| `--seq-lens` | Length of input sequences: `64,128` means test with sequences of 64 and 128 tokens |
+| `--hidden-size` | Internal dimension of the model (default: 512) |
+| `--H-cycles` | Number of times H-level (slow) processes data (default: 2) |
+| `--L-cycles` | Number of times L-level (fast) processes data (default: 2) |
+| `--H-layers` | Number of H-level transformer layers (default: 4) |
+| `--L-layers` | Number of L-level transformer layers (default: 4) |
+| `--warmup` | How many runs to discard before measuring (to let GPU settle) |
+| `--iterations` | How many actual measurements to take |
+| `--output` | Save results to JSON file |
+
+---
+
+### B. What Does `benchmark.py` Actually Measure?
+
+Inside `benchmark.py`, the `BenchmarkResult` class records these metrics:
+
+#### **Latency Metrics** (How fast things run, in microseconds `μs`)
+
+- **L-level latency**: Time for the "fast" tier to process data
+- **H-level latency**: Time for the "slow" tier to process data
+- **Total inference latency**: End-to-end prediction time (in milliseconds `ms`)
+- **h_over_l_latency_ratio**: How many times slower H-level is than L-level
+
+**Why this matters:** If H-level takes 10× longer than L-level, we know the memory hierarchy is working — slow operations vs. fast ones.
+
+#### **Memory Metrics** (How much GPU memory used)
+
+- **sram_peak_mb**: Peak memory in the "fast" tier
+- **dram_peak_mb**: Peak memory in the "slow" tier
+- **total_gpu_memory_mb**: Total GPU memory used
+
+#### **Transfer Metrics** (Cost of moving data between tiers, in microseconds)
+
+- **h_l_transfer_mean_us**: Average time to copy data from H→L
+- **l_h_transfer_mean_us**: Average time to copy data from L→H
+
+**Why this matters:** If transfers are expensive, the tiering strategy might not be worth it.
+
+#### **Hit Rate & Efficiency**
+
+- **sram_hit_rate**: How often data we tried to put in "fast" memory actually fit (0.0 = never, 1.0 = always)
+- **memory_efficiency**: Useful compute time / total time (higher is better)
+
+**Why this matters:** If hit_rate is low, data keeps spilling to slow memory and we're not getting the optimization benefit.
+
+#### **Triton Kernel Probes** (Direct measurement of memory latency)
+
+- **triton_sram_probe_latency_us**: Measured latency of "fast" memory via Triton kernels
+- **triton_dram_probe_latency_us**: Measured latency of "slow" memory via Triton kernels
+
+**Why this matters:** These are the "ground truth" measurements showing the actual speed difference between SRAM and DRAM.
+
+---
+
+### C. How Benchmarks Are Run
+
+**Flow of a benchmark:**
+
+1. **Create dummy data** (synthetic input tensors)
+2. **Warmup phase** (run N times to let GPU caches warm up)
+   - These runs are NOT counted in final results
+3. **Timing phase** (run N times with GPU timers active)
+   - Start GPU timer
+   - Run model inference
+   - Stop GPU timer
+   - Record elapsed time
+4. **Collect statistics**: min, max, mean, std dev from all the timed runs
+5. **Measure memory** (peak usage during runs)
+6. **Calculate derived metrics** (ratios, efficiency scores)
+7. **Output results** (to console and/or JSON file)
+
+---
+
+## PART 2: PyTorch Implementation — Line-by-Line
+
+### File 1: `memory_tier.py` — The Memory Manager
+
+This is the **PyTorch foundation** for the two-tier memory system. It mimics how GPU memory is organized (SRAM = fast cache vs. DRAM = main memory).
+
+#### **Header & Imports (Lines 1-22)**
+
+```python
+"""
+Memory Tier Manager for HRM SRAM/DRAM implementation.
+
+Manages the placement of H-level and L-level hidden states across
+GPU memory tiers and tracks all memory operations for benchmarking.
+
+- SRAM tier: Uses CUDA pinned memory + explicit prefetching.
+             L-level states are kept GPU-resident with minimal transfers.
+- DRAM tier: Standard GPU global memory with transfer tracking.
+             H-level states go through normal allocation paths.
+"""
+
+import time
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+from contextlib import contextmanager
+
+import torch
+```
+
+**Explanation:**
+- The docstring explains the purpose: manage two memory tiers
+- `dataclass` and `field` are used to create structured data containers
+- `contextmanager` creates a context (like `with` statements in Python)
+- `torch` is imported to use PyTorch tensor operations
+
+#### **Section 1: MemoryEvent — Tracking Individual Operations (Lines 26-31)**
+
+```python
+@dataclass
+class MemoryEvent:
+    """A single tracked memory operation."""
+    tier: str           # 'sram' or 'dram'
+    operation: str      # 'alloc', 'load', 'store', 'transfer'
+    bytes: int
+    duration_us: float  # microseconds
+    timestamp: float
+```
+
+**Explanation:**
+- A `@dataclass` is like a **lightweight container** for data
+- Each `MemoryEvent` records ONE memory operation (e.g., "allocated 1024 bytes in SRAM in 50 microseconds")
+- Fields:
+  - `tier`: Which tier (fast or slow)?
+  - `operation`: What happened? (allocate new memory, load, store, transfer between tiers)
+  - `bytes`: How much data?
+  - `duration_us`: How long did it take? (in microseconds, where 1000 μs = 1 ms)
+  - `timestamp`: When did it happen?
+
+**Real example:**
+```python
+event = MemoryEvent(
+    tier='sram',
+    operation='alloc',
+    bytes=65536,
+    duration_us=123.45,
+    timestamp=1704067200.123
+)
+# Created a record: "SRAM allocation of 65536 bytes took 123.45 microseconds"
+```
+
+#### **Section 2: TierStats — Accumulated Statistics (Lines 35-56)**
+
+```python
+@dataclass
+class TierStats:
+    """Accumulated statistics for one memory tier."""
+    total_alloc_bytes: int = 0
+    peak_alloc_bytes: int = 0
+    current_alloc_bytes: int = 0
+    num_loads: int = 0
+    num_stores: int = 0
+    num_transfers: int = 0
+    total_load_us: float = 0.0
+    total_store_us: float = 0.0
+    total_transfer_us: float = 0.0
+    hit_count: int = 0
+    miss_count: int = 0
+
+    @property
+    def hit_rate(self) -> float:
+        total = self.hit_count + self.miss_count
+        return self.hit_count / total if total > 0 else 0.0
+
+    @property
+    def avg_load_us(self) -> float:
+        return self.total_load_us / self.num_loads if self.num_loads > 0 else 0.0
+
+    @property
+    def avg_store_us(self) -> float:
+        return self.total_store_us / self.num_stores if self.num_stores > 0 else 0.0
+```
+
+**Explanation:**
+- This accumulates **total statistics** for a tier (all operations combined)
+- `total_alloc_bytes`: Total memory ever allocated
+- `peak_alloc_bytes`: Maximum memory used at any point
+- `current_alloc_bytes`: Memory currently in use
+- `num_loads`, `num_stores`: Counters for how many times data was read/written
+- `hit_count` / `miss_count`: Success/failure for fitting data in SRAM
+- **Properties** (`@property`): Computed on-the-fly from raw counts
+  - `hit_rate = hit_count / (hit_count + miss_count)` — percentage of successful SRAM fits
+  - `avg_load_us = total_load_us / num_loads` — average speed per load operation
+  - `avg_store_us = total_store_us / num_stores` — average speed per store operation
+
+**Why properties are useful:**
+Instead of storing `hit_rate` separately and having to update it constantly, we compute it whenever asked: `stats.hit_rate` automatically calculates the current rate.
+
+---
+
+#### **Section 3: MemoryTierManager — The Main Manager Class (Lines 60-90)**
+
+```python
+class MemoryTierManager:
+    """Coordinates SRAM/DRAM memory placement and tracking for HRM.
+
+    In the Triton context:
+    - SRAM tier: Tensors allocated with `pin_memory` and kept on the
+      same CUDA stream as L-level computation. Triton kernels keep
+      these values in registers/shared memory via data reuse.
+    - DRAM tier: Standard `torch.cuda` tensors. Triton kernels load
+      these from global memory each time.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        enable_tracking: bool = True,
+        sram_capacity_mb: float = 48.0,  # Typical L2 cache size
+    ):
+        self.device = device
+        self.enable_tracking = enable_tracking
+        self.sram_capacity_bytes = int(sram_capacity_mb * 1024 * 1024)
+
+        # State registries
+        self._sram_tensors: Dict[str, torch.Tensor] = {}
+        self._dram_tensors: Dict[str, torch.Tensor] = {}
+
+        # Event log
+        self._events: List[MemoryEvent] = []
+        self._sram_stats = TierStats()
+        self._dram_stats = TierStats()
+
+        # CUDA events for GPU timing
+        self._use_cuda = device.type == 'cuda'
+        if self._use_cuda:
+            self._sram_stream = torch.cuda.Stream(device=device)
+            self._dram_stream = torch.cuda.Stream(device=device)
+        else:
+            self._sram_stream = None
+            self._dram_stream = None
+```
+
+**Explanation:**
+
+- **`__init__` method**: Initializes the manager when created
+  - `device`: Where to allocate (CPU or GPU?)
+  - `enable_tracking`: Should we record all operations?
+  - `sram_capacity_mb`: How much "fast" memory is available? (Typical GPU L2 cache = 48 MB)
+  
+- **State registries:**
+  - `_sram_tensors`: Dictionary storing all tensors in SRAM (key = name, value = tensor)
+  - `_dram_tensors`: Dictionary storing all tensors in DRAM
+  - Example: `_sram_tensors['layer1_hidden'] = torch.tensor(...)`
+
+- **Event log:**
+  - `_events`: A list of `MemoryEvent` objects (every operation is recorded)
+  - `_sram_stats`, `_dram_stats`: Running statistics for each tier
+
+- **CUDA streams:**
+  - A stream is like a **"lane"** for GPU operations (operations in same lane execute sequentially, different lanes can run in parallel)
+  - Two separate streams (`_sram_stream`, `_dram_stream`) let us potentially run fast and slow operations concurrently
+
+---
+
+#### **Section 4A: SRAM Allocation (Lines 95-123)**
+
+```python
+def alloc_sram(self, name: str, shape: Tuple, dtype: torch.dtype) -> torch.Tensor:
+    """Allocate a tensor in the SRAM tier (GPU-resident, pinned)."""
+    # Calculate size in bytes
+    nbytes = torch.tensor([], dtype=dtype).element_size()
+    for s in shape:
+        nbytes *= s
+
+    # Check capacity: does it fit?
+    if self._sram_stats.current_alloc_bytes + nbytes > self.sram_capacity_bytes:
+        # Doesn't fit → spill to DRAM and record a "miss"
+        self._sram_stats.miss_count += 1
+        return self.alloc_dram(name, shape, dtype)
+
+    # It fits! Record a "hit"
+    self._sram_stats.hit_count += 1
+
+    # Time the allocation
+    t0 = self._timer_start()
+    tensor = torch.zeros(shape, dtype=dtype, device=self.device)
+
+    # Hint: keep GPU-resident (don't swap to CPU)
+    if self._use_cuda:
+        with torch.cuda.stream(self._sram_stream):
+            tensor = tensor.contiguous()
+
+    self._sram_tensors[name] = tensor
+    dur = self._timer_end(t0)
+
+    # Update statistics
+    self._sram_stats.total_alloc_bytes += nbytes
+    self._sram_stats.current_alloc_bytes += nbytes
+    self._sram_stats.peak_alloc_bytes = max(
+        self._sram_stats.peak_alloc_bytes,
+        self._sram_stats.current_alloc_bytes,
+    )
+
+    # Record the operation
+    self._record_event('sram', 'alloc', nbytes, dur)
+    return tensor
+```
+
+**Explanation (line by line):**
+
+1. **Calculate size:** Convert shape into bytes
+   - Example: shape `(128, 512)` with `float32` (4 bytes) = 128 × 512 × 4 = 262,144 bytes
+
+2. **Capacity check:** Does this allocation fit in SRAM?
+   - If `current_alloc_bytes + nbytes > sram_capacity_bytes`, it doesn't fit
+   - **Fall back to DRAM** and **count as a miss** (failed to use SRAM)
+   - Otherwise, **count as a hit** (successfully allocated in SRAM)
+
+3. **Timing:** Record when allocation starts (`t0`)
+
+4. **Create tensor:** `torch.zeros(shape, ...)` creates a tensor of zeros
+   - `shape`: dimensions
+   - `dtype`: data type (e.g., float32)
+   - `device`: 'cuda' or 'cpu'
+
+5. **GPU optimization:** `tensor.contiguous()` makes the data **contiguous in memory** (important for GPU performance)
+   - Done on the SRAM stream to suggest GPU placement
+
+6. **Store in registry:** Save for later: `_sram_tensors[name] = tensor`
+
+7. **Update stats:**
+   - Add `nbytes` to total and current
+   - Update peak if we're now using more than before
+
+8. **Record event:** Add this operation to the event log
+
+---
+
+#### **Section 4B: DRAM Allocation (Lines 126-142)**
+
+```python
+def alloc_dram(self, name: str, shape: Tuple, dtype: torch.dtype) -> torch.Tensor:
+    """Allocate a tensor in the DRAM tier (standard GPU memory)."""
+    nbytes = torch.tensor([], dtype=dtype).element_size()
+    for s in shape:
+        nbytes *= s
+
+    t0 = self._timer_start()
+    tensor = torch.zeros(shape, dtype=dtype, device=self.device)
+    self._dram_tensors[name] = tensor
+    dur = self._timer_end(t0)
+
+    self._dram_stats.total_alloc_bytes += nbytes
+    self._dram_stats.current_alloc_bytes += nbytes
+    self._dram_stats.peak_alloc_bytes = max(
+        self._dram_stats.peak_alloc_bytes,
+        self._dram_stats.current_alloc_bytes,
+    )
+
+    self._record_event('dram', 'alloc', nbytes, dur)
+    return tensor
+```
+
+**Explanation:**
+- **Almost identical to `alloc_sram`**, but:
+  - **No capacity check** (unrestricted DRAM)
+  - **No stream optimization** (standard allocation)
+  - **Stores in `_dram_tensors`** instead of `_sram_tensors`
+
+---
+
+#### **Section 5: Cross-Tier Transfers (Lines 147-177)**
+
+```python
+def transfer_sram_to_dram(self, name: str) -> torch.Tensor:
+    """Copy a tensor from SRAM tier to DRAM tier."""
+    src = self._sram_tensors[name]
+    t0 = self._timer_start()
+    dst = src.clone()
+    if self._use_cuda:
+        torch.cuda.synchronize(self.device)
+    dur = self._timer_end(t0)
+
+    self._dram_tensors[name + '_from_sram'] = dst
+    self._sram_stats.num_transfers += 1
+    self._sram_stats.total_transfer_us += dur
+    self._record_event('sram', 'transfer', src.nelement() * src.element_size(), dur)
+    return dst
+
+def transfer_dram_to_sram(self, name: str) -> torch.Tensor:
+    """Copy a tensor from DRAM tier to SRAM tier."""
+    src = self._dram_tensors[name]
+    t0 = self._timer_start()
+    dst = src.clone()
+    if self._use_cuda:
+        torch.cuda.synchronize(self.device)
+    dur = self._timer_end(t0)
+
+    self._sram_tensors[name + '_from_dram'] = dst
+    self._dram_stats.num_transfers += 1
+    self._dram_stats.total_transfer_us += dur
+    self._record_event('dram', 'transfer', src.nelement() * src.element_size(), dur)
+    return dst
+```
+
+**Explanation:**
+
+- **H→L transfer (SRAM to DRAM):**
+  1. Get source tensor from SRAM registry
+  2. `clone()` = make a complete copy
+  3. `torch.cuda.synchronize()` = wait for GPU to finish (so we time actual operation)
+  4. Calculate duration
+  5. Store in DRAM registry with new name
+  6. Update transfer counter and total transfer time
+  7. Record the event
+
+- **L→H transfer (DRAM to SRAM):** Same logic in reverse
+
+**Why clone?** We don't want to move the original; we want a copy at the destination.
+
+**Why synchronize?** GPU operations often run asynchronously (GPU schedules them but CPU moves on). Synchronize makes CPU wait, so we measure actual GPU time.
+
+---
+
+#### **Section 6: Context Managers for Streams (Lines 182-198)**
+
+```python
+@contextmanager
+def sram_context(self):
+    """Context manager that runs operations on the SRAM stream."""
+    if self._use_cuda and self._sram_stream is not None:
+        with torch.cuda.stream(self._sram_stream):
+            yield self._sram_stream
+    else:
+        yield None
+
+@contextmanager
+def dram_context(self):
+    """Context manager that runs operations on the DRAM stream."""
+    if self._use_cuda and self._dram_stream is not None:
+        with torch.cuda.stream(self._dram_stream):
+            yield self._dram_stream
+    else:
+        yield None
+```
+
+**Explanation:**
+
+A `@contextmanager` lets you use `with` statements:
+
+```python
+# Usage:
+with memory_manager.sram_context() as stream:
+    # Operations here run on the SRAM stream
+    x = torch.matmul(a, b)  # GPU runs this on _sram_stream
+    y = x + 1
+
+# When we exit the block, stream is restored
+```
+
+**Why useful?** Different operations can run on different streams in parallel, potentially overlapping computation and memory transfer.
+
+---
+
+#### **Section 7: Statistics & Reporting (Lines 203-243)**
+
+```python
+def get_stats(self) -> Dict:
+    """Return all memory tier statistics."""
+    return {
+        'sram': {
+            'peak_mb': self._sram_stats.peak_alloc_bytes / (1024 * 1024),
+            'current_mb': self._sram_stats.current_alloc_bytes / (1024 * 1024),
+            'hit_rate': self._sram_stats.hit_rate,
+            'num_loads': self._sram_stats.num_loads,
+            'num_stores': self._sram_stats.num_stores,
+            'num_transfers': self._sram_stats.num_transfers,
+            'avg_load_us': self._sram_stats.avg_load_us,
+            'avg_store_us': self._sram_stats.avg_store_us,
+            'total_transfer_us': self._sram_stats.total_transfer_us,
+        },
+        'dram': {
+            # Similar for DRAM...
+        },
+        'num_events': len(self._events),
+    }
+
+def get_events(self) -> List[MemoryEvent]:
+    """Return raw event log."""
+    return list(self._events)
+
+def reset_stats(self):
+    """Clear all statistics and event log."""
+    self._events.clear()
+    self._sram_stats = TierStats()
+    self._dram_stats = TierStats()
+
+def free_all(self):
+    """Release all managed tensors."""
+    self._sram_tensors.clear()
+    self._dram_tensors.clear()
+    self._sram_stats.current_alloc_bytes = 0
+    self._dram_stats.current_alloc_bytes = 0
+```
+
+**Explanation:**
+
+- **`get_stats()`**: Returns a dictionary with all collected metrics (ready to print or save)
+  - Converts bytes to MB (divide by 1024²)
+  - Includes hit rate, counts, timing
+
+- **`get_events()`**: Returns the raw event log (for detailed analysis)
+
+- **`reset_stats()`**: Clear counters before a new benchmark run (don't contaminate results with previous runs)
+
+- **`free_all()`**: Release memory and reset current allocations (cleanup)
+
+---
+
+#### **Section 8: Internal Timing Helpers (Lines 248-257)**
+
+```python
+def _timer_start(self) -> float:
+    if self._use_cuda:
+        torch.cuda.synchronize(self.device)
+    return time.perf_counter()
+
+def _timer_end(self, t0: float) -> float:
+    if self._use_cuda:
+        torch.cuda.synchronize(self.device)
+    return (time.perf_counter() - t0) * 1e6  # → microseconds
+
+def _record_event(self, tier: str, op: str, nbytes: int, dur_us: float):
+    if self.enable_tracking:
+        self._events.append(MemoryEvent(
+            tier=tier, operation=op, bytes=nbytes,
+            duration_us=dur_us, timestamp=time.time(),
+        ))
+```
+
+**Explanation:**
+
+- **`_timer_start()`:**
+  - Synchronize GPU first (wait for all pending operations)
+  - Record CPU time
+  - Returns the start time
+
+- **`_timer_end(t0)`:**
+  - Synchronize GPU (wait for all operations since start)
+  - Calculate elapsed time
+  - Convert to microseconds (× 1,000,000)
+
+- **`_record_event()`:**
+  - If tracking is enabled, create a `MemoryEvent` and add to log
+  - Otherwise, skip (for performance when we don't need detailed logs)
+
+---
+
+### File 2: `softmax.py` — Simple PyTorch Reference
+
+```python
+import torch
+import torch.nn.functional as F
+sample = torch.tensor([[1,2,3,4,5], [5,4,3,2,1]], dtype=torch.float32, device='cuda')
+ref = F.softmax(sample, dim=1)
+print(f"Softmax result: {ref=}")
+```
+
+**Explanation:**
+
+This is a **minimal reference implementation** showing:
+1. Create a sample tensor: 2 rows × 5 columns of numbers
+2. Apply softmax along dimension 1 (across columns)
+3. Print the result
+
+**What softmax does:** Converts scores to probabilities (sum to 1). Example:
+- Input `[1, 2, 3, 4, 5]`
+- Output ~`[0.67, 0.18, 0.05, 0.01, 0.006]` (largest inputs become largest probabilities)
+
+This is a building block used in attention mechanisms in transformers.
+
+---
+
+## PART 3: Triton Implementation — Line-by-Line
+
+### File: `triton_kernels.py` — GPU Kernels for Memory-Aware Compute
+
+Triton is a language for writing GPU kernels (low-level GPU programs). Unlike PyTorch which relies on pre-built library functions, Triton lets us control **exactly** how data flows through GPU memory.
+
+#### **Header & Imports (Lines 1-23)**
+
+```python
+"""
+Triton kernels for HRM SRAM/DRAM memory-tiered operations.
+
+Key idea: In Triton, SRAM = registers + shared memory (managed by compiler
+within a block). DRAM = global memory (GPU HBM). By structuring kernels to keep
+L-level state tile-resident (loaded once, reused many times within a block),
+we ensure L-level stays in SRAM. H-level state is loaded from global memory
+(DRAM) each cycle, paying the full memory bandwidth cost.
+
+This gives us real, measurable latency differences that map to the HRM's
+hierarchical update frequencies.
+"""
+
+import torch
+import triton
+import triton.language as tl
+import math
+```
+
+**Explanation:**
+
+- Key insight: **Data locality matters enormously** on GPUs
+  - SRAM (L1/L2 cache, registers, shared memory): ~1-4 cycles latency, 100s of GB/s bandwidth
+  - DRAM (global memory, HBM): ~200-400 cycles latency, 100s of GB/s like SRAM but way higher latency!
+
+**The trick:** Keep "important" data (L-level) in SRAM by reusing it within a kernel block. Load "less critical" data (H-level) fresh each time from DRAM.
+
+---
+
+#### **Kernel 1: SRAM-Resident RMS-Norm + Residual (Lines 28-64)**
+
+```python
+@triton.jit
+def _rms_norm_residual_fused_kernel(
+    X_ptr,           # Input tensor (residual branch)
+    Residual_ptr,    # Residual connection input
+    Out_ptr,         # Output tensor
+    N: tl.constexpr, # Hidden dimension (constexpr → compiler tiles in SRAM)
+    eps: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Fused RMS-norm + residual add.
+
+    By making N and BLOCK_N constexpr, the compiler keeps the entire hidden
+    vector in registers/shared-memory (SRAM) across the norm computation.
+    This is the kernel used for L-level (fast path).
+    """
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < N
+
+    # ---- Load both inputs into SRAM (registers) in one shot ----
+    x = tl.load(X_ptr + row * N + cols, mask=mask, other=0.0).to(tl.float32)
+    r = tl.load(Residual_ptr + row * N + cols, mask=mask, other=0.0).to(tl.float32)
+
+    # Residual add — stays in registers
+    h = x + r
+
+    # RMS norm — entirely in registers, no global memory round-trip
+    variance = tl.sum(h * h, axis=0) / N
+    h_norm = h * tl.math.rsqrt(variance + eps)
+
+    # ---- Store back to global memory ----
+    tl.store(Out_ptr + row * N + cols, h_norm.to(tl.bfloat16), mask=mask)
+```
+
+**Explanation (line by line):**
+
+1. **Function signature:**
+   - `@triton.jit`: This is GPU code (JIT-compiled, not Python)
+   - `X_ptr`, `Residual_ptr`, `Out_ptr`: **Pointers** to tensors in GPU memory
+   - `N: tl.constexpr`: "constexpr" = compiler treats as a constant (compile-time value, not runtime)
+   - `BLOCK_N: tl.constexpr`: Block size (constexpr tells compiler to optimize for fixed size)
+
+2. **Get thread ID:**
+   - `row = tl.program_id(0)`: Which "row" is this GPU thread processing?
+   - On GPU, thousands of threads run in parallel; each processes one row
+
+3. **Create index array:**
+   - `cols = tl.arange(0, BLOCK_N)`: Array `[0, 1, 2, ..., BLOCK_N-1]`
+   - `mask = cols < N`: Boolean mask to handle if N < BLOCK_N (some threads might not have data)
+
+4. **Load data into registers:**
+   - `tl.load(X_ptr + row * N + cols, ...)`: Load elements from memory
+     - `row * N`: Start at row-th row
+     - `+ cols`: Load all columns for this row
+   - `mask=mask`: Only load valid columns
+   - `other=0.0`: If invalid column, use 0.0
+   - `.to(tl.float32)`: Convert to 32-bit floats for math
+   - **KEY**: All data now in registers/shared memory (SRAM)!
+
+5. **Compute residual:**
+   - `h = x + r`: Add the two inputs (element-wise, all in registers)
+
+6. **Compute RMS norm:**
+   - RMS = sqrt(mean(x²))
+   - `variance = tl.sum(h * h, axis=0) / N`: Compute mean of squares
+   - `h_norm = h * tl.math.rsqrt(variance + eps)`: Normalize by reciprocal square root
+   - `eps`: Small number to avoid division by zero
+   - **KEY**: All arithmetic in registers!
+
+7. **Store result:**
+   - `tl.store(Out_ptr + ..., h_norm.to(tl.bfloat16), mask=mask)`
+   - Write normalized result back to global memory
+   - `.to(tl.bfloat16)`: Convert to lower precision for storage (saves memory bandwidth)
+
+**Why this is fast:**
+- Load data once → keep in SRAM → do many operations → store once
+- If we did this in PyTorch with separate ops (add, then norm), we'd load/store twice
+
+---
+
+#### **Kernel 2: DRAM-Sourced RMS-Norm + Residual (Lines 68-95)**
+
+```python
+@triton.jit
+def _rms_norm_residual_dram_kernel(
+    X_ptr,
+    Residual_ptr,
+    Out_ptr,
+    N: tl.constexpr,
+    eps: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """RMS-norm + residual for H-level.
+
+    Structurally identical but designed to be called with larger strides
+    and without re-use inside a meta-kernel. Each call does a full
+    DRAM round-trip, modeling the slower H-level memory access pattern.
+    """
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < N
+
+    # Global memory load (DRAM)
+    x = tl.load(X_ptr + row * N + cols, mask=mask, other=0.0).to(tl.float32)
+    r = tl.load(Residual_ptr + row * N + cols, mask=mask, other=0.0).to(tl.float32)
+
+    h = x + r
+    variance = tl.sum(h * h, axis=0) / N
+    h_norm = h * tl.math.rsqrt(variance + eps)
+
+    tl.store(Out_ptr + row * N + cols, h_norm.to(tl.bfloat16), mask=mask)
+```
+
+**Explanation:**
+
+**Identical kernel code**, but:
+- **Intent is different**: This is called less frequently (H-level runs slower)
+- **Memory behavior differs** via **calling context**:
+  - SRAM kernel: Data reused many times within tight loops → stays in cache
+  - DRAM kernel: Data used once then discarded → evicted from cache
+
+**In code**, they're identical because the difference is **how often and how much** they're called, not the kernel itself.
+
+---
+
+#### **Kernel 3: SwiGLU Activation (Lines 100-126)**
+
+```python
+@triton.jit
+def _swiglu_fused_sram_kernel(
+    GateUp_ptr,      # [rows, 2 * inter] — gate and up projections concatenated
+    Out_ptr,         # [rows, inter]
+    inter: tl.constexpr,
+    BLOCK_INTER: tl.constexpr,
+):
+    """Fused SiLU(gate) * up in a single kernel pass.
+
+    For L-level: The gate and up vectors are loaded once into registers
+    and the activation is computed without spilling to DRAM.
+    """
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_INTER)
+    mask = cols < inter
+
+    # Load gate and up from contiguous memory — both go into SRAM
+    gate = tl.load(GateUp_ptr + row * 2 * inter + cols, mask=mask, other=0.0).to(tl.float32)
+    up   = tl.load(GateUp_ptr + row * 2 * inter + inter + cols, mask=mask, other=0.0).to(tl.float32)
+
+    # SiLU(gate) * up — entirely in registers
+    silu_gate = gate * tl.sigmoid(gate)
+    result = silu_gate * up
+
+    tl.store(Out_ptr + row * inter + cols, result.to(tl.bfloat16), mask=mask)
+```
+
+**Explanation:**
+
+**SwiGLU** = A gating mechanism in transformers. Formula: `output = sigmoid(gate) * up`
+
+- **Input layout:** `GateUp_ptr` contains concatenated `[gate | up]` (e.g., first half is gate, second half is up)
+- **Load both halves:**
+  - Gate: `row * 2 * inter + cols` (first half)
+  - Up: `row * 2 * inter + inter + cols` (second half)
+- **Compute SiLU:**
+  - `tl.sigmoid(gate)`: Apply sigmoid (s-shaped function) to gate
+  - `silu_gate = gate * sigmoid(gate)`: SiLU = gate * sigmoid(gate)
+  - `result = silu_gate * up`: Gated output
+- **Store:** Result back to global memory
+
+**Why fused?** If done separately:
+1. Load gate
+2. Compute sigmoid
+3. Write intermediate
+4. Load intermediate
+5. Compute SiLU
+6. Write intermediate
+7. Load up
+8. Compute product
+9. Write result
+
+Fused does it in one kernel → one load, one store.
+
+---
+
+#### **Kernel 4: State Transfer (Lines 131-148)**
+
+```python
+@triton.jit
+def _state_transfer_kernel(
+    Src_ptr,
+    Dst_ptr,
+    numel: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Explicit memory copy kernel for cross-tier state transfer.
+
+    Used when H-level needs to read L-level output (or vice versa).
+    Triton compiles this into optimized async memcpy instructions.
+    """
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < numel
+
+    data = tl.load(Src_ptr + offsets, mask=mask, other=0.0)
+    tl.store(Dst_ptr + offsets, data, mask=mask)
+```
+
+**Explanation:**
+
+Simple memcpy kernel (copy data from source to destination):
+- `pid = tl.program_id(0)`: Program/thread block ID
+- `offsets = pid * BLOCK + tl.arange(0, BLOCK)`: Calculate which elements this block handles
+  - If BLOCK=1024, pid=0 handles offsets 0-1023, pid=1 handles 1024-2047, etc.
+- `mask = offsets < numel`: Don't copy past the end
+- Load and store in parallel across many threads
+
+**Why separate kernel?** GPU memcopy can be optimized by the compiler (uses memory controllers, not just compute cores).
+
+---
+
+#### **Kernel 5: Memory Latency Probe (Lines 153-176)**
+
+```python
+@triton.jit
+def _memory_latency_probe_kernel(
+    Data_ptr,
+    Out_ptr,
+    N: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    NUM_ITERS: tl.constexpr,
+):
+    """Probe kernel to measure effective memory latency.
+
+    Performs NUM_ITERS dependent loads to measure true SRAM vs DRAM latency.
+    The data dependency chain prevents compiler reordering.
+    """
+    pid = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < N
+
+    # Initial load from global memory
+    acc = tl.load(Data_ptr + pid * N + cols, mask=mask, other=0.0)
+
+    # Dependent iteration chain — forces sequential memory access
+    for _ in range(NUM_ITERS):
+        # This stays in SRAM (registers) because acc is reused
+        acc = acc * 1.00001 + 0.00001
+
+    tl.store(Out_ptr + pid * N + cols, acc, mask=mask)
+```
+
+**Explanation:**
+
+This kernel **measures latency** by creating a **dependency chain** that can't be optimized away:
+
+1. Load initial data
+2. For N iterations:
+   - Multiply by 1.00001 + add 0.00001 (cheap operations)
+   - Result depends on previous iteration (creates dependency)
+3. Store result
+
+**Why not just load/store?** The compiler could optimize away separate load-stores, but a **dependency chain** forces real latency measurement.
+
+**For SRAM data:**
+- Data stays in registers
+- All iterations hit registers (super fast)
+- Total time ≈ NUM_ITERS × 1 cycle ≈ very fast
+
+**For DRAM data:**
+- Data in global memory
+- Each iteration reloads from DRAM
+- Total time ≈ NUM_ITERS × 200-400 cycles ≈ slow!
+
+**Result:** By comparing SRAM vs DRAM probe times, we measure the latency difference.
+
+---
+
+#### **Python Wrappers (Lines 184-265)**
+
+```python
+def _next_power_of_2(n: int) -> int:
+    return 1 << (n - 1).bit_length()
+
+def triton_rms_norm_residual_sram(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    """SRAM-optimized fused RMS-norm + residual for L-level."""
+    assert x.shape == residual.shape
+    assert x.is_contiguous() and residual.is_contiguous()
+
+    rows, N = x.shape[0] * (x.shape[1] if x.ndim == 3 else 1), x.shape[-1]
+    flat_x = x.reshape(rows, N)
+    flat_r = residual.reshape(rows, N)
+    out = torch.empty_like(flat_x)
+
+    BLOCK_N = _next_power_of_2(N)
+
+    _rms_norm_residual_fused_kernel[(rows,)](
+        flat_x, flat_r, out,
+        N=N, eps=eps, BLOCK_N=BLOCK_N,
+    )
+    return out.reshape(x.shape)
+```
+
+**Explanation:**
+
+- **`_next_power_of_2(n)`**: Find smallest power of 2 ≥ n
+  - Example: `_next_power_of_2(512)` → 512, `_next_power_of_2(513)` → 1024
+  - (Bit manipulation: `(n-1).bit_length()` gives number of bits, `1 << x` is 2^x)
+  - GPUs work best with powers of 2 (thread block sizes)
+
+- **`triton_rms_norm_residual_sram(...)`**: Python wrapper to call the Triton kernel
+  - Checks inputs are same shape and contiguous
+  - Flattens to 2D (rows × hidden_size)
+  - Creates output tensor
+  - Picks block size (nearest power of 2)
+  - Launches kernel with `[(rows,)]` — one thread block per row
+  - Reshapes output back to original
+
+**Why a wrapper?** Triton kernels are GPU code; we need Python to:
+1. Prepare data (reshape, allocate output)
+2. Launch the kernel (call the JIT-compiled function)
+3. Return result to CPU
+
+---
+
+## PART 4: How PyTorch & Triton Work Together
+
+### The Full Pipeline
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ Input Data (e.g., transformer input)                            │
+└────────────────────────────────┬────────────────────────────────┘
+                                 │
+                    ┌────────────▼────────────┐
+                    │ MemoryTierManager       │
+                    │ allocates L-level ────┐ │
+                    │ allocates H-level ──┐ │ │
+                    └────────────────┬────┘ │ │
+                                     │      │ │
+                ┌────────────────────▼──┐   │ │
+                │ L-level (Fast Path)   │   │ │
+                │ Triton kernels:       │   │ │
+                │ - SRAM RMS+Residual   │   │ │
+                │ - SRAM SwiGLU         │   │ │
+                │ (Data in regs/cache)  │   │ │
+                └────────────────┬──────┘   │ │
+                                 │          │ │
+            ┌────────────────────▼────────┐ │ │
+            │ H-level (Slow Path)        │ │ │
+            │ Triton kernels:            │ │ │
+            │ - DRAM RMS+Residual       │ │ │
+            │ (Data in global memory)    │ │ │
+            └────────────────┬───────────┘ │ │
+                             │             │ │
+         ┌───────────────────▼────────────┐│ │
+         │ Triton State Transfer:         ││ │
+         │ Move L→H and H→L states ───────┼┘ │
+         │ (what MemoryTierManager times) │  │
+         └───────────────────┬────────────┘  │
+                             │               │
+                ┌────────────▼────────────┐  │
+                │ Output                  │  │
+                │ (result tensor)         │  │
+                └────────────────────────┘  │
+                                            │
+                        ┌───────────────────▼
+                        │ Statistics collected:
+                        │ - L/H latency
+                        │ - Transfer time
+                        │ - Memory usage
+                        │ - Hit rate
+```
+
+### Real Example: Forward Pass
+
+```python
+# 1. Create memory manager
+mem_mgr = MemoryTierManager(device='cuda')
+
+# 2. Allocate L-level (fast) hidden state
+L_hidden = mem_mgr.alloc_sram('L_state', shape=(batch, hidden_dim), dtype=torch.float32)
+# → Records allocation in SRAM if it fits, otherwise spills to DRAM
+
+# 3. Allocate H-level (slow) hidden state
+H_hidden = mem_mgr.alloc_dram('H_state', shape=(batch, hidden_dim), dtype=torch.float32)
+# → Records allocation in DRAM
+
+# 4. L-level computes using SRAM kernel (fast)
+with mem_mgr.sram_context():
+    L_out = triton_rms_norm_residual_sram(L_hidden, residual, eps=1e-5)
+    L_out = triton_swiglu_sram(gate_up, inter_dim)
+    # All data in registers/shared memory
+    # Time recorded: ~10-100 microseconds per operation
+
+# 5. H-level computes using DRAM kernel (slower)
+with mem_mgr.dram_context():
+    H_out = triton_rms_norm_residual_dram(H_hidden, residual, eps=1e-5)
+    # Data loaded from global memory each time
+    # Time recorded: ~100-1000 microseconds per operation
+
+# 6. Transfer L-level output to H-level
+L_to_H_output = mem_mgr.transfer_sram_to_dram('L_out')
+# → Records transfer size and time
+
+# 7. Collect metrics
+stats = mem_mgr.get_stats()
+print(f"L/H latency ratio: {stats['H_latency'] / stats['L_latency']}")
+# → Usually 10-100x difference
+```
+
+---
+
+## PART 5: Bringing It All Together
+
+### What Happens When You Run a Benchmark
+
+```
+python run_benchmark.py --mode tiered --iterations 10
+```
+
+1. **Setup:**
+   - Create MemoryTierManager
+   - Create model (HRM_Tiered)
+   - Create dummy batch of data
+
+2. **Warmup phase (first 5 runs):**
+   - Forward pass (not timed)
+   - GPU caches warm up, compilers warm up
+   - Discarded from results
+
+3. **Benchmark phase (10 timed runs):**
+   - For each iteration:
+     - Start GPU timer
+     - Forward pass:
+       - L-level uses SRAM kernels (fast)
+       - H-level uses DRAM kernels (slow)
+       - Record latencies in mem_mgr
+     - Stop GPU timer
+     - Record elapsed time
+
+4. **Statistics:**
+   - Calculate mean, min, max, std dev of times
+   - Get memory stats (peak, current usage)
+   - Get transfer stats (H↔L copy times)
+   - Calculate derived metrics (ratios, efficiency)
+   - Get Triton probe latencies (direct SRAM vs DRAM measurement)
+
+5. **Output:**
+   - Print table of results
+   - Save to JSON
+   - Optionally generate plots
+
+---
+
+## Key Insights
+
+### Why Two Implementations?
+
+| Aspect | PyTorch | Triton |
+|--------|---------|--------|
+| **What it expresses** | Algorithmic intent | Micro-architecture intent |
+| **Memory control** | Coarse (allocate tensor) | Fine (exact register usage) |
+| **Performance** | Relies on library kernels | Direct GPU control |
+| **Latency** | Can hide memory issues | Exposes memory latency differences |
+| **Usability** | Easy to write | Low-level, harder to write |
+
+### Why SRAM vs DRAM?
+
+**The 200× Latency Difference:**
+- SRAM (L2 cache): Load time ~4 cycles = ~2 nanoseconds = **very fast**
+- DRAM (HBM): Load time ~400 cycles = **200 nanoseconds = 200x slower!**
+
+By **exploiting this difference**, we can:
+- Keep "important" (L-level) data in fast SRAM
+- Allow "less important" (H-level) data to use slow DRAM
+- Model hierarchical computation: fast frequent + slow infrequent
+
+### Why Fuse Operations?
+
+**Without fusion:**
+```
+Load → Compute → Store → Load → Compute → Store
+```
+Multiple memory round-trips!
+
+**With fusion:**
+```
+Load → Compute → Compute → Compute → Store
+```
+One load, many operations, one store. **Massive speedup!**
+
+---
+
+## Testing It Yourself
+
+### Quick test (2 minutes):
+```bash
+cd test-env/HRM_optimised
+python run_benchmark.py --mode tiered --warmup 1 --iterations 3 --batch-sizes 2 --seq-lens 16
+```
+
+### Full benchmark (10 minutes):
+```bash
+python run_benchmark.py --mode compare --warmup 5 --iterations 20 --batch-sizes 1,8,32 --seq-lens 64,128
+```
+
+### Analyze results:
+```python
+import json
+with open('benchmark_results/results.json') as f:
+    results = json.load(f)
+print(f"L/H latency ratio: {results[0]['h_over_l_latency_ratio']}")
+print(f"SRAM hit rate: {results[0]['sram_hit_rate']}")
+print(f"Memory efficiency: {results[0]['memory_efficiency']}")
+```
+
+---
+
+## Summary
+
+You now understand:
+
+✅ **Benchmarks** → What metrics are collected and why  
+✅ **PyTorch layer** → How `MemoryTierManager` tracks two memory tiers  
+✅ **Triton layer** → How kernels control GPU memory access patterns  
+✅ **Integration** → How they work together to model hierarchical computation  
+✅ **Performance** → Why SRAM vs DRAM and data fusion matter for speed  
+
+The key idea: **Different parts of the model operate at different timescales (L-level fast, H-level slow), and we can model this using GPU memory hierarchy (SRAM fast, DRAM slow).**

docs/IMPLEMENTATION.md

@@ -0,0 +1,49 @@
+**Overview**
+- **Purpose**: Explains how the HRM model implements a two-tier memory hierarchy (L-level ≈ SRAM, H-level ≈ DRAM) in both PyTorch and Triton, and how the two implementations map to each other.
+
+**Files to reference**
+- `PyTorch ref`: [softmax.py](softmax.py)
+- `Triton ref`: [softmax_triton.py](softmax_triton.py)
+- `Memory tier manager`: [test-env/HRM_optimised/models/memory_tier.py](test-env/HRM_optimised/models/memory_tier.py)
+- `Triton kernels`: [test-env/HRM_optimised/models/triton_kernels.py](test-env/HRM_optimised/models/triton_kernels.py)
+- `Sparse embedding + optimizer`: [test-env/HRM_optimised/models/sparse_embedding.py](test-env/HRM_optimised/models/sparse_embedding.py)
+
+**Conceptual summary**
+- **Two-tier idea**: L-level (fast, frequently-updated state) is treated as SRAM: kept resident and reused heavily inside a compute tile. H-level (large, infrequently-updated state) is treated as DRAM: loaded from global memory when needed.
+- **Why two implementations**: PyTorch expresses algorithmic intent using tensors, streams and explicit copies; Triton encodes the micro-kernel behavior (register/shared-memory reuse, explicit loads/stores) so we can enforce the intended locality and measure latency/bandwidth differences.
+
+**PyTorch implementation (what to look for)**
+- **`MemoryTierManager` (`memory_tier.py`)**: central piece for the PyTorch-side memory-tiering.
+  - **Allocation**: `alloc_sram` and `alloc_dram` allocate tensors and track sizes. SRAM allocation is guarded by a capacity limit (spill-to-DRAM on overflow).
+  - **Contexts / Streams**: `sram_context` and `dram_context` provide CUDA stream scopes to suggest where operations should run (helps overlap/ordering and can hint locality).
+  - **Transfers**: `transfer_sram_to_dram` and `transfer_dram_to_sram` perform explicit `clone()` copies and record timing via CPU timers with `torch.cuda.synchronize()` for correctness.
+  - **Tracking & metrics**: `MemoryEvent` log and `TierStats` collect hit/miss counts, transfer counts, and timing statistics used for benchmarking.
+- **Pure-PyTorch ops**: Higher-level routines (e.g. a softmax demo in `softmax.py`) use standard `torch`/`torch.nn.functional` ops. These rely on the allocator/streams above to approximate tier behavior but cannot force register/shared-memory residency the way Triton kernels do.
+- **Sparse embedding** (`sparse_embedding.py`): shows a training path where local slices are copied into a local buffer for computation and a distributed SignSGD step reduces gradients and writes back slices — this mirrors a “local fast working set + global parameter store” design.
+
+**Triton implementation (what to look for)**
+- **Kernels are explicit**: `triton_kernels.py` contains JIT-ed kernels that declare which data is intended to remain in SRAM (registers/shared-memory) vs. be sourced from DRAM (global memory).
+  - **`constexpr` parameters** (e.g. `N`, `BLOCK_N`) cause the Triton compiler to tile loops and keep whole tiles in registers/shared memory. This is how L-level residency is enforced.
+  - **SRAM kernels**: `_rms_norm_residual_fused_kernel` and `_swiglu_fused_sram_kernel` load tile data into registers/shared memory, perform fused compute (residual add, RMS-norm, SwiGLU) entirely in-register, then store results — minimal global memory traffic.
+  - **DRAM kernels**: `_rms_norm_residual_dram_kernel` uses the same math but is intended to be invoked without tile reuse; each invocation does full DRAM round-trips.
+  - **State transfer**: `_state_transfer_kernel` is an explicit copy kernel used to move state between tiers (H↔L). This is the Triton analogue of `transfer_*` in `MemoryTierManager` but compiled down to efficient device copies.
+  - **Latency probes**: `_memory_latency_probe_kernel` builds dependent load chains to measure effective latency and expose the practical difference between SRAM-like reuse and DRAM round trips.
+  - **Wrappers** (Python functions at the bottom of `triton_kernels.py`): reshape inputs, pick grid/block sizes, and launch the kernels (e.g. `triton_rms_norm_residual_sram`, `triton_state_transfer`). These are the call sites you'd use in place of the plain PyTorch ops where you need the explicit micro-architecture behavior.
+
+**How the mapping works: PyTorch ↔ Triton**
+- **Operator mapping**: Where PyTorch would call `x + residual` followed by `F.layer_norm` or `F.softmax`, the Triton path provides fused kernels that implement the same math but with different memory patterns (fused to avoid intermediate writes and to increase register/shared-memory reuse).
+- **Memory behavior**: In PyTorch the best you can do is allocate tensors on particular streams and do explicit copies; residency in registers/shared memory is controlled by the backend and cannot be guaranteed. Triton lets you express the exact dataflow within a kernel so you can keep a tile in registers/shared memory across many operations.
+- **Performance implications**: Triton kernels reduce global memory bandwidth by reusing tile data (L-level) and reduce kernel-launch overhead by fusing ops; PyTorch code is simpler and relies on library kernels but may incur extra memory traffic and intermediate buffers.
+
+**Practical notes & where to change things**
+- To use Triton kernels in place of PyTorch ops, replace the PyTorch call sites with the wrappers in `triton_kernels.py` (for example, use `triton_rms_norm_residual_sram(...)` for L-level paths and `triton_rms_norm_residual_dram(...)` for H-level).
+- Tune tile sizes (`BLOCK_N`, `BLOCK_INTER`) by setting `constexpr` arguments and the `_next_power_of_2` helper. These control register/shared-memory usage vs parallelism.
+- Use `triton_state_transfer` to implement the same explicit copies that `MemoryTierManager.transfer_*` performs; Triton copy kernels can be more efficient than `.clone()` in tight loops.
+
+**References / next steps**
+- Inspect the call sites in the model that route L-level vs H-level computation and swap in the Triton wrappers to validate behavior and measure latency.
+- Use `MemoryTierManager.get_events()` and the Triton `memory_latency_probe` to compare the effective behavior end-to-end.
+
+If you want, I can:
+- add small call-site examples showing the PyTorch call and the Triton replacement, or
+- run a micro-benchmark that compares the two paths and append results to this doc.

eval_dummy.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""
+eval_dummy.py — Evaluate Tiered vs Baseline HRM using dummy (random) tensors.
+
+No datasets or checkpoints needed. This measures pure model throughput,
+latency, and memory usage on synthetic inputs.
+
+Usage:
+    source venv/bin/activate
+    python eval_dummy.py
+    python eval_dummy.py --batch-size 64 --seq-len 256 --hidden-size 1024
+    python eval_dummy.py --iterations 50 --plot
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from dataclasses import dataclass, asdict
+
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from models.memory_tier import MemoryTierManager
+from models.hrm.hrm_tiered import HRM_Tiered
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
+
+
+# =====================================================================
+#  Dummy batch generator
+# =====================================================================
+
+def make_batch(batch_size, seq_len, vocab_size, device):
+    return {
+        "inputs": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
+        "labels": torch.randint(0, vocab_size, (batch_size, seq_len), device=device),
+        "puzzle_identifiers": torch.arange(batch_size, device=device),
+    }
+
+
+# =====================================================================
+#  Single-model evaluation
+# =====================================================================
+
+@dataclass
+class EvalResult:
+    model_name: str
+    batch_size: int
+    seq_len: int
+    hidden_size: int
+    param_count: int
+    latency_mean_ms: float
+    latency_std_ms: float
+    throughput_sps: float
+    gpu_memory_mb: float
+
+
+def eval_model(model, model_name, batch, device, warmup=5, iterations=20):
+    """Run inference on dummy data and collect timing stats."""
+    model.eval()
+    bs = batch["inputs"].shape[0]
+    sl = batch["inputs"].shape[1]
+    param_count = sum(p.numel() for p in model.parameters())
+
+    # Warmup
+    with torch.no_grad():
+        for _ in range(warmup):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+            model(carry, batch)
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+        torch.cuda.synchronize()
+
+    # Timed iterations
+    latencies = []
+    with torch.no_grad():
+        for _ in range(iterations):
+            carry = model.initial_carry(batch)
+            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+            carry.steps = carry.steps.to(device)
+            carry.halted = carry.halted.to(device)
+            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+
+            if torch.cuda.is_available():
+                start = torch.cuda.Event(enable_timing=True)
+                end = torch.cuda.Event(enable_timing=True)
+                start.record()
+
+            model(carry, batch)
+
+            if torch.cuda.is_available():
+                end.record()
+                torch.cuda.synchronize()
+                latencies.append(start.elapsed_time(end))
+            else:
+                pass  # CPU fallback handled by perf_counter
+
+    mean_ms = sum(latencies) / len(latencies)
+    std_ms = (sum((x - mean_ms)**2 for x in latencies) / max(len(latencies)-1, 1)) ** 0.5
+    throughput = bs / (mean_ms / 1000) if mean_ms > 0 else 0
+    gpu_mem = torch.cuda.max_memory_allocated(device) / (1024**2) if torch.cuda.is_available() else 0
+
+    return EvalResult(
+        model_name=model_name,
+        batch_size=bs, seq_len=sl,
+        hidden_size=model.config.hidden_size,
+        param_count=param_count,
+        latency_mean_ms=mean_ms,
+        latency_std_ms=std_ms,
+        throughput_sps=throughput,
+        gpu_memory_mb=gpu_mem,
+    )
+
+
+# =====================================================================
+#  Main comparison
+# =====================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="HRM Dummy Tensor Evaluation")
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--seq-len", type=int, default=81, help="Sudoku=81, or any length")
+    parser.add_argument("--hidden-size", type=int, default=512)
+    parser.add_argument("--num-heads", type=int, default=8)
+    parser.add_argument("--H-cycles", type=int, default=2)
+    parser.add_argument("--L-cycles", type=int, default=2)
+    parser.add_argument("--H-layers", type=int, default=4)
+    parser.add_argument("--L-layers", type=int, default=4)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--iterations", type=int, default=20)
+    parser.add_argument("--output", type=str, default="benchmark_results/eval_dummy.json")
+    parser.add_argument("--plot", action="store_true")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    vocab_size = 32
+
+    config_dict = {
+        "batch_size": args.batch_size,
+        "seq_len": args.seq_len,
+        "puzzle_emb_ndim": 0,
+        "num_puzzle_identifiers": args.batch_size,
+        "vocab_size": vocab_size,
+        "H_cycles": args.H_cycles,
+        "L_cycles": args.L_cycles,
+        "H_layers": args.H_layers,
+        "L_layers": args.L_layers,
+        "hidden_size": args.hidden_size,
+        "expansion": 4.0,
+        "num_heads": args.num_heads,
+        "pos_encodings": "rope",
+        "halt_max_steps": 1,
+        "halt_exploration_prob": 0.0,
+    }
+
+    batch = make_batch(args.batch_size, args.seq_len, vocab_size, device)
+
+    print(f"\n{'='*64}")
+    print(f"  HRM Dummy Tensor Evaluation")
+    print(f"  Device:      {device} ({torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'})")
+    print(f"  Batch:       {args.batch_size}")
+    print(f"  Seq Len:     {args.seq_len}")
+    print(f"  Hidden:      {args.hidden_size}")
+    print(f"  H/L cycles:  {args.H_cycles}/{args.L_cycles}")
+    print(f"  H/L layers:  {args.H_layers}/{args.L_layers}")
+    print(f"  Iterations:  {args.iterations} (warmup: {args.warmup})")
+    print(f"{'='*64}")
+
+    results = []
+
+    # ── Baseline ──
+    print(f"\n  → Evaluating Baseline (hrm_act_v1)...")
+    baseline_model = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
+    r_baseline = eval_model(baseline_model, "Baseline", batch, device, args.warmup, args.iterations)
+    results.append(r_baseline)
+    del baseline_model
+    torch.cuda.empty_cache()
+
+    # ── Tiered ──
+    print(f"  → Evaluating Tiered (SRAM/DRAM Triton)...")
+    mem_mgr = MemoryTierManager(device=device, enable_tracking=True)
+    tiered_model = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
+    r_tiered = eval_model(tiered_model, "Tiered", batch, device, args.warmup, args.iterations)
+
+    # Grab tier-specific stats
+    tier_timing = tiered_model.get_timing_stats()
+    tier_mem = mem_mgr.get_stats()
+    results.append(r_tiered)
+    del tiered_model
+    torch.cuda.empty_cache()
+
+    # ── Print comparison table ──
+    print(f"\n{'='*64}")
+    print(f"  {'Model':<12} {'Params':>10} {'Latency(ms)':>13} {'±σ':>8} {'Throughput':>12} {'GPU MB':>8}")
+    print(f"  {'-'*58}")
+    for r in results:
+        print(f"  {r.model_name:<12} {r.param_count/1e6:>9.1f}M {r.latency_mean_ms:>13.2f} {r.latency_std_ms:>8.2f} {r.throughput_sps:>12.1f} {r.gpu_memory_mb:>8.1f}")
+
+    # Speedup
+    speedup = r_baseline.latency_mean_ms / r_tiered.latency_mean_ms if r_tiered.latency_mean_ms > 0 else 0
+    mem_diff = r_baseline.gpu_memory_mb - r_tiered.gpu_memory_mb
+    print(f"\n  Speedup:        {speedup:.2f}x")
+    print(f"  Memory saving:  {mem_diff:.1f} MB")
+
+    # Tier-specific stats
+    if tier_timing:
+        l_us = tier_timing.get("L_forward_us", {}).get("mean_us", 0)
+        h_us = tier_timing.get("H_forward_us", {}).get("mean_us", 0)
+        ratio = h_us / l_us if l_us > 0 else 0
+        print(f"\n  L-level (SRAM): {l_us:.1f} μs")
+        print(f"  H-level (DRAM): {h_us:.1f} μs")
+        print(f"  H/L ratio:      {ratio:.2f}x")
+
+    if tier_mem:
+        print(f"  SRAM hit rate:  {tier_mem['sram']['hit_rate']:.2%}")
+    print(f"{'='*64}\n")
+
+    # ── Save ──
+    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+    with open(args.output, "w") as f:
+        json.dump([asdict(r) for r in results], f, indent=2)
+    print(f"  Results saved → {args.output}")
+
+    # ── Plots ──
+    if args.plot:
+        try:
+            import matplotlib
+            matplotlib.use("Agg")
+            import matplotlib.pyplot as plt
+
+            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+            fig.suptitle("HRM Dummy Eval: Baseline vs Tiered", fontweight="bold")
+            names = [r.model_name for r in results]
+            colors = ["#e74c3c", "#2ecc71"]
+
+            # Latency
+            axes[0].bar(names, [r.latency_mean_ms for r in results], color=colors)
+            axes[0].set_ylabel("Latency (ms)")
+            axes[0].set_title("Inference Latency")
+            axes[0].grid(axis="y", alpha=0.3)
+
+            # Throughput
+            axes[1].bar(names, [r.throughput_sps for r in results], color=colors)
+            axes[1].set_ylabel("Samples/sec")
+            axes[1].set_title("Throughput")
+            axes[1].grid(axis="y", alpha=0.3)
+
+            # Memory
+            axes[2].bar(names, [r.gpu_memory_mb for r in results], color=colors)
+            axes[2].set_ylabel("GPU Memory (MB)")
+            axes[2].set_title("Peak Memory")
+            axes[2].grid(axis="y", alpha=0.3)
+
+            plt.tight_layout()
+            plot_dir = os.path.dirname(args.output) or "benchmark_results"
+            plot_path = os.path.join(plot_dir, "eval_dummy_comparison.png")
+            plt.savefig(plot_path, dpi=150)
+            plt.close()
+            print(f"  Plot saved → {plot_path}")
+        except ImportError:
+            print("  (matplotlib not found — skipping plots)")
+
+    print("  Done!\n")
+
+
+if __name__ == "__main__":
+    main()

evaluate.py

@@ -0,0 +1,68 @@
+from typing import List
+import yaml
+import os
+
+import torch
+import torch.distributed as dist
+
+import pydantic
+from omegaconf import OmegaConf
+from pretrain import PretrainConfig, init_train_state, evaluate, create_dataloader
+
+
+class EvalConfig(pydantic.BaseModel):
+    checkpoint: str
+    
+    save_outputs: List[str] = ["inputs", "labels", "puzzle_identifiers", "logits", "q_halt_logits", "q_continue_logits"]
+
+
+def launch():
+    eval_cfg = EvalConfig(**OmegaConf.to_container(OmegaConf.from_cli()))  # type: ignore
+    
+    RANK = 0
+    WORLD_SIZE = 1
+    # Initialize distributed training if in distributed environment (e.g. torchrun)
+    if "LOCAL_RANK" in os.environ:
+        # Initialize distributed, default device and dtype
+        dist.init_process_group(backend="nccl")
+
+        RANK = dist.get_rank()
+        WORLD_SIZE = dist.get_world_size()
+
+        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+    with open(os.path.join(os.path.dirname(eval_cfg.checkpoint), "all_config.yaml"), "r") as f:
+        config = PretrainConfig(**yaml.safe_load(f))
+
+        config.eval_save_outputs = eval_cfg.save_outputs
+        config.checkpoint_path = os.path.dirname(eval_cfg.checkpoint)
+
+    # Dataloader
+    train_loader, train_metadata = create_dataloader(config, "train", test_set_mode=False, epochs_per_iter=1, global_batch_size=config.global_batch_size, rank=RANK, world_size=WORLD_SIZE)
+    eval_loader,  eval_metadata  = create_dataloader(config, "test", test_set_mode=True, epochs_per_iter=1, global_batch_size=config.global_batch_size, rank=RANK, world_size=WORLD_SIZE)
+
+    # Models
+    train_state = init_train_state(config, train_metadata, world_size=WORLD_SIZE)
+    # Try unwrap torch.compile
+    try:
+        train_state.model.load_state_dict(torch.load(eval_cfg.checkpoint, map_location="cuda"), assign=True)
+    except:
+        train_state.model.load_state_dict({k.removeprefix("_orig_mod."): v for k, v in torch.load(eval_cfg.checkpoint, map_location="cuda").items()}, assign=True)
+    
+    train_state.step = 0
+    ckpt_filename = os.path.basename(eval_cfg.checkpoint)
+    if ckpt_filename.startswith("step_"):
+        train_state.step = int(ckpt_filename.removeprefix("step_"))
+
+    # Evaluate
+    print ("Starting evaluation")
+    
+    train_state.model.eval()
+    metrics = evaluate(config, train_state, eval_loader, eval_metadata, rank=RANK, world_size=WORLD_SIZE)
+
+    if metrics is not None:
+        print (metrics)
+
+
+if __name__ == "__main__":
+    launch()

fusedevals.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""
+Evaluate Baseline (Traditional) vs Tiered (Fused) HRM models on the specified dataset.
+Both models are evaluated using the EXACT SAME weights to verify correctness and compare speed.
+"""
+
+import os
+import sys
+import yaml
+import time
+import argparse
+
+import torch
+import numpy as np
+from safetensors.torch import load_file
+from omegaconf import OmegaConf
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from pretrain import PretrainConfig, create_dataloader
+from models.hrm.hrm_tiered import HRM_Tiered
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
+from models.losses import ACTLossHead
+from models.memory_tier import MemoryTierManager
+
+def build_model(arch, config_dict, device):
+    if arch == "tiered":
+        mm = MemoryTierManager(device=device, enable_tracking=True)
+        model = HRM_Tiered(config_dict, memory_manager=mm)
+    else:
+        model = HierarchicalReasoningModel_ACTV1(config_dict)
+    
+    # Wrap in ACTLossHead exactly as pretrain.py does
+    loss_head = ACTLossHead(model, loss_type="stablemax_cross_entropy").to(device)
+    loss_head.eval()
+    return loss_head, (mm if arch == "tiered" else None)
+
+@torch.no_grad()
+def benchmark_model(model_name, model, dataloader, metadata, device):
+    print(f"\n[{model_name}] Starting Evaluation on dataset...")
+    
+    model.eval()
+    all_metrics = []
+    
+    start_time = time.perf_counter()
+    total_samples = 0
+    
+    # We will accumulate the exact accuracy matching evaluate.py
+    total_accuracy = 0
+    total_exact_accuracy = 0
+    total_count = 0
+    
+    for set_name, batch, batch_size in dataloader:
+        batch = {k: v.to(device) for k, v in batch.items()}
+        # ACTLossHead wraps initial_carry
+        with torch.device(device):
+            carry = model.initial_carry(batch)
+        
+        while True:
+            carry, loss, metrics, _, all_finish = model(carry=carry, batch=batch, return_keys=[])
+            if all_finish:
+                break
+                
+        total_accuracy += metrics["accuracy"].item()
+        total_exact_accuracy += metrics["exact_accuracy"].item()
+        total_count += metrics["count"].item()
+        total_samples += batch_size
+
+    # Synchronize GPU to ensure timing is correct
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+        
+    end_time = time.perf_counter()
+    duration = end_time - start_time
+    
+    acc = total_accuracy / max(total_count, 1)
+    exact_acc = total_exact_accuracy / max(total_count, 1)
+    throughput = total_samples / duration
+    
+    print(f"[{model_name}] Results:")
+    print(f"  Duration:   {duration:.2f}s")
+    print(f"  Throughput: {throughput:.1f} samples/sec")
+    print(f"  Token Acc:  {acc*100:.2f}%")
+    print(f"  Exact Acc:  {exact_acc*100:.2f}%")
+    
+    return {
+        "duration_s": duration,
+        "throughput": throughput,
+        "token_acc": acc,
+        "exact_acc": exact_acc
+    }
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from  matplotlib.gridspec import GridSpec
+import multiprocessing as mp
+
+def run_evaluation(arch, model_cfg, state_dict, data_path, global_batch_size, gpu_id, result_queue):
+    device = torch.device(f"cuda:{gpu_id}")
+    torch.cuda.set_device(device)
+    
+    # Needs to recreate dataloader per process
+    cfg_container = {
+        "arch": model_cfg,
+        "data_path": data_path,
+        "global_batch_size": global_batch_size,
+        "epochs": 1, "lr": 7e-5, "lr_min_ratio": 1.0, "lr_warmup_steps": 2000,
+        "weight_decay": 1.0, "beta1": 0.9, "beta2": 0.95,
+        "puzzle_emb_lr": 7e-5, "puzzle_emb_weight_decay": 1.0,
+        "seed": 0
+    }
+    config = PretrainConfig(**cfg_container)
+    eval_loader, eval_metadata = create_dataloader(
+        config, "test", test_set_mode=True, epochs_per_iter=1, 
+        global_batch_size=config.global_batch_size, rank=0, world_size=1
+    )
+
+    model_cfg = model_cfg.copy()
+    model_cfg.update({
+        "batch_size": global_batch_size,
+        "vocab_size": eval_metadata.vocab_size,
+        "seq_len": eval_metadata.seq_len,
+        "num_puzzle_identifiers": eval_metadata.num_puzzle_identifiers,
+        "causal": False
+    })
+    
+    model_name = "Traditional HRM" if arch == "baseline" else "Fused Tiered HRM"
+    model, _ = build_model(arch, model_cfg, device)
+    
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except:
+        model.load_state_dict(state_dict, strict=False)
+        
+    res = benchmark_model(model_name, model, eval_loader, eval_metadata, device)
+    res["arch"] = arch
+    result_queue.put(res)
+
+def create_comparison_plots(base_res, tier_res, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    
+    c_base, c_tier = "#4A90D9", "#E85D75"
+    bg, text, grid = "#1a1a2e", "#e0e0e0", "#333355"
+
+    plt.rcParams.update({
+        "figure.facecolor": bg, "axes.facecolor": "#16213e",
+        "axes.edgecolor": grid, "axes.labelcolor": text,
+        "text.color": text, "xtick.color": text, "ytick.color": text,
+        "grid.color": grid, "grid.alpha": 0.3,
+        "font.family": "sans-serif", "font.size": 11,
+    })
+
+    fig = plt.figure(figsize=(15, 6))
+    fig.suptitle("Sudoku Extreme: Traditional vs Fused Tiered HRM", fontsize=16, fontweight="bold", y=0.98)
+    gs = GridSpec(1, 3, figure=fig, wspace=0.3)
+    labels = ["Traditional (v1)", "Fused Tiered"]
+
+    def bar_plot(ax, title, ylabel, vals, fmt=".2f", is_percent=False):
+        bars = ax.bar(labels, vals, color=[c_base, c_tier], edgecolor="white", width=0.5)
+        ax.set_title(title, fontweight="bold")
+        ax.set_ylabel(ylabel)
+        for b, v in zip(bars, vals):
+            val_str = f"{v*100:{fmt}}%" if is_percent else f"{v:{fmt}}"
+            ax.text(b.get_x() + b.get_width()/2, b.get_height() * 1.02,
+                    val_str, ha="center", fontsize=11, color=text, fontweight="bold")
+        ax.grid(axis="y")
+        if is_percent: ax.set_ylim(0, 1.1)
+
+    # 1. Throughput
+    bar_plot(fig.add_subplot(gs[0, 0]), "Inference Throughput", "Samples / Second",
+             [base_res["throughput"], tier_res["throughput"]], fmt=".1f")
+
+    # 2. Token Accuracy
+    bar_plot(fig.add_subplot(gs[0, 1]), "Token Accuracy", "Accuracy",
+             [base_res["token_acc"], tier_res["token_acc"]], is_percent=True)
+
+    # 3. Exact Match Accuracy
+    bar_plot(fig.add_subplot(gs[0, 2]), "Exact Puzzle Accuracy", "Accuracy",
+             [base_res["exact_acc"], tier_res["exact_acc"]], is_percent=True)
+
+    path = os.path.join(output_dir, "eval_fused_vs_v1_comparison.png")
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"\n  Plot saved → {path}")
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--weights", type=str, default="hf_upload/tiered_hrm_sram_dram/model.safetensors")
+    parser.add_argument("--output-dir", type=str, default="benchmark_results")
+    args = parser.parse_args()
+    
+    # Base Configuration
+    model_cfg = {
+        "name": "hrm.hrm_tiered@HRM_Tiered",
+        "loss": {"name": "losses@ACTLossHead", "loss_type": "stablemax_cross_entropy"},
+        "hidden_size": 512, "num_heads": 8, "expansion": 4,
+        "H_layers": 4, "L_layers": 4, "H_cycles": 2, "L_cycles": 2,
+        "halt_max_steps": 16, "halt_exploration_prob": 0.1,
+        "pos_encodings": "rope", "puzzle_emb_ndim": 512,
+        "batch_size": 384, "vocab_size": 32, "seq_len": 81,
+        "num_puzzle_identifiers": 384, "causal": False
+    }
+    data_path = "data/sudoku-extreme-1k-aug-1000"
+    
+    print(f"Loading weights from: {args.weights}")
+    try:
+        state_dict = load_file(args.weights)
+    except Exception as e:
+        print(f"Could not load as safetensors, falling back to torch.load... ({e})")
+        raw_state_dict = torch.load(args.weights, map_location="cpu", weights_only=True)
+        # Strip torch.compile prefix just in case as evaluate.py does
+        state_dict = {k.removeprefix("_orig_mod."): v for k, v in raw_state_dict.items()}
+        
+    if "model.inner.embed_tokens.embedding_weight" not in state_dict and "inner.embed_tokens.embedding_weight" in state_dict:
+        state_dict = {f"model.{k}": v for k, v in state_dict.items()}
+
+    mp.set_start_method('spawn', force=True)
+    ctx = mp.get_context('spawn')
+    queue = ctx.Queue()
+
+    print("\nStarting Parallel Evaluation on 2 GPUs...")
+    print("  Traditional HRM -> GPU 0")
+    print("  Fused Tiered HRM  -> GPU 1")
+    
+    # Launch parallel processes
+    p1 = ctx.Process(target=run_evaluation, args=("baseline", model_cfg.copy(), state_dict, data_path, 384, 0, queue))
+    p2 = ctx.Process(target=run_evaluation, args=("tiered", model_cfg.copy(), state_dict, data_path, 384, 1, queue))
+    
+    p1.start()
+    p2.start()
+    
+    p1.join()
+    p2.join()
+
+    # Collect Results
+    results = {}
+    while not queue.empty():
+        res = queue.get()
+        results[res["arch"]] = res
+
+    if "baseline" in results and "tiered" in results:
+        base_res = results["baseline"]
+        tier_res = results["tiered"]
+
+        print("\n" + "="*50)
+        print("      FINAL COMPARISON OVERVIEW")
+        print("="*50)
+        print(f"Token Accuracy:  Traditional {base_res['token_acc']*100:.2f}%  vs  Fused {tier_res['token_acc']*100:.2f}%")
+        print(f"Exact Accuracy:  Traditional {base_res['exact_acc']*100:.2f}%  vs  Fused {tier_res['exact_acc']*100:.2f}%")
+        print(f"Throughput:      Traditional {base_res['throughput']:.1f} samp/s  vs  Fused {tier_res['throughput']:.1f} samp/s")
+        print(f"Speedup Margin:  {tier_res['throughput'] / base_res['throughput']:.2f}x")
+        
+        # Plotting
+        create_comparison_plots(base_res, tier_res, args.output_dir)
+    else:
+        print("\nEvaluation failed. One or both models did not return results.")
+    
+if __name__ == "__main__":
+    main()

hf_upload_new/README.md

@@ -0,0 +1,10 @@
+---
+tags:
+- hrm
+- sudoku
+- pytorch
+- safetensors
+---
+# HRM Multi-GPU Trained Model
+Latest checkpoint (step 26040) from the multi-GPU training run.
+Accuracy: 60% Exact Match.

hf_upload_new/config.yaml

@@ -0,0 +1,35 @@
+arch:
+  H_cycles: 2
+  H_layers: 4
+  L_cycles: 2
+  L_layers: 4
+  expansion: 4
+  halt_exploration_prob: 0.1
+  halt_max_steps: 16
+  hidden_size: 512
+  loss:
+    loss_type: stablemax_cross_entropy
+    name: losses@ACTLossHead
+  name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
+  num_heads: 8
+  pos_encodings: rope
+  puzzle_emb_ndim: 512
+beta1: 0.9
+beta2: 0.95
+checkpoint_every_eval: true
+checkpoint_path: checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/multi_gpu_new_run
+data_path: data/sudoku-extreme-1k-aug-1000
+epochs: 20000
+eval_interval: 2000
+eval_save_outputs: []
+global_batch_size: 768
+lr: 0.0001
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+project_name: Sudoku-extreme-1k-aug-1000 ACT-torch
+puzzle_emb_lr: 0.0001
+puzzle_emb_weight_decay: 1.0
+run_name: multi_gpu_new_run
+seed: 0
+skip_eval: false
+weight_decay: 1.0

hf_upload_new/hrm_act_v1.py

@@ -0,0 +1,283 @@
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from pydantic import BaseModel
+
+from models.common import trunc_normal_init_
+from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
+from models.sparse_embedding import CastedSparseEmbedding
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1Carry:
+    inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry
+    
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+
+class HierarchicalReasoningModel_ACTV1Config(BaseModel):
+    batch_size: int
+    seq_len: int
+    puzzle_emb_ndim: int = 0
+    num_puzzle_identifiers: int
+    vocab_size: int
+
+    H_cycles: int
+    L_cycles: int
+
+    H_layers: int
+    L_layers: int
+
+    # Transformer config
+    hidden_size: int
+    expansion: float
+    num_heads: int
+    pos_encodings: str
+
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    
+    # Halting Q-learning config
+    halt_max_steps: int
+    halt_exploration_prob: float
+
+    forward_dtype: str = "bfloat16"
+
+
+class HierarchicalReasoningModel_ACTV1Block(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+
+        self.self_attn = Attention(
+            hidden_size=config.hidden_size,
+            head_dim=config.hidden_size // config.num_heads,
+            num_heads=config.num_heads,
+            num_key_value_heads=config.num_heads,
+            causal=False
+        )
+        self.mlp = SwiGLU(
+            hidden_size=config.hidden_size,
+            expansion=config.expansion,
+        )
+        self.norm_eps = config.rms_norm_eps
+
+    def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Post Norm
+        # Self Attention
+        hidden_states = rms_norm(hidden_states + self.self_attn(cos_sin=cos_sin, hidden_states=hidden_states), variance_epsilon=self.norm_eps)
+        # Fully Connected
+        hidden_states = rms_norm(hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps)
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
+    def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Input injection (add)
+        hidden_states = hidden_states + input_injection
+        # Layers
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+        self.config = config
+        self.forward_dtype = getattr(torch, self.config.forward_dtype)
+
+        # I/O
+        self.embed_scale  = math.sqrt(self.config.hidden_size)
+        embed_init_std = 1.0 / self.embed_scale
+
+        self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        self.lm_head      = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.q_head       = CastedLinear(self.config.hidden_size, 2, bias=True)
+
+        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  # ceil div
+        if self.config.puzzle_emb_ndim > 0:
+            # Zero init puzzle embeddings
+            self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
+                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)
+
+        # LM Blocks
+        if self.config.pos_encodings == "rope":
+            self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads,
+                                              max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
+                                              base=self.config.rope_theta)
+        elif self.config.pos_encodings == "learned":
+            self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        else:
+            raise NotImplementedError()
+
+        # Reasoning Layers
+        self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)])
+        self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)])
+        
+        # Initial states
+        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+
+        # Q head special init
+        # Init Q to (almost) zero for faster learning during bootstrapping
+        with torch.no_grad():
+            self.q_head.weight.zero_()
+            self.q_head.bias.fill_(-5)  # type: ignore
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor):
+        # Token embedding
+        embedding = self.embed_tokens(input.to(torch.int32))
+
+        # Puzzle embeddings
+        if self.config.puzzle_emb_ndim > 0:
+            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
+            
+            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
+            if pad_count > 0:
+                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))
+
+            embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2)
+
+        # Position embeddings
+        if self.config.pos_encodings == "learned":
+            # scale by 1/sqrt(2) to maintain forward variance
+            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))
+
+        # Scale
+        return self.embed_scale * embedding
+
+    def empty_carry(self, batch_size: int):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+            z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+        )
+        
+    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
+            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
+        )
+
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        seq_info = dict(
+            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
+        )
+
+        # Input encoding
+        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])
+
+        # Forward iterations
+        with torch.no_grad():
+            z_H, z_L = carry.z_H, carry.z_L
+
+            for _H_step in range(self.config.H_cycles):
+                for _L_step in range(self.config.L_cycles):
+                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
+                        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+
+                if not (_H_step == self.config.H_cycles - 1):
+                    z_H = self.H_level(z_H, z_L, **seq_info)
+
+        assert not z_H.requires_grad and not z_L.requires_grad
+
+        # 1-step grad
+        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.H_level(z_H, z_L, **seq_info)
+
+        # LM Outputs
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+
+        # Q head
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+        
+        return new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+
+
+class HierarchicalReasoningModel_ACTV1(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
+        self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return HierarchicalReasoningModel_ACTV1Carry(
+            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
+            
+            steps=torch.zeros((batch_size, ), dtype=torch.int32),
+            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
+            
+            current_data={k: torch.empty_like(v) for k, v in batch.items()}
+        )
+        
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+        # Update data, carry (removing halted sequences)
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}
+
+        # Forward inner model
+        new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits
+        }
+        
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            
+            halted = is_last_step
+
+            # if training, and ACT is enabled
+            if self.training and (self.config.halt_max_steps > 1):
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                halted = halted | (q_halt_logits > q_continue_logits)
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
+
+                halted = halted & (new_steps >= min_halt_steps)
+
+                # Compute target Q
+                # NOTE: No replay buffer and target networks for computing target Q-value.
+                # As batch_size is large, there're many parallel envs.
+                # Similar concept as PQN https://arxiv.org/abs/2407.04811
+                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-1]
+                
+                outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits)))
+
+        return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs

hf_upload_new/losses.py

@@ -0,0 +1,131 @@
+from typing import Any, Tuple, Dict, Sequence, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+# This ID tells the code to "skip" or ignore certain data points (like padding)
+IGNORE_LABEL_ID = -100
+
+
+def s(x, epsilon=1e-30):
+    """
+    A helper function called 'StableMax' transformation.
+    Instead of standard Softmax, this is more 'stable' for the GPU.
+    If a number is negative, it squashes it; if positive, it slightly increases it.
+    """
+    return torch.where(
+        x < 0,
+        1 / (1 - x + epsilon),
+        x + 1
+    )
+
+
+def log_stablemax(x, dim=-1):
+    """
+    Calculates the 'Log-Probability'.
+    In AI, we prefer working with 'Logs' because adding small numbers is 
+    much safer for computers than multiplying them (which can lead to 0).
+    """
+    s_x = s(x)
+    return torch.log(s_x / torch.sum(s_x, dim=dim, keepdim=True))
+
+
+def stablemax_cross_entropy(logits, labels, ignore_index: int = -100):
+    """
+    The main 'Comparison' function. 
+    It compares the model's guess (logits) to the real answer (labels).
+    The higher the score, the worse the model's guess was.
+    """
+    # 1. Convert the model's numbers into probabilities
+    logprobs = log_stablemax(logits.to(torch.float64), dim=-1)
+
+    # 2. figure out which cells actually have numbers (and aren't just empty padding)
+    valid_mask = labels != ignore_index
+    transformed_labels = torch.where(valid_mask, labels, 0)
+    
+    # 3. Pick out the probability for the correct answer
+    prediction_logprobs = torch.gather(logprobs, index=transformed_labels.to(torch.long).unsqueeze(-1), dim=-1).squeeze(-1)
+
+    # 4. Return the negative of that probability (our 'Punishment' score)
+    return -torch.where(valid_mask, prediction_logprobs, 0)
+
+
+def softmax_cross_entropy(logits, labels, ignore_index: int = -100):
+    # Cast logits to f32
+    # Flatten logits
+    return F.cross_entropy(logits.to(torch.float32).view(-1, logits.shape[-1]), labels.to(torch.long).view(-1), ignore_index=ignore_index, reduction="none").view(labels.shape)
+
+
+class ACTLossHead(nn.Module):
+    """
+    The 'Reasoning Manager'. 
+    This class wraps the model and decides:
+    1. Did the model solve it?
+    2. Should it stop thinking?
+    3. How many steps did it take?
+    """
+    def __init__(self, model: nn.Module, loss_type: str):
+        super().__init__()
+        self.model = model
+        self.loss_fn = globals()[loss_type]
+        
+    def initial_carry(self, *args, **kwargs):
+        """Prepares the model's 'Internal Memory' for a new puzzle."""
+        return self.model.initial_carry(*args, **kwargs)  # type: ignore
+
+    def forward(
+        self,
+        return_keys: Sequence[str],
+        # Model args
+        **model_kwargs,
+    ) -> Tuple[Any, torch.Tensor, Dict[str, torch.Tensor], Optional[Dict[str, torch.Tensor]], torch.Tensor]:
+        # Model logits
+        # B x SeqLen x D
+        new_carry, outputs = self.model(**model_kwargs)
+        labels = new_carry.current_data["labels"]
+
+        # 2. Check the results (without keeping track of 'learning math' here)
+        with torch.no_grad():
+            mask = labels != IGNORE_LABEL_ID
+            loss_counts = mask.sum(-1)
+            loss_divisor = loss_counts.clamp_min(1).unsqueeze(-1)
+
+            # Did the model pick the right numbers?
+            is_correct = mask & (torch.argmax(outputs["logits"], dim=-1) == labels)
+            # Is the WHOLE puzzle correct?
+            seq_is_correct = is_correct.sum(-1) == loss_counts
+            
+            # --- Store metrics for the human to read(When it halts) ---
+            valid_metrics = new_carry.halted & (loss_counts > 0)
+            metrics = {
+                "count": valid_metrics.sum(),
+                "accuracy":       torch.where(valid_metrics, (is_correct.to(torch.float32) / loss_divisor).sum(-1), 0).sum(),
+                "exact_accuracy": (valid_metrics & seq_is_correct).sum(),
+                "q_halt_accuracy": (valid_metrics & ((outputs["q_halt_logits"] >= 0) == seq_is_correct)).sum(),
+                "steps":          torch.where(valid_metrics, new_carry.steps, 0).sum(),
+            }
+
+        # 3. Calculate the learning punishments (Losses)
+        # lm_loss: Punishment for getting digits wrong
+        lm_loss = (self.loss_fn(outputs["logits"], labels, ignore_index=IGNORE_LABEL_ID) / loss_divisor).sum()
+        
+        # q_halt_loss: Punishment for stopping too early (or too late!)
+        q_halt_loss = F.binary_cross_entropy_with_logits(outputs["q_halt_logits"], seq_is_correct.to(outputs["q_halt_logits"].dtype), reduction="sum")
+
+        metrics.update({
+            "lm_loss": lm_loss.detach(),
+            "q_halt_loss": q_halt_loss.detach(),
+        })
+
+        # 4. Q-Continue Loss: Bootstrapping (The 'Look-Ahead' logic)
+        q_continue_loss = 0
+        if "target_q_continue" in outputs:
+            # Encourages the model to continue if thinking more will lead to a correct answer
+            q_continue_loss = F.binary_cross_entropy_with_logits(outputs["q_continue_logits"], outputs["target_q_continue"], reduction="sum")
+            metrics["q_continue_loss"] = q_continue_loss.detach()
+
+        # Filter outputs for return
+        detached_outputs = {k: outputs[k].detach() for k in return_keys if k in outputs}
+
+        return new_carry, lm_loss + 0.5 * (q_halt_loss + q_continue_loss), metrics, detached_outputs, new_carry.halted.all()

hf_upload_new/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e2976804165b5cca1f539d66c1726200b919bdb4845d768c73345e5c5892970
+size 109109832

latency_plot_trained_model.py

@@ -0,0 +1,417 @@
+#!/usr/bin/env python3
+"""
+latency_plot_trained_model.py — Load trained checkpoints, evaluate accuracy
+on Sudoku test data, measure inference latency, and generate combined plots.
+
+Usage:
+    source venv/bin/activate
+    python latency_plot_trained_model.py \
+        --baseline "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 belligerent-squirrel/step_52080" \
+        --tiered "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HRM_Tiered realistic-dalmatian/step_52080"
+"""
+
+import argparse
+import json
+import os
+import sys
+import yaml
+
+# Disable torch.compile — avoids 10+ min compilation during eval
+# and prevents inference_mode/compile conflicts
+os.environ["DISABLE_COMPILE"] = "1"
+
+import torch
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from pretrain import PretrainConfig, init_train_state, evaluate, create_dataloader
+
+
+# ═══════════════════════════════════════════════════════════
+#  Load a trained checkpoint
+# ═══════════════════════════════════════════════════════════
+
+def load_trained_model(ckpt_path, device="cuda"):
+    """Load checkpoint, return (train_state, config, eval_loader, latency_loader, eval_metadata).
+
+    Model hierarchy: torch.compile → ACTLossHead → ACTV1/HRM_Tiered → _Inner
+    Returns TWO eval loaders: one for accuracy (consumed by evaluate()), one for latency.
+    """
+    ckpt_dir = os.path.dirname(ckpt_path)
+    config_path = os.path.join(ckpt_dir, "all_config.yaml")
+
+    with open(config_path, "r") as f:
+        content = f.read()
+
+    if "!!python/object" not in content:
+        raw = yaml.safe_load(content)
+    else:
+        # Fallback for the irreparably mangled tiered config dump
+        print("    [Warning] Tiered config YAML is mangled, using robust fallback.")
+        raw = {
+            "arch": {
+                "name": "hrm.hrm_tiered@HRM_Tiered",
+                "hidden_size": 512,
+                "num_heads": 8,
+                "puzzle_emb_ndim": 512,
+                "pos_encodings": "rope",
+                "H_layers": 4, "H_cycles": 2,
+                "L_layers": 4, "L_cycles": 2,
+                "expansion": 4,
+                "halt_max_steps": 16,
+                "halt_exploration_prob": 0.1,
+                "memory_tier": {"sram_capacity_mb": 48, "enable_tracking": True},
+                "loss": {"loss_type": "stablemax_cross_entropy", "name": "losses@ACTLossHead"}
+            },
+            "global_batch_size": 384,
+            "skip_eval": False,
+            "eval_save_outputs": [],
+            "checkpoint_path": ckpt_dir,
+            "epochs": 20000,
+            "lr": 7.0e-05,
+            "lr_min_ratio": 1.0,
+            "lr_warmup_steps": 2000,
+            "weight_decay": 1.0,
+            "beta1": 0.9,
+            "beta2": 0.95,
+            "puzzle_emb_lr": 7.0e-05,
+            "puzzle_emb_weight_decay": 1.0,
+            "eval_interval": 2000,
+            "data_path": "data/sudoku-extreme-1k-aug-1000",
+            "project_name": "Sudoku-extreme-1k-aug-1000 ACT-torch",
+            "run_name": "HRM_Tiered realistic-dalmatian",
+            "checkpoint_every_eval": True
+        }
+        
+    config = PretrainConfig(**raw)
+    config.checkpoint_path = ckpt_dir
+
+    # Build dataloaders — need TWO because evaluate() consumes its loader
+    _, train_metadata = create_dataloader(
+        config, "train", test_set_mode=False, epochs_per_iter=1,
+        global_batch_size=config.global_batch_size, rank=0, world_size=1,
+    )
+    eval_loader, eval_metadata = create_dataloader(
+        config, "test", test_set_mode=True, epochs_per_iter=1,
+        global_batch_size=config.global_batch_size, rank=0, world_size=1,
+    )
+    latency_loader, _ = create_dataloader(
+        config, "test", test_set_mode=True, epochs_per_iter=1,
+        global_batch_size=config.global_batch_size, rank=0, world_size=1,
+    )
+
+    # Build model (torch.compile → ACTLossHead → model) and load weights
+    train_state = init_train_state(config, train_metadata, world_size=1)
+    try:
+        train_state.model.load_state_dict(
+            torch.load(ckpt_path, map_location=device, weights_only=True), assign=True
+        )
+    except Exception:
+        state = torch.load(ckpt_path, map_location=device, weights_only=True)
+        train_state.model.load_state_dict(
+            {k.removeprefix("_orig_mod."): v for k, v in state.items()}, assign=True
+        )
+
+    ckpt_name = os.path.basename(ckpt_path)
+    if ckpt_name.startswith("step_"):
+        train_state.step = int(ckpt_name.removeprefix("step_"))
+
+    train_state.model.eval()
+    return train_state, config, eval_loader, latency_loader, eval_metadata
+
+
+def unwrap_model(compiled_model):
+    """Unwrap torch.compile + ACTLossHead to get the ACTV1/HRM_Tiered wrapper.
+
+    Hierarchy: OptimizedModule._orig_mod = ACTLossHead.model = ACTV1/HRM_Tiered
+    """
+    model = compiled_model
+    # Unwrap torch.compile
+    if hasattr(model, '_orig_mod'):
+        model = model._orig_mod
+    # Unwrap ACTLossHead to get to the ACTV1/HRM_Tiered wrapper
+    if hasattr(model, 'model'):
+        model = model.model
+    return model
+
+
+# ═══════════════════════════════════════════════════════════
+#  Evaluate accuracy on real Sudoku test set
+# ═══════════════════════════════════════════════════════════
+
+def eval_accuracy(config, train_state, eval_loader, eval_metadata, limit_batches=20):
+    """Run the real evaluation on a subset of batches and return metrics dict."""
+    import itertools
+    class LimitedLoader:
+        def __init__(self, loader, limit):
+            self.loader = loader
+            self.limit = limit
+        def __iter__(self):
+            return itertools.islice(self.loader, self.limit)
+            
+    limited_eval_loader = LimitedLoader(eval_loader, limit_batches)
+    metrics = evaluate(config, train_state, limited_eval_loader, eval_metadata, rank=0, world_size=1)
+    if metrics is None:
+        return {}
+
+    # Flatten and convert to floats (skip nested dicts / non-numeric)
+    result = {}
+    for k, v in metrics.items():
+        if isinstance(v, torch.Tensor):
+            result[k] = v.item()
+        elif isinstance(v, (int, float)):
+            result[k] = float(v)
+        elif isinstance(v, dict):
+            for kk, vv in v.items():
+                if isinstance(vv, torch.Tensor):
+                    result[f"{k}/{kk}"] = vv.item()
+                elif isinstance(vv, (int, float)):
+                    result[f"{k}/{kk}"] = float(vv)
+    return result
+
+
+# ═══════════════════════════════════════════════════════════
+#  Measure inference latency on real data
+# ═══════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def measure_latency(compiled_model, eval_loader, device, warmup=3, iterations=20):
+    """Time forward pass using the unwrapped model wrapper (ACTV1/HRM_Tiered).
+
+    The unwrapped model has:
+      - initial_carry(batch) → carry
+      - forward(carry, batch) → (new_carry, outputs)
+    """
+    # Unwrap torch.compile + ACTLossHead
+    model = unwrap_model(compiled_model)
+    model.eval()
+
+    # Collect batches from the eval loader
+    batches = []
+    for set_name, batch, global_bs in eval_loader:
+        batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                 for k, v in batch.items()}
+        batches.append(batch)
+        if len(batches) >= warmup + iterations:
+            break
+
+    if not batches:
+        return {"latency_ms": 0, "latency_std": 0, "throughput": 0}
+
+    # Helper: create carry and move all tensors to device
+    def make_carry(batch):
+        carry = model.initial_carry(batch)
+        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
+        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
+        carry.steps = carry.steps.to(device)
+        carry.halted = carry.halted.to(device)
+        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
+        return carry
+
+    # Warmup
+    for i in range(min(warmup, len(batches))):
+        batch = batches[i]
+        carry = make_carry(batch)
+        model(carry, batch)
+    torch.cuda.synchronize()
+
+    # Timed runs
+    latencies = []
+    bs_total = 0
+    n_iters = min(iterations, max(1, len(batches) - warmup))
+    for i in range(n_iters):
+        batch = batches[(warmup + i) % len(batches)]
+        carry = make_carry(batch)
+
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        model(carry, batch)
+        end.record()
+        torch.cuda.synchronize()
+        latencies.append(start.elapsed_time(end))
+        bs_total += batch["inputs"].shape[0]
+
+    lat = np.array(latencies)
+    avg_bs = bs_total / len(latencies) if latencies else 1
+    return {
+        "latency_ms": float(np.mean(lat)),
+        "latency_std": float(np.std(lat)),
+        "throughput": float(avg_bs / (np.mean(lat) / 1000)) if lat.mean() > 0 else 0,
+    }
+
+
+# ═══════════════════════════════════════════════════════════
+#  Generate combined plots
+# ═══════════════════════════════════════════════════════════
+
+def create_combined_plot(base_data, tier_data, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+
+    c_base, c_tier = "#4A90D9", "#E85D75"
+    bg, text, grid = "#1a1a2e", "#e0e0e0", "#333355"
+
+    plt.rcParams.update({
+        "figure.facecolor": bg, "axes.facecolor": "#16213e",
+        "axes.edgecolor": grid, "axes.labelcolor": text,
+        "text.color": text, "xtick.color": text, "ytick.color": text,
+        "grid.color": grid, "grid.alpha": 0.3,
+        "font.family": "sans-serif", "font.size": 11,
+    })
+
+    fig = plt.figure(figsize=(18, 10))
+    fig.suptitle("HRM Trained Model Comparison: Baseline vs Tiered",
+                 fontsize=18, fontweight="bold", y=0.98)
+    gs = GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.35)
+    labels = ["Baseline", "Tiered"]
+
+    def bar_ax(ax, title, ylabel, vals, fmt=".2f"):
+        bars = ax.bar(labels, vals, color=[c_base, c_tier],
+                      edgecolor="white", linewidth=0.5, width=0.5)
+        ax.set_title(title, fontweight="bold")
+        ax.set_ylabel(ylabel)
+        for b, v in zip(bars, vals):
+            ax.text(b.get_x() + b.get_width()/2, b.get_height() * 1.02,
+                    f"{v:{fmt}}", ha="center", fontsize=11, color=text)
+        ax.grid(axis="y")
+
+    # Extract metrics with safe defaults
+    def get_acc(data, key, normalize=True):
+        count = data["accuracy"].get("eval/count", 1)
+        val = data["accuracy"].get(key, 0)
+        if normalize and count > 0:
+            return val / count * 100
+        return val
+
+    # 1. Exact Accuracy
+    bar_ax(fig.add_subplot(gs[0, 0]), "Exact Accuracy (Sudoku)", "%",
+           [get_acc(base_data, "eval/exact_accuracy"),
+            get_acc(tier_data, "eval/exact_accuracy")])
+
+    # 2. Cell Accuracy
+    bar_ax(fig.add_subplot(gs[0, 1]), "Cell-level Accuracy", "%",
+           [get_acc(base_data, "eval/accuracy"),
+            get_acc(tier_data, "eval/accuracy")])
+
+    # 3. Avg Reasoning Steps
+    bar_ax(fig.add_subplot(gs[0, 2]), "Avg Reasoning Steps (ACT)", "steps",
+           [get_acc(base_data, "eval/steps"),
+            get_acc(tier_data, "eval/steps")], fmt=".1f")
+
+    # 4. Inference Latency
+    bar_ax(fig.add_subplot(gs[1, 0]), "Inference Latency", "ms",
+           [base_data["latency"]["latency_ms"],
+            tier_data["latency"]["latency_ms"]])
+
+    # 5. Throughput
+    bar_ax(fig.add_subplot(gs[1, 1]), "Throughput", "samples/sec",
+           [base_data["latency"]["throughput"],
+            tier_data["latency"]["throughput"]], fmt=".0f")
+
+    # 6. Summary
+    ax6 = fig.add_subplot(gs[1, 2])
+    speedup = (base_data["latency"]["latency_ms"] / tier_data["latency"]["latency_ms"]
+               if tier_data["latency"]["latency_ms"] > 0 else 0)
+    base_exact = get_acc(base_data, "eval/exact_accuracy")
+    tier_exact = get_acc(tier_data, "eval/exact_accuracy")
+    summary = (
+        f"Exact Accuracy:\n"
+        f"  Baseline: {base_exact:.1f}%\n"
+        f"  Tiered:   {tier_exact:.1f}%\n\n"
+        f"Speedup: {speedup:.2f}x\n"
+        f"Throughput:\n"
+        f"  {tier_data['latency']['throughput']:.0f} vs "
+        f"{base_data['latency']['throughput']:.0f}/s"
+    )
+    ax6.text(0.5, 0.5, summary, transform=ax6.transAxes,
+             ha="center", va="center", fontsize=13, fontfamily="monospace",
+             bbox=dict(boxstyle="round,pad=0.5", facecolor="#0f3460", alpha=0.8))
+    ax6.set_title("Summary", fontweight="bold")
+    ax6.axis("off")
+
+    path = os.path.join(output_dir, "trained_model_comparison.png")
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"  Plot saved → {path}")
+    return path
+
+
+# ═══════════════════════════════════════════════════════════
+#  Main
+# ═══════════════════════════════════════════════════════════
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate trained Baseline vs Tiered HRM")
+    parser.add_argument("--baseline", type=str, required=True, help="Baseline checkpoint path")
+    parser.add_argument("--tiered", type=str, required=True, help="Tiered checkpoint path")
+    parser.add_argument("--latency-iters", type=int, default=20)
+    parser.add_argument("--output-dir", type=str, default="benchmark_results")
+    args = parser.parse_args()
+
+    device = "cuda"
+
+    print("=" * 64)
+    print("  Trained Model Comparison: Baseline vs Tiered")
+    print(f"  Device: {torch.cuda.get_device_name(0)}")
+    print("=" * 64)
+
+    results = {}
+
+    # ── Baseline ──
+    print("\n  [1/4] Loading Baseline checkpoint...")
+    base_state, base_cfg, base_eval_loader, base_lat_loader, base_eval_meta = load_trained_model(args.baseline, device)
+    n_params = sum(p.numel() for p in base_state.model.parameters()) / 1e6
+    print(f"        Step: {base_state.step}, Params: {n_params:.1f}M")
+
+    print("  [2/4] Evaluating Baseline accuracy + latency...")
+    base_acc = eval_accuracy(base_cfg, base_state, base_eval_loader, base_eval_meta)
+    print(f"        Accuracy metrics: {base_acc}")
+    base_lat = measure_latency(base_state.model, base_lat_loader, device,
+                               iterations=args.latency_iters)
+    print(f"        Latency: {base_lat['latency_ms']:.2f} ms ± {base_lat['latency_std']:.2f}")
+    results["baseline"] = {"accuracy": base_acc, "latency": base_lat}
+
+    # Free memory
+    del base_state, base_eval_loader, base_lat_loader
+    torch.cuda.empty_cache()
+
+    # ── Tiered ──
+    print("\n  [3/4] Loading Tiered checkpoint...")
+    tier_state, tier_cfg, tier_eval_loader, tier_lat_loader, tier_eval_meta = load_trained_model(args.tiered, device)
+    n_params = sum(p.numel() for p in tier_state.model.parameters()) / 1e6
+    print(f"        Step: {tier_state.step}, Params: {n_params:.1f}M")
+
+    print("  [4/4] Evaluating Tiered accuracy + latency...")
+    tier_acc = eval_accuracy(tier_cfg, tier_state, tier_eval_loader, tier_eval_meta)
+    print(f"        Accuracy metrics: {tier_acc}")
+    tier_lat = measure_latency(tier_state.model, tier_lat_loader, device,
+                               iterations=args.latency_iters)
+    print(f"        Latency: {tier_lat['latency_ms']:.2f} ms ± {tier_lat['latency_std']:.2f}")
+    results["tiered"] = {"accuracy": tier_acc, "latency": tier_lat}
+
+    del tier_state, tier_eval_loader, tier_lat_loader
+    torch.cuda.empty_cache()
+
+    # ── Plots ──
+    print("\n  Generating comparison plots...")
+    create_combined_plot(results["baseline"], results["tiered"], args.output_dir)
+
+    # ── Save JSON ──
+    json_path = os.path.join(args.output_dir, "trained_model_results.json")
+    os.makedirs(args.output_dir, exist_ok=True)
+    with open(json_path, "w") as f:
+        json.dump(results, f, indent=2, default=str)
+    print(f"  Results saved → {json_path}")
+
+    print("\n" + "=" * 64)
+    print("  Done!")
+    print("=" * 64)
+
+
+if __name__ == "__main__":
+    main()

models/common.py

@@ -0,0 +1,32 @@
+import math
+
+import torch
+from torch import nn
+
+
+def trunc_normal_init_(tensor: torch.Tensor, std: float = 1.0, lower: float = -2.0, upper: float = 2.0):
+    # NOTE: PyTorch nn.init.trunc_normal_ is not mathematically correct, the std dev is not actually the std dev of initialized tensor
+    # This function is a PyTorch version of jax truncated normal init (default init method in flax)
+    # https://github.com/jax-ml/jax/blob/main/jax/_src/random.py#L807-L848
+    # https://github.com/jax-ml/jax/blob/main/jax/_src/nn/initializers.py#L162-L199
+
+    with torch.no_grad():
+        if std == 0:
+            tensor.zero_()
+        else:
+            sqrt2 = math.sqrt(2)
+            a = math.erf(lower / sqrt2)
+            b = math.erf(upper / sqrt2)
+            z = (b - a) / 2
+
+            c = (2 * math.pi) ** -0.5
+            pdf_u = c * math.exp(-0.5 * lower ** 2)
+            pdf_l = c * math.exp(-0.5 * upper ** 2)
+            comp_std = std / math.sqrt(1 - (upper * pdf_u - lower * pdf_l) / z - ((pdf_u - pdf_l) / z) ** 2)
+
+            tensor.uniform_(a, b)
+            tensor.erfinv_()
+            tensor.mul_(sqrt2 * comp_std)
+            tensor.clip_(lower * comp_std, upper * comp_std)
+
+    return tensor

models/fused_hierarchical_scan.py

@@ -0,0 +1,186 @@
+"""
+Fused Hierarchical Scan for the Tiered Memory HRM.
+
+This module implements the algorithm proposed in Section 4.2 of the project proposal:
+'The Tiered Memory Hierarchical Scan'.
+
+It replaces the Python-level `for` loop over L_cycles with a single OpenAI Triton
+kernel invocation. The kernel loads the slow state (z_H) once into SRAM, executes
+the entire fast state (z_L) recurrence for T steps completely within SRAM (registers/shared memory),
+and only writes the final states back to HBM (DRAM) at the cycle boundary.
+
+Algorithm Steps Implemented:
+1. Block-wise Loading: Sequence is chunked by T (L_cycles).
+2. DRAM -> SRAM Fetch: z_H and inputs loaded once.
+3. SRAM-Resident Recurrence: T steps computed without global memory access.
+4. Boundary Sync: High-level state updated at the end.
+5. SRAM -> DRAM Write: Only final states are materialized.
+"""
+
+import torch
+import triton
+import triton.language as tl
+import math
+
+
+@triton.jit
+def _fused_hierarchical_scan_kernel(
+    # --- Data Pointers ---
+    X_ptr,          # Inputs sequence [batch, seq_len, hidden_size]
+    Z_L_in_ptr,     # Initial Fast State (L-level) [batch, hidden_size]
+    Z_H_in_ptr,     # Initial Slow State (H-level) [batch, hidden_size]
+    Z_L_out_ptr,    # Final Fast State Output [batch, hidden_size]
+    Z_H_out_ptr,    # Final Slow State Output [batch, hidden_size]
+    # --- Matrix Weights (Simplified Recurrence for demonstration) ---
+    W_L_ptr,        # Weights for L-level update [hidden_size, hidden_size]
+    W_H_ptr,        # Weights for H-level update [hidden_size, hidden_size]
+    # --- Shapes & Strides ---
+    stride_batch_x, stride_seq_x, stride_dim_x,
+    T: tl.constexpr,               # Number of fast steps per slow step (L_cycles)
+    HIDDEN_SIZE: tl.constexpr,     # Hidden dimension (must fit in SRAM)
+    BLOCK_DIM: tl.constexpr,       # Power of 2 for memory alignment
+):
+    """
+    SRAM-Resident Hierarchical Scan Kernel.
+    
+    This kernel is designed so that the inner loop (t = 0...T) operates entirely
+    on registers (`z_L` and `z_H` variables inside the kernel). 
+    """
+    batch_idx = tl.program_id(0)
+    
+    # Offsets for the hidden dimension
+    dim_offsets = tl.arange(0, BLOCK_DIM)
+    mask = dim_offsets < HIDDEN_SIZE
+
+    # -----------------------------------------------------------------
+    # Step 2: DRAM -> SRAM Fetch
+    # Load the initial states into registers (SRAM) for this batch
+    # -----------------------------------------------------------------
+    z_L = tl.load(Z_L_in_ptr + batch_idx * HIDDEN_SIZE + dim_offsets, mask=mask, other=0.0)
+    z_H = tl.load(Z_H_in_ptr + batch_idx * HIDDEN_SIZE + dim_offsets, mask=mask, other=0.0)
+
+    # Note: In a full transformer, W_L would be large. For a true SRAM-resident kernel,
+    # the parameter matrices must either be small enough to stay in shared memory,
+    # or the recurrence is element-wise (like Mamba/SSMs). 
+    # Here we simulate the update using an element-wise/diagonal approximation 
+    # of the weights to ensure it stays in SRAM, mimicking a Diagonal State Space Model.
+    w_L = tl.load(W_L_ptr + dim_offsets, mask=mask, other=0.0)
+    w_H = tl.load(W_H_ptr + dim_offsets, mask=mask, other=0.0)
+
+    # -----------------------------------------------------------------
+    # Step 3: SRAM-Resident Recurrence
+    # Execute the L-level loop entirely in SRAM without hitting DRAM
+    # -----------------------------------------------------------------
+    for t in range(T):
+        # 3a. Load input `x_t` for the current step (DRAM -> SRAM)
+        x_t_ptr = X_ptr + batch_idx * stride_batch_x + t * stride_seq_x + dim_offsets * stride_dim_x
+        x_t = tl.load(x_t_ptr, mask=mask, other=0.0)
+
+        # 3b. L-Level Update Rule: f_L(z_L, z_H, x)
+        # E.g., z_L = act(W_L * z_L + z_H + x_t)
+        # All computation here is Register-to-Register (Zero HBM cost)
+        pre_act = (w_L * z_L) + z_H + x_t
+        
+        # Simple non-linearity (e.g., SiLU/Swish)
+        z_L = pre_act * tl.sigmoid(pre_act)
+
+    # -----------------------------------------------------------------
+    # Step 4: Boundary Sync
+    # Compute the new high-level state using the final z_L
+    # -----------------------------------------------------------------
+    # f_H(z_H, z_L)
+    pre_act_H = (w_H * z_H) + z_L
+    z_H_new = pre_act_H * tl.sigmoid(pre_act_H)
+
+    # -----------------------------------------------------------------
+    # Step 5: SRAM -> DRAM Write
+    # Write only the final chunk boundaries back to Global Memory (HBM)
+    # -----------------------------------------------------------------
+    tl.store(Z_L_out_ptr + batch_idx * HIDDEN_SIZE + dim_offsets, z_L, mask=mask)
+    tl.store(Z_H_out_ptr + batch_idx * HIDDEN_SIZE + dim_offsets, z_H_new, mask=mask)
+
+
+# =====================================================================
+# PyTorch Wrapper
+# =====================================================================
+
+def next_power_of_2(n: int) -> int:
+    """Returns the next power of 2 greater than or equal to n."""
+    return 1 if n == 0 else 2**(n - 1).bit_length()
+
+class FusedHierarchicalScanLayer(torch.nn.Module):
+    """
+    PyTorch module that encapsulates the Tiered Memory Hierarchical Scan.
+    
+    This replaces the Python loop over T steps with the Triton kernel,
+    achieving the memory I/O reduction outlined in the paper.
+    """
+    def __init__(self, hidden_size: int, T_steps: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.T = T_steps
+        
+        # Diagonal weight matrices for the recurrence (SSM-style)
+        self.w_L = torch.nn.Parameter(torch.randn(hidden_size) / math.sqrt(hidden_size))
+        self.w_H = torch.nn.Parameter(torch.randn(hidden_size) / math.sqrt(hidden_size))
+
+    def forward(self, x_chunk: torch.Tensor, z_L: torch.Tensor, z_H: torch.Tensor):
+        """
+        Args:
+            x_chunk: Tensor of shape (batch, T, hidden_size) containing the inputs for this chunk.
+            z_L: Tensor of shape (batch, hidden_size) containing the initial L-state.
+            z_H: Tensor of shape (batch, hidden_size) containing the initial H-state.
+            
+        Returns:
+            z_L_new, z_H_new: The updated states after T steps.
+        """
+        batch_size, seq_len, dim = x_chunk.shape
+        
+        assert seq_len == self.T, f"Expected chunk of size T={self.T}, got {seq_len}"
+        assert dim == self.hidden_size, "Dimension mismatch"
+        assert z_L.shape == (batch_size, self.hidden_size)
+        assert z_H.shape == (batch_size, self.hidden_size)
+        assert x_chunk.is_contiguous()
+
+        # Allocate output tensors in HBM
+        z_L_out = torch.empty_like(z_L)
+        z_H_out = torch.empty_like(z_H)
+
+        # Determine Triton block size
+        BLOCK_DIM = next_power_of_2(self.hidden_size)
+
+        # 1D Grid: one program per batch element
+        grid = (batch_size,)
+
+        # Launch the fused kernel
+        _fused_hierarchical_scan_kernel[grid](
+            x_chunk, z_L, z_H, z_L_out, z_H_out,
+            self.w_L, self.w_H,
+            x_chunk.stride(0), x_chunk.stride(1), x_chunk.stride(2),
+            T=self.T,
+            HIDDEN_SIZE=self.hidden_size,
+            BLOCK_DIM=BLOCK_DIM
+        )
+
+        return z_L_out, z_H_out
+
+# Example conceptual usage:
+if __name__ == "__main__":
+    batch = 32
+    T = 8 # Number of L-cycles in one H-cycle
+    dim = 256
+    
+    # Initialize the fused scan module
+    scanner = FusedHierarchicalScanLayer(hidden_size=dim, T_steps=T).cuda()
+    
+    # Dummy data
+    x_chunk = torch.randn(batch, T, dim, device='cuda')
+    z_L_init = torch.randn(batch, dim, device='cuda')
+    z_H_init = torch.randn(batch, dim, device='cuda')
+    
+    # Execute the fused scan!
+    z_L_final, z_H_final = scanner(x_chunk, z_L_init, z_H_init)
+    
+    print(f"Successfully executed Tiered Memory Hierarchical Scan.")
+    print(f"z_L moved from {z_L_init.shape} -> {z_L_final.shape}")
+    print(f"z_H moved from {z_H_init.shape} -> {z_H_final.shape}")

models/hrm/hrm_act_v1.py

@@ -0,0 +1,283 @@
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from pydantic import BaseModel
+
+from models.common import trunc_normal_init_
+from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear
+from models.sparse_embedding import CastedSparseEmbedding
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1InnerCarry:
+    z_H: torch.Tensor
+    z_L: torch.Tensor
+
+
+@dataclass
+class HierarchicalReasoningModel_ACTV1Carry:
+    inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry
+    
+    steps: torch.Tensor
+    halted: torch.Tensor
+    
+    current_data: Dict[str, torch.Tensor]
+
+
+class HierarchicalReasoningModel_ACTV1Config(BaseModel):
+    batch_size: int
+    seq_len: int
+    puzzle_emb_ndim: int = 0
+    num_puzzle_identifiers: int
+    vocab_size: int
+
+    H_cycles: int
+    L_cycles: int
+
+    H_layers: int
+    L_layers: int
+
+    # Transformer config
+    hidden_size: int
+    expansion: float
+    num_heads: int
+    pos_encodings: str
+
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    
+    # Halting Q-learning config
+    halt_max_steps: int
+    halt_exploration_prob: float
+
+    forward_dtype: str = "bfloat16"
+
+
+class HierarchicalReasoningModel_ACTV1Block(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+
+        self.self_attn = Attention(
+            hidden_size=config.hidden_size,
+            head_dim=config.hidden_size // config.num_heads,
+            num_heads=config.num_heads,
+            num_key_value_heads=config.num_heads,
+            causal=False
+        )
+        self.mlp = SwiGLU(
+            hidden_size=config.hidden_size,
+            expansion=config.expansion,
+        )
+        self.norm_eps = config.rms_norm_eps
+
+    def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Post Norm
+        # Self Attention
+        hidden_states = rms_norm(hidden_states + self.self_attn(cos_sin=cos_sin, hidden_states=hidden_states), variance_epsilon=self.norm_eps)
+        # Fully Connected
+        hidden_states = rms_norm(hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps)
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module):
+    def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]):
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Input injection (add)
+        hidden_states = hidden_states + input_injection
+        # Layers
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+
+        return hidden_states
+
+
+class HierarchicalReasoningModel_ACTV1_Inner(nn.Module):
+    def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None:
+        super().__init__()
+        self.config = config
+        self.forward_dtype = getattr(torch, self.config.forward_dtype)
+
+        # I/O
+        self.embed_scale  = math.sqrt(self.config.hidden_size)
+        embed_init_std = 1.0 / self.embed_scale
+
+        self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        self.lm_head      = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.q_head       = CastedLinear(self.config.hidden_size, 2, bias=True)
+
+        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  # ceil div
+        if self.config.puzzle_emb_ndim > 0:
+            # Zero init puzzle embeddings
+            self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
+                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)
+
+        # LM Blocks
+        if self.config.pos_encodings == "rope":
+            self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads,
+                                              max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
+                                              base=self.config.rope_theta)
+        elif self.config.pos_encodings == "learned":
+            self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
+        else:
+            raise NotImplementedError()
+
+        # Reasoning Layers
+        self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)])
+        self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)])
+        
+        # Initial states
+        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
+
+        # Q head special init
+        # Init Q to (almost) zero for faster learning during bootstrapping
+        with torch.no_grad():
+            self.q_head.weight.zero_()
+            self.q_head.bias.fill_(-5)  # type: ignore
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor):
+        # Token embedding
+        embedding = self.embed_tokens(input.to(torch.int32))
+
+        # Puzzle embeddings
+        if self.config.puzzle_emb_ndim > 0:
+            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
+            
+            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
+            if pad_count > 0:
+                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))
+
+            embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2)
+
+        # Position embeddings
+        if self.config.pos_encodings == "learned":
+            # scale by 1/sqrt(2) to maintain forward variance
+            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))
+
+        # Scale
+        return self.embed_scale * embedding
+
+    def empty_carry(self, batch_size: int):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+            z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
+        )
+        
+    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
+            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
+        )
+
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        seq_info = dict(
+            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
+        )
+
+        # Input encoding
+        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])
+
+        # Forward iterations
+        with torch.no_grad():
+            z_H, z_L = carry.z_H, carry.z_L
+
+            for _H_step in range(self.config.H_cycles):
+                for _L_step in range(self.config.L_cycles):
+                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
+                        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+
+                if not (_H_step == self.config.H_cycles - 1):
+                    z_H = self.H_level(z_H, z_L, **seq_info)
+
+        assert not z_H.requires_grad and not z_L.requires_grad
+
+        # 1-step grad
+        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.H_level(z_H, z_L, **seq_info)
+
+        # LM Outputs
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+
+        # Q head
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+        
+        return new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+
+
+class HierarchicalReasoningModel_ACTV1(nn.Module):
+    """ACT wrapper."""
+
+    def __init__(self, config_dict: dict):
+        super().__init__()
+        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
+        self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+
+        return HierarchicalReasoningModel_ACTV1Carry(
+            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
+            
+            steps=torch.zeros((batch_size, ), dtype=torch.int32),
+            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
+            
+            current_data={k: torch.empty_like(v) for k, v in batch.items()}
+        )
+        
+    def forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+        # Update data, carry (removing halted sequences)
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+
+        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}
+
+        # Forward inner model
+        new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits
+        }
+        
+        with torch.no_grad():
+            # Step
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            
+            halted = is_last_step
+
+            # if training, and ACT is enabled
+            if self.training and (self.config.halt_max_steps > 1):
+                # Halt signal
+                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
+                halted = halted | (q_halt_logits > q_continue_logits)
+
+                # Exploration
+                min_halt_steps = (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
+
+                halted = halted & (new_steps >= min_halt_steps)
+
+                # Compute target Q
+                # NOTE: No replay buffer and target networks for computing target Q-value.
+                # As batch_size is large, there're many parallel envs.
+                # Similar concept as PQN https://arxiv.org/abs/2407.04811
+                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-1]
+                
+                outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits)))
+
+        return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs

models/hrm/hrm_tiered.py

@@ -0,0 +1,466 @@
+"""
+HRM with SRAM/DRAM Memory Tiering — Triton-accelerated.
+
+This module wraps the original HRM architecture and adds explicit
+memory tier placement:
+  - L-level state (z_L): SRAM tier (Triton kernels keep in registers/shared mem)
+  - H-level state (z_H): DRAM tier (standard global memory path)
+
+The Triton kernels in `triton_kernels.py` handle the fused operations,
+and `MemoryTierManager` tracks all memory events for benchmarking.
+
+EXPLANATION FOR BEGINNERS:
+Think of this model like a Human brain solving a hard puzzle:
+- SRAM (L-level): Your 'Short-term memory' or Scratchpad. It's super fast but tiny.
+- DRAM (H-level): Your 'Long-term planning'. It's slower but holds the big picture.
+
+The 'Triton' kernels are special pieces of code that talk directly to the GPU 
+hardware to make these two memory 'tiers' run as fast as possible.
+"""
+
+
+from typing import Tuple, List, Dict, Optional
+from dataclasses import dataclass
+import math
+import time
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from models.common import trunc_normal_init_
+from models.layers import (
+    rms_norm, SwiGLU, Attention, RotaryEmbedding,
+    CosSin, CastedEmbedding, CastedLinear,
+)
+from models.sparse_embedding import CastedSparseEmbedding
+from models.memory_tier import MemoryTierManager
+from models.triton_kernels import (
+    triton_rms_norm_residual_sram,
+    triton_rms_norm_residual_dram,
+    triton_swiglu_sram,
+    triton_state_transfer,
+)
+
+# Try importing the original HRM components for carry/config reuse
+from models.hrm.hrm_act_v1 import (
+    HierarchicalReasoningModel_ACTV1Config,
+    HierarchicalReasoningModel_ACTV1InnerCarry,
+    HierarchicalReasoningModel_ACTV1Carry,
+)
+
+
+# ===================================================================
+# Tiered Transformer Block — uses Triton kernels per tier
+# ===================================================================
+
+class TieredBlock(nn.Module):
+    """A transformer block that uses tier-specific Triton kernels.
+
+    A single 'Building Block' of the transformer.
+    It can live in either the 'fast' (SRAM) or 'slow' (DRAM) tier.
+    """
+
+    def __init__(
+        self,
+        config: HierarchicalReasoningModel_ACTV1Config,
+        tier: str,  # 'sram' (fast) or 'dram' (slow)
+    ):
+        super().__init__()
+        self.tier = tier
+        self.norm_eps = config.rms_norm_eps
+
+        # Standard transformer components: Attention (looking at other cells) 
+        # and MLP (thinking about the current cell)
+        self.self_attn = Attention(
+            hidden_size=config.hidden_size,
+            head_dim=config.hidden_size // config.num_heads,
+            num_heads=config.num_heads,
+            num_key_value_heads=config.num_heads,
+            causal=False,
+        )
+        self.mlp = SwiGLU(
+            hidden_size=config.hidden_size,
+            expansion=config.expansion,
+        )
+
+        # Track whether we can use Triton (requires CUDA + contiguous bf16)
+        self._use_triton = torch.cuda.is_available()
+
+    def forward(
+        self,
+        cos_sin: CosSin,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        """Processes the state through one layer of math."""
+        # 1. Attention: Which other cells in the puzzle are relevant right now?
+        attn_out = self.self_attn(cos_sin=cos_sin, hidden_states=hidden_states)
+
+        # 2. Use specialized Triton 'kernels' to speed up the memory access
+        if self._use_triton and hidden_states.is_contiguous():
+            if self.tier == 'sram':
+                hidden_states = triton_rms_norm_residual_sram(
+                    attn_out, hidden_states, eps=self.norm_eps,
+                )
+            else:
+                hidden_states = triton_rms_norm_residual_dram(
+                    attn_out, hidden_states, eps=self.norm_eps,
+                )
+        else:
+            # Fallback to PyTorch
+            hidden_states = rms_norm(
+                hidden_states + attn_out,
+                variance_epsilon=self.norm_eps,
+            )
+
+        # 3. MLP: Update the 'knowledge' of each cell based on the new info
+        mlp_out = self.mlp(hidden_states)
+
+        if self._use_triton and hidden_states.is_contiguous():
+            if self.tier == 'sram':
+                hidden_states = triton_rms_norm_residual_sram(
+                    mlp_out, hidden_states, eps=self.norm_eps,
+                )
+            else:
+                hidden_states = triton_rms_norm_residual_dram(
+                    mlp_out, hidden_states, eps=self.norm_eps,
+                )
+        else:
+            hidden_states = rms_norm(
+                hidden_states + mlp_out,
+                variance_epsilon=self.norm_eps,
+            )
+
+        return hidden_states
+
+
+class TieredReasoningModule(nn.Module):
+    """
+    A collection of layers that performs one 'thought' step.
+    This corresponds to the H-level or L-level of the model.
+    """
+
+    def __init__(
+        self,
+        layers: List[TieredBlock],
+        tier: str,
+    ):
+        super().__init__()
+        self.tier = tier
+        self.layers = nn.ModuleList(layers)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_injection: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + input_injection
+
+        for layer in self.layers:
+            hidden_states = layer(hidden_states=hidden_states, **kwargs)
+
+        return hidden_states
+
+
+# ===================================================================
+# Main Tiered HRM Model
+# ===================================================================
+
+class HRM_Tiered_Inner(nn.Module):
+    """HRM inner model with SRAM/DRAM memory tiering.
+
+    - L-level: SRAM tier (Triton fused kernels, register-resident)
+    - H-level: DRAM tier (Triton DRAM-path kernels)
+    """
+
+    def __init__(
+        self,
+        config: HierarchicalReasoningModel_ACTV1Config,
+        memory_manager: Optional[MemoryTierManager] = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.forward_dtype = getattr(torch, self.config.forward_dtype)
+        self.memory_manager = memory_manager
+
+        # --- Embedding: Converting digits (0-9) into high-dimensional vectors ---
+        self.embed_scale = math.sqrt(self.config.hidden_size)
+        embed_init_std = 1.0 / self.embed_scale
+
+        self.embed_tokens = CastedEmbedding(
+            self.config.vocab_size, self.config.hidden_size,
+            init_std=embed_init_std, cast_to=self.forward_dtype,
+        )
+        self.lm_head = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.q_head = CastedLinear(self.config.hidden_size, 2, bias=True)
+
+        # --- Puzzle Embedding: Special memory for each specific puzzle ---
+        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)
+        if self.config.puzzle_emb_ndim > 0:
+            self.puzzle_emb = CastedSparseEmbedding(
+                self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
+                batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype,
+            )
+
+        # ---- Position encodings ----
+        if self.config.pos_encodings == "rope":
+            self.rotary_emb = RotaryEmbedding(
+                dim=self.config.hidden_size // self.config.num_heads,
+                max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
+                base=self.config.rope_theta,
+            )
+        elif self.config.pos_encodings == "learned":
+            self.embed_pos = CastedEmbedding(
+                self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size,
+                init_std=embed_init_std, cast_to=self.forward_dtype,
+            )
+        else:
+            raise NotImplementedError()
+
+        # ---- Tiered Reasoning Layers ----
+        # H-level: DRAM tier (slow updates)
+        self.H_level = TieredReasoningModule(
+            layers=[TieredBlock(self.config, tier='dram') for _ in range(self.config.H_layers)],
+            tier='dram',
+        )
+        # L-level: SRAM tier (fast updates)
+        self.L_level = TieredReasoningModule(
+            layers=[TieredBlock(self.config, tier='sram') for _ in range(self.config.L_layers)],
+            tier='sram',
+        )
+
+        # ---- Initial states ----
+        self.H_init = nn.Buffer(
+            trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1),
+            persistent=True,
+        )
+        self.L_init = nn.Buffer(
+            trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1),
+            persistent=True,
+        )
+
+        # Q head special init
+        with torch.no_grad():
+            self.q_head.weight.zero_()
+            self.q_head.bias.fill_(-5)
+
+        # ---- Timing storage for benchmarks ----
+        self._timing: Dict[str, List[float]] = {
+            'L_forward_us': [],
+            'H_forward_us': [],
+            'H_L_transfer_us': [],
+            'L_H_transfer_us': [],
+        }
+
+    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor):
+        embedding = self.embed_tokens(input.to(torch.int32))
+
+        if self.config.puzzle_emb_ndim > 0:
+            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
+            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
+            if pad_count > 0:
+                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))
+            embedding = torch.cat(
+                (puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding),
+                dim=-2,
+            )
+
+        if self.config.pos_encodings == "learned":
+            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))
+
+        return self.embed_scale * embedding
+
+    def empty_carry(self, batch_size: int):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.empty(
+                batch_size, self.config.seq_len + self.puzzle_emb_len,
+                self.config.hidden_size, dtype=self.forward_dtype,
+            ),
+            z_L=torch.empty(
+                batch_size, self.config.seq_len + self.puzzle_emb_len,
+                self.config.hidden_size, dtype=self.forward_dtype,
+            ),
+        )
+
+    def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry):
+        return HierarchicalReasoningModel_ACTV1InnerCarry(
+            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
+            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
+        )
+
+    @torch.compiler.disable
+    def _gpu_timer_start(self):
+        if torch.cuda.is_available():
+            start = torch.cuda.Event(enable_timing=True)
+            start.record()
+            return start
+        return time.perf_counter()
+
+    @torch.compiler.disable
+    def _gpu_timer_end(self, start) -> float:
+        if isinstance(start, torch.cuda.Event):
+            end = torch.cuda.Event(enable_timing=True)
+            end.record()
+            torch.cuda.synchronize()
+            return start.elapsed_time(end) * 1000  # ms → μs
+        return (time.perf_counter() - start) * 1e6
+
+    @torch.compiler.disable
+    def _timed_forward(self, module, hidden_states, injection, seq_info, tier_key, memory_ctx=None):
+        """Run a reasoning module forward with optional timing (disabled under torch.compile)."""
+        t0 = self._gpu_timer_start() if not self.training else None
+
+        if memory_ctx is not None:
+            with memory_ctx:
+                out = module(hidden_states, injection, **seq_info)
+        else:
+            out = module(hidden_states, injection, **seq_info)
+
+        if t0 is not None:
+            self._timing[tier_key].append(self._gpu_timer_end(t0))
+        return out
+
+    def forward(
+        self,
+        carry: HierarchicalReasoningModel_ACTV1InnerCarry,
+        batch: Dict[str, torch.Tensor],
+    ) -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+
+        seq_info = dict(
+            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None,
+        )
+
+        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])
+
+        # ---- Hierarchical reasoning with tiered memory ----
+        with torch.no_grad():
+            z_H, z_L = carry.z_H, carry.z_L
+
+            for _H_step in range(self.config.H_cycles):
+                for _L_step in range(self.config.L_cycles):
+                    if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)):
+                        # L-level: SRAM tier (fast path)
+                        sram_ctx = self.memory_manager.sram_context() if self.memory_manager else None
+                        z_L = self._timed_forward(
+                            self.L_level, z_L, z_H + input_embeddings,
+                            seq_info, 'L_forward_us', sram_ctx,
+                        )
+
+                if not (_H_step == self.config.H_cycles - 1):
+                    # H-level: DRAM tier (slow path)
+                    z_L_for_H = z_L.contiguous()
+                    dram_ctx = self.memory_manager.dram_context() if self.memory_manager else None
+                    z_H = self._timed_forward(
+                        self.H_level, z_H, z_L_for_H,
+                        seq_info, 'H_forward_us', dram_ctx,
+                    )
+
+        assert not z_H.requires_grad and not z_L.requires_grad
+
+        # 1-step grad (with tier-aware kernels)
+        z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
+        z_H = self.H_level(z_H, z_L, **seq_info)
+
+        new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())
+        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
+        q_logits = self.q_head(z_H[:, 0]).to(torch.float32)
+
+        return new_carry, output, (q_logits[..., 0], q_logits[..., 1])
+
+    def get_timing_stats(self) -> Dict[str, Dict[str, float]]:
+        """Get aggregated timing statistics for benchmarking."""
+        stats = {}
+        for key, values in self._timing.items():
+            if values:
+                stats[key] = {
+                    'mean_us': sum(values) / len(values),
+                    'min_us': min(values),
+                    'max_us': max(values),
+                    'count': len(values),
+                    'total_us': sum(values),
+                }
+            else:
+                stats[key] = {'mean_us': 0, 'min_us': 0, 'max_us': 0, 'count': 0, 'total_us': 0}
+        return stats
+
+    def reset_timing(self):
+        for key in self._timing:
+            self._timing[key] = []
+
+
+# ===================================================================
+# ACT Wrapper (same logic as original, uses tiered inner model)
+# ===================================================================
+
+class HRM_Tiered(nn.Module):
+    """ACT wrapper for the memory-tiered HRM model."""
+
+    def __init__(self, config_dict: dict, memory_manager: Optional[MemoryTierManager] = None):
+        super().__init__()
+        self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict)
+        self.memory_manager = memory_manager
+        self.inner = HRM_Tiered_Inner(self.config, memory_manager=memory_manager)
+
+    @property
+    def puzzle_emb(self):
+        return self.inner.puzzle_emb
+
+    def initial_carry(self, batch: Dict[str, torch.Tensor]):
+        batch_size = batch["inputs"].shape[0]
+        return HierarchicalReasoningModel_ACTV1Carry(
+            inner_carry=self.inner.empty_carry(batch_size),
+            steps=torch.zeros((batch_size,), dtype=torch.int32),
+            halted=torch.ones((batch_size,), dtype=torch.bool),
+            current_data={k: torch.empty_like(v) for k, v in batch.items()},
+        )
+
+    def forward(
+        self,
+        carry: HierarchicalReasoningModel_ACTV1Carry,
+        batch: Dict[str, torch.Tensor],
+    ) -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:
+        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
+        new_steps = torch.where(carry.halted, 0, carry.steps)
+        new_current_data = {
+            k: torch.where(carry.halted.view((-1,) + (1,) * (batch[k].ndim - 1)), batch[k], v)
+            for k, v in carry.current_data.items()
+        }
+
+        new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)
+
+        outputs = {
+            "logits": logits,
+            "q_halt_logits": q_halt_logits,
+            "q_continue_logits": q_continue_logits,
+        }
+
+        with torch.no_grad():
+            new_steps = new_steps + 1
+            is_last_step = new_steps >= self.config.halt_max_steps
+            halted = is_last_step
+
+            if self.training and (self.config.halt_max_steps > 1):
+                halted = halted | (q_halt_logits > q_continue_logits)
+                min_halt_steps = (
+                    (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob)
+                    * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
+                )
+                halted = halted & (new_steps >= min_halt_steps)
+
+                next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-1]
+                outputs["target_q_continue"] = torch.sigmoid(
+                    torch.where(
+                        is_last_step,
+                        next_q_halt_logits,
+                        torch.maximum(next_q_halt_logits, next_q_continue_logits),
+                    )
+                )
+
+        return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs
+
+    def get_timing_stats(self) -> Dict:
+        return self.inner.get_timing_stats()
+
+    def reset_timing(self):
+        self.inner.reset_timing()

models/layers.py

@@ -0,0 +1,174 @@
+from typing import Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+_USE_NATIVE_ATTN = False
+
+try:
+    from flash_attn_interface import flash_attn_func  # type: ignore[import]
+except ImportError:
+    try:
+        from flash_attn import flash_attn_func  # type: ignore[import]
+    except ImportError:
+        # Fallback to PyTorch native scaled dot-product attention
+        _USE_NATIVE_ATTN = True
+
+        def flash_attn_func(q, k, v, causal=False):
+            """PyTorch-native fallback for FlashAttention.
+
+            Input shape:  [batch, seq_len, num_heads, head_dim]
+            SDPA expects:  [batch, num_heads, seq_len, head_dim]
+            """
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
+            return out.transpose(1, 2)
+
+from models.common import trunc_normal_init_
+
+
+CosSin = Tuple[torch.Tensor, torch.Tensor]
+
+
+def _find_multiple(a, b):
+    return (-(a // -b)) * b
+
+
+def rotate_half(x: torch.Tensor):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    # q, k: [bs, seq_len, num_heads, head_dim]
+    # cos, sin: [seq_len, head_dim]
+    orig_dtype = q.dtype
+    q = q.to(cos.dtype)
+    k = k.to(cos.dtype)
+
+    q_embed = (q * cos.unsqueeze(-2)) + (rotate_half(q) * sin.unsqueeze(-2))
+    k_embed = (k * cos.unsqueeze(-2)) + (rotate_half(k) * sin.unsqueeze(-2))
+
+    return q_embed.to(orig_dtype), k_embed.to(orig_dtype)
+
+
+class CastedLinear(nn.Module):
+    def __init__(self,
+                 in_features: int,
+                 out_features: int,
+                 bias: bool):
+        super().__init__()
+        # Truncated LeCun normal init
+        self.weight = nn.Parameter(
+            trunc_normal_init_(torch.empty((out_features, in_features)), std=1.0 / (in_features ** 0.5))
+        )
+        self.bias = None
+        if bias:
+            # Zero init bias
+            self.bias = nn.Parameter(torch.zeros((out_features, )))
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.linear(input, self.weight.to(input.dtype), bias=self.bias.to(input.dtype) if self.bias is not None else None)
+
+
+class CastedEmbedding(nn.Module):
+    def __init__(self,
+                 num_embeddings: int,
+                 embedding_dim: int,
+                 init_std: float,
+                 cast_to: torch.dtype):
+        super().__init__()
+        self.cast_to = cast_to
+
+        # Truncated LeCun normal init
+        self.embedding_weight = nn.Parameter(
+            trunc_normal_init_(torch.empty((num_embeddings, embedding_dim)), std=init_std)
+        )
+        
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.embedding(input, self.embedding_weight.to(self.cast_to))
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings, base, device=None):
+        super().__init__()
+
+        # RoPE
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+        t = torch.arange(max_position_embeddings, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = nn.Buffer(emb.cos(), persistent=False)
+        self.sin_cached = nn.Buffer(emb.sin(), persistent=False)
+
+    def forward(self):
+        return self.cos_cached, self.sin_cached
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_size, head_dim, num_heads, num_key_value_heads, causal=False):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+        self.output_size = head_dim * num_heads
+        self.num_heads = num_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.causal = causal
+
+        self.qkv_proj = CastedLinear(self.hidden_size, (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, bias=False)
+        self.o_proj = CastedLinear(self.output_size, self.hidden_size, bias=False)
+
+    def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # hidden_states: [bs, seq_len, num_heads, head_dim]
+        qkv = self.qkv_proj(hidden_states)
+
+        # Split head
+        qkv = qkv.view(batch_size, seq_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+        query = qkv[:, :, :self.num_heads]
+        key = qkv[:, :, self.num_heads: self.num_heads + self.num_key_value_heads]
+        value = qkv[:, :, self.num_heads + self.num_key_value_heads:]
+
+        # RoPE
+        if cos_sin is not None:
+            cos, sin = cos_sin
+            query, key = apply_rotary_pos_emb(query, key, cos, sin)
+
+        # flash attn
+        attn_output = flash_attn_func(q=query, k=key, v=value, causal=self.causal)
+        if isinstance(attn_output, tuple):  # fa2 and fa3 compatibility
+            attn_output = attn_output[0]
+
+        attn_output = attn_output.view(batch_size, seq_len, self.output_size)  # type: ignore
+        return self.o_proj(attn_output)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_size: int, expansion: float):
+        super().__init__()
+        inter = _find_multiple(round(expansion * hidden_size * 2 / 3), 256)
+
+        self.gate_up_proj = CastedLinear(hidden_size, inter * 2, bias=False)
+        self.down_proj    = CastedLinear(inter, hidden_size, bias=False)
+
+    def forward(self, x):
+        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+def rms_norm(hidden_states: torch.Tensor, variance_epsilon: float) -> torch.Tensor:
+    input_dtype = hidden_states.dtype
+    hidden_states = hidden_states.to(torch.float32)
+
+    variance = hidden_states.square().mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
+    return hidden_states.to(input_dtype)