1st Place
Axiom
We implemented 4 GPU kernels using Helion DSL targeting NVIDIA B200 (Blackwell) GPUs: 1.
Core trick
Writes Mamba-2 and Gated DeltaNet kernels in Helion with per-shape block configs, tf32 dot products, and exp2 on B200.
Description
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.
README excerpt
<div align="center"> # Axiom ### Fast kernels. Measured on real silicon. Helion kernels for Mamba-2 and Gated DeltaNet, shaped for NVIDIA B200. [**1st Place · PyTorch Helion Hackathon**](https://cerebralvalley.ai/e/helion-hackathon/hackathon/gallery?project=9) [Results](#results) · [Quick start](#quick-start) · [Kernels](#the-kernels) · [Deep dives](#deep-dives) </div> <br> <p align="center"> <img src="assets/b200_gpu_utilization.png" alt="NVIDIA B200 utilization during Axiom optimizer trials" width="920"> </p> <p align="center"><sub>Live B200 optimizer trials: compile, verify, benchmark, repeat.</sub></p> <br> ## A smaller path to faster sequence models Axiom turns four expensive sequence operations into specialized GPU kernels. Each implementation is written in the [Helion](https://github.com/pytorch/helion) DSL, compiled for the target shape, checked against a PyTorch reference, and timed with the same evaluation harness. ## Results Measured on the `helion` H200 environment with `python eval.py benchmark`. The baseline and optimized kernels use the same harness and benchmark shapes. <p align="center"> <img src="assets/benchmark_speedup.svg" alt="Axiom geomean benchmark speedups on NVIDIA H200" width="920"> </p> | Kernel | Baseline | Axiom | Geomean speedup | |:--|--:|--:|--:| | Causal Conv1D | 249.5 µs | 9.5 µs | **29.0×** | | Recompute W/U | 359.0 µs | 5.6 µs | **57.8×** | <sub>Values above are the fastest minimum latency shown for each kernel; the sp