ITADN

Support register-backed tensor caching across `hl.static_range` iterations (equivalent to Triton's `VAR_ARGS_ARRAY`)

#2041Openmengluy0125 创建于 2026-04-17
M
mengluy0125commented
**Is your feature request related to a problem? Please describe.** When writing two-pass kernels (e.g., layernorm: pass 1 gathers data + accumulates stats, pass 2 normalizes + writes output), there's no way to cache tensors in registers across two `hl.static_range` loops. This forces a re-gather from global memory in pass 2, adding ~67% more memory traffic compared to the equivalent Triton kernel. In Triton, this is solved with `VAR_ARGS_ARRAY`: ```python out_values: "VAR_ARGS_ARRAY" for i in range(len(input_group_ptrs)): out_values[i] = tl.load(...) # store in registers # Later, zero-cost register read: for i in range(len(input_group_ptrs)): value = out_values[i] # no memory op ``` In Helion, all attempts to cache tensors across `hl.static_range` iterations fail: - `list.append(tensor)` → not supported by tracer - `pre_allocated_list[i] = tensor` → `TypeInferenceError: Subscript assignment not supported` (SymInt key into SequenceType) - Temporary global memory buffer → works but adds write+read round-trip, negating the benefit **Describe the solution you'd like** Support storing and retrieving tensors from a list-like structure inside `hl.static_range` loops. Since `static_range` is unrolled at trace time, each `gathered_values[i]` resolves to a distinct tensor reference. The type propagation system should track this. Possible API options: ```python # Option A: Allow list subscript assignment with static_range index gathered_values = [None] * G for i in hl.static_range(G): gathered_values[i] = input_group[i][idx, tile_d] # Option B: Dedicated API gathered_values = hl.register_tensor_list(G, shape=[tile_o, tile_d], dtype=dtype) for i in hl.static_range(G): gathered_values[i] = input_group[i][idx, tile_d] ``` **Describe alternatives you've considered** 1. **Re-gather from global memory** (current approach): Works but 1.8x slower than Triton for the fused group_index_select + layernorm kernel (G=10, D=48, N=4096, O=4096 on H100). 2. **Temporary buffer in global memory**: Allocate `torch.empty([O, FULL_D])`, write in pass 1, read in pass 2. Compiles but adds 2G extra memory ops (G stores + G loads) — same total traffic as re-gather, with worse L2 sharing. **Additional context** - Benchmark: Triton fwd 0.0068ms (1204 GB/s) vs Helion fwd 0.0123ms (666 GB/s) for `G=10, D=48, N=4096, O=4096` on H100 bf16 - The gap is entirely due to the extra memory traffic in pass 2 (Triton: 3 mem ops/group, Helion: 5 mem ops/group) - This pattern is common in any two-pass reduction kernel (layernorm, softmax, etc.) where intermediate values need reuse
1 条评论