models/fused_hierarchical_scan.py

"""
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}")