CuteDSL - Accelerated ML Inference with Custom Triton and CUDA Kernels

CuteDSL converts popular ML models into optimized versions with fused operations, custom attention kernels, and reduced memory allocations – all while maintaining output equivalence.

The goal is to build the fastest possible frontier model implementations and catalog them, much like the transformers/diffusers ecosystem. A lot of this is created by autoresearch-style bots that try to maintain evals while speeding up models – fusing kernels with CuteDSL has been working well.

CuteChronos2 – 24x Faster Time Series Forecasting

The first target is Amazon Chronos-2, a state-of-the-art time series forecasting model. CuteChronos2 is a from-scratch reimplementation with custom Triton kernels for every major operation:

• Unscaled tiled attention (FlashAttention-style, avoids materializing the S*S attention matrix)
• Fused RoPE (inv_freq + cos/sin + Q/K rotation in one kernel)
• Fused RMS LayerNorm + Linear (eliminates normalized intermediate tensors)
• Fused MLP (two-layer MLP without materializing the 3072-wide hidden)
• Fused preprocessing (NaN-aware normalize + arcsinh + patch + time encoding)
• C++/CUDA preprocessing kernels for NaN-aware normalization and patching
• torch.compile support with reduce-overhead mode

Benchmark Results

On an RTX 5090 with Chronos-2 base (768 d_model, 12 layers), batch=1, length=512:

The compiled mode uses torch.compile(mode="reduce-overhead") which captures CUDA graphs for near-zero kernel launch overhead.

Quick Start

1
2
3
4
5
6

pip install uv
uv venv && source .venv/bin/activate
uv pip install -e .

# Convert and benchmark
python -m cutechronos.convert --benchmark --benchmark-compiled

1
2
3
4
5
6
7
8
9
10
11
12

import torch
from cutechronos.model import CuteChronos2Model

model = CuteChronos2Model.from_pretrained_compiled(
    "amazon/chronos-bolt-base",
    compile_mode="reduce-overhead",
)
model = model.to("cuda", torch.bfloat16)

context = torch.randn(1, 512, device="cuda")
with torch.inference_mode():
    quantile_preds = model(context)  # (batch, 21_quantiles, prediction_length)

CuteZImage – Accelerated Text-to-Image

The second model is Z-Image Turbo, a fast text-to-image diffusion model. CuteZImage reimplements the transformer backbone with:

• Fused SiLU-gated FFN – eliminates the 10240-wide intermediate allocation
• Fused AdaLN + RMS Norm – timestep conditioning fused with normalization
• Complex-valued RoPE kernel – fused reshape + complex multiply + flatten
• from_diffusers() weight loading – load from any HuggingFace Z-Image checkpoint

Architecture: 30 main layers + 2 refiner layers, dim=3840, 30 heads, SiLU-gated FFN (hidden=10240).

The Triton Kernel Approach

Each kernel fuses multiple PyTorch operations into a single GPU kernel launch, eliminating intermediate tensor allocations and memory bandwidth bottlenecks:

Adding New Models

The pattern for accelerating any model:

1. Profile the original model to identify bottleneck operations
2. Write Triton kernels that fuse multiple operations
3. Create a model class that loads original weights and uses fused kernels
4. Validate output equivalence within tight tolerance (max abs error < 1e-4)
5. Benchmark to confirm speedup

Check out the project on GitHub: github.com/lee101/cutedsl

Comment and share