Skip to main content

MOLTEN

Write the math. Get the kernel.

PyPI License Speedup


Molten turns mathematical operation specs into fused, portable CUDA kernels.

No tile loops. No schedules. No framework lock-in. The output is a .cu file. It compiles with nvcc. It runs without PyTorch.

Built by Tushar Sharma at ALIA Labs.

Install

pip install alia-molten

30 Seconds to a Fused Kernel

from molten import ZeroCompiler
from molten.ir import DataflowGraph, TensorShape

g = DataflowGraph("fused_rmsnorm")
x = g.add_input("x", TensorShape([2048, 5120]))
w = g.add_input("w", TensorShape([5120]))
out = g.rms_norm(x, w, "norm")
g.add_output(out)

compiler = ZeroCompiler()
kernels = compiler.compile(g)        # 3 ops -> 1 kernel
compiler.save(kernels, "output/")    # standalone .cu file

That's it. Three operations. One kernel. Zero CUDA written by hand.

What Happens Under the Hood

Math Spec -> DataflowGraph -> Optimizer -> Fusion Engine -> CUDA Codegen -> .cu

The fusion engine knows six rules:

Pattern What It Does
Elementwise chain Fuses N ops into 1. Kills N-1 memory round-trips.
MatMul + bias + activation Epilogue fusion. One kernel does matmul, adds bias, applies GELU.
RMSNorm Fuses reduce + normalize + scale. One pass over the data.
Softmax Fuses max + exp + sum + divide. Three passes become one.

Benchmarks

RTX 4090, torch 2.6.0+cu124, CUDA 12.8. Reproduce with python benchmarks/bench_molten_generated.py.

Molten-generated RMSNorm (zero hand-written CUDA):

Eager torch.compile Molten vs Eager
decode (1 token) 26.6 us n/a 16.2 us 1.65x
prefill (2048 tokens) 224.8 us n/a 95.9 us 2.34x
long (8192 tokens) 1547.7 us n/a 748.9 us 2.07x

The comparison is against eager, not torch.compile: the installed Triton is incompatible with this torch build, so no torch.compile baseline could be measured. The benchmark reports n/a rather than omitting the column, because a missing baseline must not read as a win.

These numbers supersede an earlier table that reported 4.6x against torch.compile. That measurement was invalid: the generated RMSNorm kernel dropped the weight tensor entirely (see CHANGELOG), so it was timed doing strictly less work than the reference it was compared against. The corrected kernel loads and applies the weight.

Correctness is verified separately by benchmarks/validate_correctness.py (20/20 against PyTorch eager; the Molten-generated RMSNorm checks use non-unity weights, the hand-written reference checks use w=ones) and every generated kernel is compile-checked by benchmarks/compile_check.py (7/7 under nvcc).

Hand-written fused RMSNorm+SiLU*gate (the target Molten is closing in on):

Eager (3 ops) Fused (1 kernel) Speedup
decode 207 us 27 us 7.6x
prefill 347 us 97 us 3.6x
long 1327 us 403 us 3.3x

Why Not torch.compile?

torch.compile generates Triton code tied to PyTorch. You can't deploy it without the full Python + PyTorch + Triton stack.

Molten generates a .cu file. Ship it to TensorRT, ONNX Runtime, a C++ server, a Jetson, whatever. It's just CUDA.

Tested On

RTX 4090 (Ada, sm_89) — every number in this README was measured here.

Not currently verifiable: an RTX 5090 is present in the development machine but the installed torch (cu124, built for sm_50–sm_90) cannot execute kernels on sm_120, and no H100 is available. Earlier RTX 5090 and H100 figures have been removed rather than carried forward unverified. Re-adding them requires a torch build for the target architecture and committed result artifacts.

Citation

@article{sharma2026molten,
  title={Molten: Fused GPU Kernel Generation from Mathematical Specifications},
  author={Sharma, Tushar},
  year={2026},
  url={https://github.com/TxsharDev/molten}
}

Roadmap

v0.1 (current) - IR, fusion engine, CUDA codegen, JIT runtime. RMSNorm and elementwise fusion proven. Scalar memory access.

v0.2 - Vectorized loads (float4/half2). This closes the gap where torch.compile currently wins at long sequences. fp16 I/O benchmarked end-to-end. @zero decorator dispatches generated kernels directly.

v0.3 - Attention fusion (Q@K softmax @V as one kernel). RoPE integration. Polyhedral loop optimization for complex fusion patterns. Auto-tuning via hardware counter feedback.

License

Apache-2.0 | ALIA Labs

Metadata

Release files for alia-molten 0.1.1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for alia-molten 0.1.1
File Size Uploaded Provenance Actions
alia_molten-0.1.1.tar.gz 24.6 kB Without provenance Details for alia_molten-0.1.1.tar.gz

Built distribution (wheel)

Table of built distributions (wheels) for alia-molten 0.1.1
File Interpreter ABI Platform Actions
alia_molten-0.1.1-py3-none-any.whl Python 3 none any Details for alia_molten-0.1.1-py3-none-any.whl

Total release size: 49.4 kB

Release files / alia_molten-0.1.1.tar.gz

Download URL alia_molten-0.1.1.tar.gz
Size 24.6 kB
Tags Source
SHA-256 checksum
How to use checksums
4ff07954c97b148954b278a9cd36307ea816c609caaf53056ecf4d85c72b97b0
BLAKE2b-256 checksum
How to use checksums
a6157989e621cc9b03cc27fd2a90cc5aa56a3169a57dc97388cee80e5ad66d6a
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/6.2.0 CPython/3.12.9

Release files / alia_molten-0.1.1-py3-none-any.whl

Download URL alia_molten-0.1.1-py3-none-any.whl
Size 24.8 kB
Tags Python 3
SHA-256 checksum
How to use checksums
91e610e6566c0f2d91fb13262a8679e4f938d6d715c1bfe649379b582d940f06
BLAKE2b-256 checksum
How to use checksums
ad9b3d3f0a0f29ddf1d6d94f26f17c76baee9731a2350711e2889687ebf5a067
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/6.2.0 CPython/3.12.9

Security for alia-molten 0.1.1

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Release history Release notifications | RSS feed

This release

0.1.1 This release

2 release files No provenance

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page