Axiom

We implemented 4 GPU kernels using Helion DSL targeting NVIDIA B200 (Blackwell) GPUs: 1. causal_conv1d: Depthwise causal 1D convolution (core Mamba/Mamba-2 component). Key optimizations: eliminated padding allocation by using hl.load with extra_mask for inline zero-padding, per-shape configs with large block_sizes (2048-4096), deep pipelining (num_stages=7), and L2 cache grouping. 2. gated_deltanet_chunk_fwd_h: Inter-chunk state recurrence for Gated DeltaNet. Sequential bottleneck kernel with two hl.dot operations per chunk. Key optimizations: dot_precision="tf32" (~1.5x), exp2 instead of exp (hardware fast-path), gating diff instead of k (fewer FLOPs), fused hl.dot(acc=state) accumulation, per-shape configs with higher warps for tiny BH=1 shapes, and range_num_stages=[0,3] for inner loop pipelining (~1.3x). 3. gated_deltanet_recompute_w_u: Recomputes W and U via two matrix multiplications per chunk. Achieved 2.4x speedup using persistent_blocked PID with num_sm_multiplier=16, maxnreg=32, aggressive range_unroll_factors=[4], and l2_groupings=[16]. 4. gated_deltanet_chunk_fwd_o: Intra-chunk forward pass with 3 dot products (local QK attention + global state projection). Optimized with dot_precision="tf32", exp2 fast math, and num_warps=16.