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