ZzCompute
Built at PyTorch Helion Hackathon · Mar 14, 2026 · San Francisco, CA
1. Causal Conv1d — Depthwise Causal Convolution Implements the core Mamba/Mamba-2 causal convolution where each channel is convolved independently with past context. The kernel tiles across batch, channels, and sequence length with filter width (W=4) specialized at compile time for loop unrolling. Boundary handling uses masked loads (hl.load with extra_mask) to avoid host-side padding. Autotuned configs optimize block sizes (1024–4096), L2 cache grouping, and pipeline stages for the memory-bound pattern on B200. 2. Gated DeltaNet Chunk Forward H — Inter-Chunk Recurrence Implements the sequential bottleneck (Eq. 8), maintaining a K×V state across 64-timestep chunks. Parallelized across batch×head while processing chunks sequentially. Each step applies delta correction, exponential decay gating, and state accumulation. Uses IEEE float32 dot precision to prevent numerical drift. Dynamic shape fallback ensures correctness. 3. Gated DeltaNet Chunk Forward O — Output Computes per-chunk outputs combining inter-chunk state queries and intra-chunk causal attention (Eq. 9). Chunks run fully in parallel. Optimizations include exp2, fused hl.dot(..., acc=) accumulation, and B200 ACF scheduling. 4. Gated DeltaNet Recompute W/U — WY Transform Computes WY-transformed keys/values via two batched matmuls per chunk (Eq. 4–7). Uses exp2 gating and persistent-interleaved PID scheduling for high SM utilization. Single config passes all 12 shapes within the 12-minute compile limit.