Enable Partial Tile Store/Loading

# Feature Request: Support Partial Tile Store for LCE + Slice Fusion ## Summary When fusing a slice operation into a Helion tiled matmul kernel, the current approach can only write **complete tiles** that fall entirely within the slice boundary. If `slice_size` is not a multiple of the tile size, the boundary tile is either dropped or incorrectly handled. We need **partial tile store** support in Helion to properly fuse slice operations into tiled kernels. ## Motivation In production models (e.g., LCE – Linear Combination Embedding), a common pattern is: ``` output = A^T @ B^T + bias # matmul sliced = output[:, :slice_size, :] # followed immediately by a slice ``` Fusing the slice into the matmul kernel avoids materializing the full output tensor and reading it back just to slice — saving both memory bandwidth and a kernel launch. However, because the kernel tiles over the N dimension in fixed-size blocks, the slice boundary rarely aligns with tile boundaries, making correct fusion impossible without partial tile store support. ## Simplified Example Below is a simplified version of the kernel ### PyTorch Reference ```python @torch.compile def matmul_with_slice( A: torch.Tensor, # [Batch, K, M] B: torch.Tensor, # [N, K] bias: torch.Tensor, # [N] slice_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ Computes C = (A^T @ B^T + bias) and returns both the full output and a sliced view output[:, :slice_size, :]. """ A = A.to(torch.bfloat16) B = B.to(torch.bfloat16) bias = bias.to(torch.bfloat16) # A is [B, K, M] -> A^T is [B, M, K] # C = A^T @ B^T : [B, M, K] @ [K, N] = [B, M, N] output = torch.matmul(A.transpose(-2, -1), B.T) # Transpose output to [B, N, M] output = output.transpose(-2, -1).contiguous() # Add bias: [N] -> [N, 1] for broadcasting over M output = (output + bias.unsqueeze(-1)).contiguous() # Slice along the N dimension sliced_output = output[:, :slice_size, :] return output, sliced_output ``` ### Helion Kernel (Current Approach) ```python @hl.kernel() def matmul_with_slice( A: torch.Tensor, # [Batch, K, M] B: torch.Tensor, # [N, K] bias: torch.Tensor, # [N] slice_size: hl.constexpr, ) -> tuple[torch.Tensor, torch.Tensor]: Batch = A.size(0) K = A.size(1) M = hl.specialize(A.size(2)) N = hl.specialize(B.size(0)) output = torch.empty((Batch, N, M), device=A.device, dtype=torch.bfloat16) sliced_output = torch.empty( (Batch, slice_size, M), device=A.device, dtype=torch.bfloat16 ) for tile_b in hl.tile(Batch, block_size=1): for tile_m, tile_n in hl.tile([M, N]): acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) for tile_k in hl.tile(K): a_tile = A[tile_b.begin, tile_k, tile_m].to(torch.bfloat16) b_tile = B[tile_n, tile_k].to(torch.bfloat16) acc = torch.addmm(acc, a_tile.t(), b_tile.t()) # bias is [N] — broadcast over M bias_tile = bias[tile_n].to(torch.bfloat16) acc = acc + bias_tile[None, :] result = acc.to(torch.bfloat16).t() # [tile_n, tile_m] output[tile_b.begin, tile_n, tile_m] = result # ⚠️ Slice fusion: can only write tiles FULLY within the slice boundary if tile_n.end <= slice_size: sliced_output[tile_b.begin, tile_n, tile_m] = result return output, sliced_output ``` ## The Problem The slice fusion guard `if tile_n.end <= slice_size` only writes tiles that are **entirely** within the slice boundary. When `slice_size` is not a multiple of the N-dimension tile size, the boundary tile is silently dropped. **Concrete example:** `N = 192`, `tile_size_n = 64`, `slice_size = 48` | Tile range | `tile_n.end <= 48`? | Written to `sliced_output`? | |------------|---------------------|-----------------------------| | `[0, 64)` | 64 > 48 → **No** | ❌ Not written | | `[64, 128)`| 128 > 48 → **No** | ❌ Not written | | `[128, 192)`| 192 > 48 → **No** | ❌ Not written | **Result:** `sliced_output` is entirely zeros — completely wrong. **Another example:** `N = 192`, `tile_size_n = 64`, `slice_size = 100` | Tile range | `tile_n.end <= 100`? | Written to `sliced_output`? | |------------|----------------------|-----------------------------| | `[0, 64)` | 64 ≤ 100 → **Yes** | ✅ Written (64 elements) | | `[64, 128)`| 128 > 100 → **No** | ❌ Not written | | `[128, 192)`| 192 > 100 → **No** | ❌ Not written | **Result:** Only 64 out of 100 elements are written — the remaining 36 elements (indices 64–99) are missing. ## What We Need: Partial Tile Store To correctly fuse the slice, we need the ability to **partially store a tile** — writing only the rows that fall within the slice boundary. Conceptually: ```python # ✅ Desired behavior (pseudocode) if tile_n.begin < slice_size: # Determine how many rows of this tile fall within the slice valid_rows = min(tile_n.end, slice_size) - tile_n.begin # Store only the first `valid_rows` rows of `result` sliced_output[tile_b.begin, tile_n.begin : tile_n.begin + valid_rows, tile_m] = ( result[:valid_rows, :] ) ``` Without partial tile store support in Helion, this pattern cannot be expressed, and the slice must remain as a separate post-kernel operation — defeating the purpose of the fusion.