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
• Triton ref: softmax_triton.py
• Memory tier manager: test-env/HRM_optimised/models/memory_tier.py
• Triton kernels: test-env/HRM_optimised/models/triton_kernels.py
• Sparse embedding + optimizer: 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.