# Axiom

- **Event:** [PyTorch Helion Hackathon](https://cerebralvalley.ai/e/helion-hackathon)
- **When:** Sat, Mar 14 at 9:00 AM – 10:00 PM (PDT)
- **Where:** Shack15, San Francisco, CA
- **Placement:** 1st Place
- **Team:** [Yahya Alhinai](https://cerebralvalley.ai/u/yhinai)
- **GitHub:** https://github.com/yhinai/axiom
- **Gallery:** https://cerebralvalley.ai/e/helion-hackathon/hackathon/gallery
- **Page:** https://cerebralvalley.ai/e/helion-hackathon/hackathon/gallery/9

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.

---

Markdown version of https://cerebralvalley.ai/e/helion-hackathon/hackathon/gallery/9. Site index for agents: https://cerebralvalley.ai/llms.txt · full text: https://cerebralvalley.ai/llms-full.txt