WarpDrive

Here's how each kernel works: 1. FP8 Quantization (fp8_quant) Memory-bound, fully parallel. Takes float32 activations, groups them (64-128 elements per group), finds the absolute max per group, computes a scale factor (absmax / 448.0), then divides and clamps to [-448, 448]. Our optimized version removes the baseline's 3x redundant abs/amax/divide operations — single pass per group. 2. Causal Conv1d (causal_conv1d) Memory-bound, fully parallel. Depthwise 1D convolution with causal (left) padding, used in Mamba/SSM architectures. Each channel is convolved independently with a tiny filter (W=3-4 taps): out[b,d,t] = bias[d] + sum_{k} weight[d,k] * x[b,d,t-W+1+k]. The filter is so small it fits in registers — the bottleneck is memory bandwidth. 3. GDN chunk_fwd_h (gated_deltanet_chunk_fwd_h) Compute-bound, sequential across chunks. The state recurrence kernel — maintains a hidden state h [K,V] and updates it chunk by chunk (C=64). Per chunk: (1) store h, (2) compute v_new = u - w @ h, (3) gate with exp(g_last - g_t), (4) decay h *= exp(g_last), (5) update h += k^T @ v_gated. The current VG6 version uses hl.dot with fp32 inputs and acc=state*decay to fuse the decay into the accumulator. Sequential bottleneck: NT chunks must run in order. 4. GDN chunk_fwd_o (gated_deltanet_chunk_fwd_o) Compute-bound, fully parallel. Computes outputs per chunk using the stored hidden states from chunk_fwd_h. Three matmuls per chunk: (1) inter = q @ h * exp(g) (cross-chunk state contribution), (2) intra = causal_mask(q @ k^T * exp(g_diff)) @ v_new (within-chunk attention), (3) output = inter + intra. The causal mask is built via o_t[:, None] >= o_t[None, :]. Embarrassingly parallel across all chunks. 5. GDN recompute_w_u (gated_deltanet_recompute_w_u) Compute-bound, fully parallel. Computes the WY-transformed keys and values needed by chunk_fwd_h. Per chunk: u = A @ (v * beta) and w = A @ (k * beta * exp(g)), where A is a [64,64] WY matrix. Two matmuls per chunk — load A once, loop over V chunks then K chunks. Feeds directly into chunk_fwd_h. Pipeline They chain together: recompute_w_u → chunk_fwd_h → chunk_fwd_o. All use chunk size C=64, K/V typically 64-128. FP8 quant and causal conv1d are independent standalone kernels.