docs/IMPLEMENTATION.md

**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.