lee101/cutedsl

models accelerated with cutedsl

Python

3

60 commits

updated Sep 8, 2026

See the code

README

CuteDSL

Accelerated model inference via custom CuTeDSL and CUDA kernels. CuteDSL converts popular ML models into optimized versions with fused operations, custom attention kernels, and reduced memory allocations while maintaining output equivalence.

This is not just CuteDSL related the goal is to just build the fastest and best froneir models possible and catalog them, much like transformers/diffusers etc.... doesnt need to be cutedsl only!

A lot of this is created by autoresearch style bots that try to maintain evals while speeding up models and fusing kernels with cutedsl seems to be working well right now.

cutedsl.png

memecoin

Assiciated with this memecoin for raising funds for this project CuTeDSL memecoin

Hopefully this project can baloon out into a kind of transformers or ggml style repo with many useful accelerated models!

Please any help welcome!

Diffusion Frontier

The Qwen Edit / RTX 5090 notes describe an opt-in Diffusers CuTe normalization adapter, CUDA C correctness fixes, and reproducible kernel benchmarks. Full-checkpoint image parity remains an explicit evaluation step before enabling custom kernels by default.

diffusion_frontier is a no-surprises benchmark for comparing public text-to-image models on the same hardware. It catalogs Z-Image Turbo, FLUX.1 Schnell, Proteus v0.2/v0.4, RealVisXL v4, and Krea-2 Turbo; preflights pipeline, cache, disk, RAM, and VRAM requirements without downloading; supports standard Diffusers LoRAs; records synchronized warm latency and peak VRAM; and emits same-model retention metrics plus an opaque human-review manifest.

uv pip install -e '.[diffusion]'
cutedsl-diffusion-frontier --list-models
cutedsl-diffusion-frontier --preflight --models krea2-turbo flux-schnell realvisxl-v4
cutedsl-diffusion-frontier --models realvisxl-v4 --allow-download --steps 4 8 20 40
cutedsl-sdxl-frontier --models realvisxl-v4 --allow-download
cutedsl-diffusion-frontier --models flux-schnell --offload cuda \
  --quantization torchao-fp8dq --cache-prompt-embeds --release-text-encoders
cutedsl-consolidate-frontier /path/to/run-a/results.json /path/to/run-b/results.json

Uncached weights are never downloaded unless --allow-download is present. Krea-2 currently requires Diffusers from source plus Transformers 5.13 or newer in an isolated environment. For a quantized CUDA run with cached prompt embeddings, the runner stages Qwen conditioning before loading the resident denoiser so the 4B BF16 encoder and 13B FP8 transformer do not compete for a 32 GB GPU. Pixel metrics are deliberately labeled as same-model retention, not cross-model quality; use review_manifest.json for blind preference scores. The scheduler-aware SDXL runner supports native40, LCM, Hyper-SD, and combined style-LoRA profiles. diffusion_frontier.flow_resume provides exact same-prompt sigma-tail resume for FLUX/Krea; cross-prompt latent reuse remains an explicitly quality-gated learned experiment.

Latent Teleportation for Z-Image

The Latent Teleportation paper studies schedule-calibrated forecasting, aligned denoiser anchors, motion-pruned trajectory retrieval, variance-aware cache shrinkage, and a small learned cache adapter for Z-Image Turbo.

The experimental Z-Image latent-teleport cache adapter is a 7,312-parameter (31 KB) residual head. On its held-out fold it reduces latent forecast relL2 from 0.2320 for calibrated momentum and 0.2250 for the scalar cache gate to 0.2225. In the six-prompt long-warp image grid it improves PSNR from 21.48 dB to 21.92 dB; the variance-shrunk non-neural cache path gives the highest measured SSIM at 0.7883.

# Refit the focused cache adapter from stored trajectories.
python scripts/spec_train_cache_adapter.py --device auto

# Reproduce the teacher/local/cache/shrinkage/adapter image grid.
python scripts/spec_knn_visual.py \
  --top-k 8 --include-shrinkage \
  --adapter-dir /path/to/downloaded/adapter \
  --output-dir results/speculative/cache-procedure-visual \
  --figure paper/figures/cache_procedure_grid.png

The released adapter targets the recorded 512px, 16-step Z-Image Turbo schedule. It is not a standalone image model; new solvers, resolutions, and step schedules require trajectory capture and refitting.

Models

Shared-host interpolation benchmark

python -m scripts.bench_quantile_interpolation runs a CPU-only matched microbenchmark without downloading or loading models. It alternates baseline and candidate order, checks exact parity, and reports all round timings. On the shared host on 2026-09-05, FP32 [8,5,64] interpolation measured 212.35 → 120.76 microseconds (1.76x; median of seven 200-call rounds). This is not an end-to-end forecasting or GPU throughput measurement. The candidate computes float32 quantile metadata on the CPU, avoiding four GPU scalar reads per interpolated quantile while retaining the existing arithmetic. Direct quantile selection is unchanged. Use --device cuda only with shared-GPU headroom; verify with python -m pytest tests/test_quantile_interpolation.py.

CuteChronos2

Amazon Chronos-2 is a state-of-the-art time series forecasting model. CuteChronos2 is a from-scratch reimplementation with:

  • Custom Triton kernels for unscaled attention, RoPE, RMS LayerNorm, fused LN+Linear, fused MLP, fused preprocessing, and fused output transform
  • C++/CUDA preprocessing kernels for NaN-aware normalization and patching
  • Memory optimizations: in-place residual adds, cached position_ids, .reshape() instead of .contiguous().view()
  • torch.compile support with reduce-overhead mode for additional speedup
  • 1.4x+ speedup over the original implementation with identical outputs (max abs error < 1e-4)

Architecture

CuteChronos2 reimplements the full Chronos-2 encoder-only transformer:

Input Series -> InstanceNorm -> Patching -> InputPatchEmbedding
    -> [EncoderBlock x 12] -> FinalLayerNorm -> OutputPatchEmbedding
    -> Inverse InstanceNorm -> Quantile Predictions

Each EncoderBlock:
    TimeSelfAttention (RoPE, unscaled attention)
    GroupSelfAttention (cross-series attention)
    FeedForward (RMS LN + MLP)

Triton Kernels

KernelWhat it fusesSpeedup source
unscaled_attentionQK^T + mask + softmax + V multiply (no 1/sqrt(d_k) scaling)Avoids materializing S*S attention matrix
rms_layernormT5-style RMS norm (FP32 variance)Single kernel, no intermediate tensors
ropeinv_freq + cos/sin + Q/K rotationSingle kernel for both Q and K
fused_rms_norm_linearRMS LayerNorm + linear projectionEliminates normalized intermediate
fused_rms_norm_qkvRMS LayerNorm + Q/K/V projectionsOne normalization, three projections
fused_linear_reluLinear + ReLU activationEliminates pre-ReLU intermediate
fused_mlp_reluTwo-layer MLP (linear + relu + linear)Avoids materializing 3072-wide hidden
fused_preprocessNaN-aware normalize + arcsinh + patch + time_encTwo-phase: reduce then transform
fused_output_transformRearrange + sinh + unscaleSingle pass over output tensor

Quick Start

Requirements

  • Python >= 3.10
  • PyTorch >= 2.2.0
  • NVIDIA GPU with CUDA support
  • Triton >= 3.0.0

Installation

pip install uv
uv venv
source .venv/bin/activate

# Install the package
uv pip install -e .

# For validation against original model (optional)
uv pip install -e ".[chronos]"

# For development
uv pip install -e ".[dev]"

Download and Convert Chronos-2

# Convert the default model (amazon/chronos-2)
python -m cutechronos.convert --benchmark

# Convert with validation and compiled benchmark
python -m cutechronos.convert \
    --model-id amazon/chronos-2 \
    --benchmark \
    --benchmark-compiled \
    --output-json results.json

# The model is automatically downloaded from HuggingFace Hub
# and converted to CuteChronos2 format.

Use in Code

import torch
from cutechronos.model import CuteChronos2Model

# Load from HuggingFace checkpoint directory
model = CuteChronos2Model.from_pretrained("path/to/chronos-bolt-base")
model = model.to("cuda", torch.bfloat16)

# Or with torch.compile for even faster inference
model = CuteChronos2Model.from_pretrained_compiled(
    "path/to/chronos-bolt-base",
    compile_mode="reduce-overhead",
)
model = model.to("cuda", torch.bfloat16)

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

Pipeline API (matches upstream)

from cutechronos.pipeline import CuteChronos2Pipeline

pipe = CuteChronos2Pipeline.from_pretrained(
    "amazon/chronos-bolt-base",
    device="cuda",
    dtype=torch.bfloat16,
)

# Single series
context = torch.randn(512)
predictions = pipe.predict(context, prediction_length=30)
# predictions[0].shape: (1, 21, 30)

# Batch of variable-length series
contexts = [torch.randn(300), torch.randn(400), torch.randn(512)]
predictions = pipe.predict(contexts, prediction_length=30)

# Batch of multivariate series
multi_context = torch.randn(8, 3, 512)
quantiles, mean = pipe.predict_quantiles(
    multi_context,
    prediction_length=24,
    quantile_levels=[0.1, 0.5, 0.9],
)
# quantiles[0].shape: (3, 24, 3)
# mean[0].shape: (3, 24)

CuteChronos2Pipeline now supports:

  • 1-D tensors for a single univariate series
  • 2-D tensors for a batch of univariate series
  • 3-D tensors with shape (batch, n_variates, history_length) for multivariate inference
  • lists mixing univariate 1-D items and multivariate 2-D items

Dictionary-style covariate inputs are still not implemented in Cute. For past_covariates and future_covariates, use the upstream Chronos-2 pipeline for now.

Run Benchmarks

# Quick benchmark with synthetic data
python -m cutechronos.benchmark --device cuda --n-runs 20

# Benchmark with real data
python -m cutechronos.benchmark \
    --data-dir /path/to/csv/data \
    --symbols BTCUSD ETHUSD \
    --context-length 512 \
    --prediction-length 30

Parakeet STT Experiments

CuteDSL also includes a Parakeet ASR experiment harness for sweeping latency, real-time factor, and transcript drift across inference settings.

Install the optional ASR stack:

uv pip install -e ".[parakeet]"

Run a sweep over batch size, timestamp extraction, input mode, and autocast:

python -m cuteparakeet.benchmark \
    --device cuda \
    --audio-dir /home/lee/code/voicetype/test_audio \
    --reference-json /home/lee/code/voicetype/test_audio/ground_truth.json \
    --batch-sizes 1 4 8 12 16 \
    --timestamps off on \
    --input-modes path array \
    --amp-dtypes none bf16 \
    --allow-tf32 \
    --output-json /tmp/parakeet_results.json

Current direction from local experiments on the bundled voicetype samples:

  • batch_size=12 with input_mode=array, bf16, timestamps=off was the best throughput point so far.
  • Preloading audio arrays was slightly faster and slightly more accurate than path-based input on this corpus.
  • Enabling timestamps reduced throughput materially.
  • torch.compile on the Parakeet encoder was a large regression and is not recommended.
  • The first GPU call is materially slower than steady-state because NeMo does decoder graph setup on first use, so benchmark first-call and steady-state separately.

There is also an ONNX export utility:

python -m cuteparakeet.export_onnx \
    --device cuda \
    --batch-size 1 \
    --audio-seconds 4 \
    --output-dir /tmp/parakeet_onnx \
    --simplify

CuteChronos2 also has a TensorRT probe/export path. It exports the model through the pure PyTorch fallback kernels, writes a JSON report with torch._dynamo and ONNX details, and optionally builds a TensorRT engine when tensorrt is installed on the target host:

python -m cutechronos.export_tensorrt \
    --model-path amazon/chronos-2 \
    --output-dir /tmp/cutechronos_trt \
    --batch-size 8 \
    --context-length 512 \
    --max-batch-size 16 \
    --max-context-length 1024 \
    --build-engine \
    --fp16

To turn the JSON report into a short Markdown brief:

python /home/lee/code/dotfiles/trt_export_markdown.py \
    /tmp/cutechronos_trt/report.json \
    --output /tmp/cutechronos_trt/report.md

Cross-Language Wrapper Examples

There is also a wrapper benchmark under examples/chronos_wrappers/ that drives the same CuteChronos2 inference path from:

  • C
  • Go
  • Python

and compares those wrapper timings against upstream Chronos-2 Python.

make -C examples/chronos_wrappers

/home/lee/code/stock/.venv/bin/python examples/chronos_wrappers/benchmark.py \
    --device cuda \
    --runs 5 \
    --warmup 1 \
    --stock-data-dir /home/lee/code/stock/trainingdatadailybinance \
    --symbols BTCUSD ETHUSD SOLUSD

The benchmark includes the toy sequence [2, 4, 6, 8, 12] -> [14, 16, 18] plus a few real stock/crypto CSV series, and reports both absolute MAE and percentage MAPE%.

There is also a tuner at examples/chronos_wrappers/tune_cutechronos.py that sweeps pre-augmentation strategy, context length, quantile, inference ensemble strategy, dtype, and compile mode over rolling holdout windows, ranking configs by MAPE% by default.

For LoRA fine-tuning, there is a frontier runner at examples/lora_frontier.py. It drives the stock repo's Chronos-2 LoRA trainer, measures wall-clock training time, records validation/test MAE%, and emits the Pareto frontier for train_seconds versus val_mae_percent.

/home/lee/code/stock/.venv/bin/python examples/lora_frontier.py \
    --symbols BTCUSD \
    --preaugs percent_change log_returns baseline \
    --context-lengths 128 256 \
    --learning-rates 5e-5 1e-4 \
    --num-steps 50 100 250

To verify an exported LoRA checkpoint can be loaded back through CuteChronos inference, use examples/eval_exported_lora.py:

/home/lee/code/stock/.venv/bin/python examples/eval_exported_lora.py \
    --model-id /path/to/finetuned-ckpt \
    --symbol BTCUSD \
    --preaug percent_change \
    --context-length 128 \
    --prediction-length 24

For bulk hourly-crypto tuning, use examples/crypto_hourly_frontier.py. It stages a broad per-symbol inference sweep first, optionally tries dilation ensembles only on laggards, then launches bounded LoRA frontier runs only where the base model is still weak.

/home/lee/code/stock/.venv/bin/python examples/crypto_hourly_frontier.py \
    --data-root /home/lee/code/stock/trainingdatahourly/crypto \
    --run-dilation-pass \
    --run-lora-pass

To test whether a mixed-symbol crypto LoRA base transfers to harder targets, use examples/mixed_crypto_lora_transfer.py. It trains a LoRA on a pool of hourly crypto symbols, exports a normal finetuned-ckpt, and reports target-symbol validation/test MAPE before and after the mixed training.

/home/lee/code/stock/.venv/bin/python examples/mixed_crypto_lora_transfer.py \
    --max-train-symbols 8 \
    --eval-symbols SOLUSDT ETHUSDT

To test raw multivariate pair forecasting on hourly crypto data, use examples/multivariate_pair_eval.py. It now supports small frontier-style searches over partner pools, context lengths, strides, and cross_learning, so it is useful both for quick pair checks and for finding lower-MAE multivariate setups before porting them into native services.

/home/lee/code/stock/.venv/bin/python examples/multivariate_pair_eval.py \
    --backends cute \
    --target-symbols TAOUSD SOLUSDT BTCUSDT ETHUSDT \
    --partner-symbols ETHUSDT SOLUSDT BTCUSDT BTCUSD ETHUSD SOLUSD \
    --partner-pool-sizes 1 2 \
    --context-lengths 128 256 \
    --strides 12 24 \
    --cross-learning-options off on \
    --windows 64 \
    --warmup-windows 4 \
    --max-multivariate-runs 64

This emits both JSON and CSV summaries plus a Pareto-style frontier ranked by avg_latency_ms versus mean_mape_pct.

Run Tests

# All tests (requires chronos-forecasting for model comparison tests)
pytest cutechronos/tests/ -v

# Tests that don't require the original model
pytest cutechronos/tests/ -v -m "not model_required"

# Individual kernel tests
pytest cutechronos/tests/test_attention.py -v
pytest cutechronos/tests/test_rms_layernorm.py -v
pytest cutechronos/tests/test_rope.py -v

Benchmark Results

On RTX 5090, Chronos-2 base (768 d_model, 12 layers), B=1 L=512:

ImplementationLatency (ms)SpeedupGPU Memory (MB)
Original Chronos2Pipeline30.9baseline248
CuteChronos2Model (eager)24.01.3x247
CuteChronos2Model (torch.compile)1.324.4x237

The compiled mode uses torch.compile(mode="reduce-overhead") which captures CUDA graphs for near-zero kernel launch overhead. This is the recommended mode for production inference with fixed input shapes.

Second Model: CuteZImage

Z-Image Turbo is 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
  • Fused output epilogue: RMSNorm + AdaLN gate + residual add in one launch
  • Complex-valued RoPE kernel: Fused reshape + complex multiply + flatten
  • RMS LayerNorm kernel: Triton-accelerated T5-style normalization
  • PyTorch SDPA attention: Leverages FlashAttention v2 automatically
  • Regional torch.compile: Compiles the repeated refiner/transformer blocks without placing variable-length pipeline orchestration in the initial graph
  • 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 pipeline defaults to ZIMAGE_COMPILE_MODE=regional, which keeps the variable-length orchestration eager and compiles the repeated denoiser blocks. Use regional:max-autotune for a longer compile with additional tuning. Whole- transformer modes remain available explicitly, but regional compilation is the safe production setting because the model returns Python lists and uses custom CUDA ops.

For quality-preserving fresh generations, keep the default accelerated QKV path and automatic SDPA, leave CUTEZIMAGE_USE_TRITON_ROPE=0 and leave the default CUTEZIMAGE_FUSED_RESIDUAL=1; do not enable latent teleportation/Taylor skips. The latter are useful research controls, but the 16-step image sweep shows visible retention loss even at k=1. Use the cache adapter only for its validated offline cache procedure, or use exact replay for an identical cached request.

The output epilogue is enabled by default after qualification on an RTX 3090 Ti. At the production (1, 2048, 3840) shape its isolated operation is 3.86x faster, and a production-width block is 5.8% faster. A 10-block bf16 model probe measured 0.75% relative output RMSE (maximum absolute error 0.018). Disable it for strict baseline numerics:

CUTEZIMAGE_FUSED_RESIDUAL=0 \
  .venv/bin/python -m cutezimage.benchmark --pipeline

There is also a pure stable-diffusion.cpp Z-Image benchmark path for testing the "all C/C++ runtime" direction against the Python backends:

./scripts/setup_external_zimage.sh --build-sdcpp
python scripts/download_zimage_lowvram_weights.py

python -m zimageaccelerated.sdcpp_benchmark \
  --offload-to-cpu \
  --diffusion-fa \
  --vae-tiling \
  --clip-on-cpu

That writes a benchmark JSON file plus the generated image under benchmark_images/sdcpp/, and gives us a concrete baseline before attempting deeper C/CUDA rewrites of Z-Image-specific kernels or cache paths.

Third Model: CuteAnima

Anima 2.9B is an anime/illustration diffusion transformer: a Cosmos Predict2 backbone (40 layers, dim 2048, 16 heads) expanded to 2.9B, with a Qwen3 text encoder and an LLM adapter that cross-attends T5 token embeddings. CuteAnima keeps the reference numerics and attacks the parts around the GEMMs:

  • Meta-device load: the reference randomly initialises 2.9B fp32 parameters on CPU before overwriting them, and separately loads a 3.9 GB base transformer it then discards. CuteAnima instantiates on meta, adopts the bf16 checkpoint with assign=True, and never reads the base transformer - load drops from minutes to ~6-20 s.
  • Batched classifier-free guidance: one batch-2 transformer call per step instead of two batch-1 calls.
  • Fused adaLN modulation: LayerNorm + (1 + scale) * x + shift in one Triton pass, and a fused x + gate * delta residual. Each block runs that pair three times.
  • Prompt-embedding cache: production traffic reuses one negative prompt, so it is encoded once.

Measured on an RTX 3090 Ti at 832x1216, 28 steps, guidance 4.0, bf16, minimum of two runs on a shared desktop GPU (python -m cuteanima.benchmark):

variantlatencyvs sequential CFGPSNR vs reference
sequential CFG (reference loop)24.9 s1.00x-
batched CFG25.3 s0.98x35.0 dB
batched CFG + fused kernels25.7 s0.97x35.7 dB
sequential CFG + torch.compile19.3 s1.29x-
batched CFG + fused + torch.compile19.1 s1.30x31.2 dB

The honest read: at this resolution Anima is GEMM-bound, so the pointwise fusions land inside run-to-run noise and only torch.compile moves the wall clock. The fused kernels are kept because they are the layer that matters at small resolutions and for the compile-free path, and benchmark_kernels.py reports their isolated numbers. PSNR differences are floating-point association, not quality: the images are the same picture from the same seed.

uv pip install -e '.[anima]'
python -m cuteanima.benchmark --variants reference batched fused fused_compile --repeat 2
python -m cuteanima.benchmark_kernels
pytest cuteanima/tests -q

Project Structure

cutedsl/
  cutechronos/
    model.py              # Full CuteChronos2Model (encoder + weight loading)
    pipeline.py           # Drop-in pipeline wrapper
    convert.py            # Model download + conversion CLI
    benchmark.py          # End-to-end benchmarking
    kernels.py            # C++/CUDA extension loader
    modules/
      _fallbacks.py       # Pure PyTorch fallback implementations
      time_attention.py   # Fused TimeSelfAttention module
      group_attention.py  # Fused GroupSelfAttention module
      feedforward.py      # Fused FeedForward module
      output.py           # Fused output head module
    triton_kernels/
      attention.py        # Unscaled tiled attention (FlashAttention-style)
      rms_layernorm.py    # T5-style RMS LayerNorm
      rope.py             # Fused RoPE (inv_freq + cos/sin + apply)
      fused_layernorm_linear.py  # LN + Linear / LN + QKV
      fused_mlp.py        # Linear+ReLU / full MLP fusion
      fused_preprocess.py # NaN-aware preprocessing pipeline
      fused_output.py     # Rearrange + sinh + unscale
    cpp/
      preprocessing.cpp   # C++ preprocessing kernels
      preprocessing.cu    # CUDA preprocessing kernels
    tests/
      test_model.py       # Model equivalence tests
      test_pipeline.py    # Pipeline API tests
      test_attention.py   # Attention kernel tests
      ...
  cutezimage/
    model.py              # CuteZImageTransformer (30 layers + weight loading)
    triton_kernels/
      rms_norm.py         # RMS LayerNorm
      fused_silu_gate_ffn.py  # SiLU + gating + FFN fusion
      fused_adaln_norm.py # AdaLN + RMS norm fusion
      rope_complex.py     # Complex-valued RoPE
    tests/
      test_model.py       # Model component tests
      test_kernels.py     # Triton kernel correctness tests
  cuteanima/
    loader.py             # Meta-device Anima 2.9B load + pinned reference pipeline
    runner.py             # Batched-CFG denoise loop + prompt-embedding cache
    patch.py              # Binds fused block forward onto Cosmos blocks
    benchmark.py          # End-to-end variants + image quality metrics
    benchmark_kernels.py  # Isolated kernel timings
    triton_kernels/
      fused_adaln.py      # adaLN modulation + gated residual
    tests/
      test_kernels.py     # Kernel equivalence vs PyTorch reference

Adding New Models

CuteDSL is designed to accelerate more models beyond Chronos-2. The pattern:

  1. Identify the bottleneck operations in the original model
  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
  5. Benchmark to confirm speedup

License

Apache-2.0

lee101/cutedsl

models accelerated with cutedsl

Python

3

60 commits

updated Sep 8, 2026

See the code

README

CuteDSL

Accelerated model inference via custom CuTeDSL and CUDA kernels. CuteDSL converts popular ML models into optimized versions with fused operations, custom attention kernels, and reduced memory allocations while maintaining output equivalence.

This is not just CuteDSL related the goal is to just build the fastest and best froneir models possible and catalog them, much like transformers/diffusers etc.... doesnt need to be cutedsl only!

A lot of this is created by autoresearch style bots that try to maintain evals while speeding up models and fusing kernels with cutedsl seems to be working well right now.

cutedsl.png

memecoin

Assiciated with this memecoin for raising funds for this project CuTeDSL memecoin

Hopefully this project can baloon out into a kind of transformers or ggml style repo with many useful accelerated models!

Please any help welcome!

Diffusion Frontier

The Qwen Edit / RTX 5090 notes describe an opt-in Diffusers CuTe normalization adapter, CUDA C correctness fixes, and reproducible kernel benchmarks. Full-checkpoint image parity remains an explicit evaluation step before enabling custom kernels by default.

diffusion_frontier is a no-surprises benchmark for comparing public text-to-image models on the same hardware. It catalogs Z-Image Turbo, FLUX.1 Schnell, Proteus v0.2/v0.4, RealVisXL v4, and Krea-2 Turbo; preflights pipeline, cache, disk, RAM, and VRAM requirements without downloading; supports standard Diffusers LoRAs; records synchronized warm latency and peak VRAM; and emits same-model retention metrics plus an opaque human-review manifest.

uv pip install -e '.[diffusion]'
cutedsl-diffusion-frontier --list-models
cutedsl-diffusion-frontier --preflight --models krea2-turbo flux-schnell realvisxl-v4
cutedsl-diffusion-frontier --models realvisxl-v4 --allow-download --steps 4 8 20 40
cutedsl-sdxl-frontier --models realvisxl-v4 --allow-download
cutedsl-diffusion-frontier --models flux-schnell --offload cuda \
  --quantization torchao-fp8dq --cache-prompt-embeds --release-text-encoders
cutedsl-consolidate-frontier /path/to/run-a/results.json /path/to/run-b/results.json

Uncached weights are never downloaded unless --allow-download is present. Krea-2 currently requires Diffusers from source plus Transformers 5.13 or newer in an isolated environment. For a quantized CUDA run with cached prompt embeddings, the runner stages Qwen conditioning before loading the resident denoiser so the 4B BF16 encoder and 13B FP8 transformer do not compete for a 32 GB GPU. Pixel metrics are deliberately labeled as same-model retention, not cross-model quality; use review_manifest.json for blind preference scores. The scheduler-aware SDXL runner supports native40, LCM, Hyper-SD, and combined style-LoRA profiles. diffusion_frontier.flow_resume provides exact same-prompt sigma-tail resume for FLUX/Krea; cross-prompt latent reuse remains an explicitly quality-gated learned experiment.

Latent Teleportation for Z-Image

The Latent Teleportation paper studies schedule-calibrated forecasting, aligned denoiser anchors, motion-pruned trajectory retrieval, variance-aware cache shrinkage, and a small learned cache adapter for Z-Image Turbo.

The experimental Z-Image latent-teleport cache adapter is a 7,312-parameter (31 KB) residual head. On its held-out fold it reduces latent forecast relL2 from 0.2320 for calibrated momentum and 0.2250 for the scalar cache gate to 0.2225. In the six-prompt long-warp image grid it improves PSNR from 21.48 dB to 21.92 dB; the variance-shrunk non-neural cache path gives the highest measured SSIM at 0.7883.

# Refit the focused cache adapter from stored trajectories.
python scripts/spec_train_cache_adapter.py --device auto

# Reproduce the teacher/local/cache/shrinkage/adapter image grid.
python scripts/spec_knn_visual.py \
  --top-k 8 --include-shrinkage \
  --adapter-dir /path/to/downloaded/adapter \
  --output-dir results/speculative/cache-procedure-visual \
  --figure paper/figures/cache_procedure_grid.png

The released adapter targets the recorded 512px, 16-step Z-Image Turbo schedule. It is not a standalone image model; new solvers, resolutions, and step schedules require trajectory capture and refitting.

Models

Shared-host interpolation benchmark

python -m scripts.bench_quantile_interpolation runs a CPU-only matched microbenchmark without downloading or loading models. It alternates baseline and candidate order, checks exact parity, and reports all round timings. On the shared host on 2026-09-05, FP32 [8,5,64] interpolation measured 212.35 → 120.76 microseconds (1.76x; median of seven 200-call rounds). This is not an end-to-end forecasting or GPU throughput measurement. The candidate computes float32 quantile metadata on the CPU, avoiding four GPU scalar reads per interpolated quantile while retaining the existing arithmetic. Direct quantile selection is unchanged. Use --device cuda only with shared-GPU headroom; verify with python -m pytest tests/test_quantile_interpolation.py.

CuteChronos2

Amazon Chronos-2 is a state-of-the-art time series forecasting model. CuteChronos2 is a from-scratch reimplementation with:

  • Custom Triton kernels for unscaled attention, RoPE, RMS LayerNorm, fused LN+Linear, fused MLP, fused preprocessing, and fused output transform
  • C++/CUDA preprocessing kernels for NaN-aware normalization and patching
  • Memory optimizations: in-place residual adds, cached position_ids, .reshape() instead of .contiguous().view()
  • torch.compile support with reduce-overhead mode for additional speedup
  • 1.4x+ speedup over the original implementation with identical outputs (max abs error < 1e-4)

Architecture

CuteChronos2 reimplements the full Chronos-2 encoder-only transformer:

Input Series -> InstanceNorm -> Patching -> InputPatchEmbedding
    -> [EncoderBlock x 12] -> FinalLayerNorm -> OutputPatchEmbedding
    -> Inverse InstanceNorm -> Quantile Predictions

Each EncoderBlock:
    TimeSelfAttention (RoPE, unscaled attention)
    GroupSelfAttention (cross-series attention)
    FeedForward (RMS LN + MLP)

Triton Kernels

KernelWhat it fusesSpeedup source
unscaled_attentionQK^T + mask + softmax + V multiply (no 1/sqrt(d_k) scaling)Avoids materializing S*S attention matrix
rms_layernormT5-style RMS norm (FP32 variance)Single kernel, no intermediate tensors
ropeinv_freq + cos/sin + Q/K rotationSingle kernel for both Q and K
fused_rms_norm_linearRMS LayerNorm + linear projectionEliminates normalized intermediate
fused_rms_norm_qkvRMS LayerNorm + Q/K/V projectionsOne normalization, three projections
fused_linear_reluLinear + ReLU activationEliminates pre-ReLU intermediate
fused_mlp_reluTwo-layer MLP (linear + relu + linear)Avoids materializing 3072-wide hidden
fused_preprocessNaN-aware normalize + arcsinh + patch + time_encTwo-phase: reduce then transform
fused_output_transformRearrange + sinh + unscaleSingle pass over output tensor

Quick Start

Requirements

  • Python >= 3.10
  • PyTorch >= 2.2.0
  • NVIDIA GPU with CUDA support
  • Triton >= 3.0.0

Installation

pip install uv
uv venv
source .venv/bin/activate

# Install the package
uv pip install -e .

# For validation against original model (optional)
uv pip install -e ".[chronos]"

# For development
uv pip install -e ".[dev]"

Download and Convert Chronos-2

# Convert the default model (amazon/chronos-2)
python -m cutechronos.convert --benchmark

# Convert with validation and compiled benchmark
python -m cutechronos.convert \
    --model-id amazon/chronos-2 \
    --benchmark \
    --benchmark-compiled \
    --output-json results.json

# The model is automatically downloaded from HuggingFace Hub
# and converted to CuteChronos2 format.

Use in Code

import torch
from cutechronos.model import CuteChronos2Model

# Load from HuggingFace checkpoint directory
model = CuteChronos2Model.from_pretrained("path/to/chronos-bolt-base")
model = model.to("cuda", torch.bfloat16)

# Or with torch.compile for even faster inference
model = CuteChronos2Model.from_pretrained_compiled(
    "path/to/chronos-bolt-base",
    compile_mode="reduce-overhead",
)
model = model.to("cuda", torch.bfloat16)

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

Pipeline API (matches upstream)

from cutechronos.pipeline import CuteChronos2Pipeline

pipe = CuteChronos2Pipeline.from_pretrained(
    "amazon/chronos-bolt-base",
    device="cuda",
    dtype=torch.bfloat16,
)

# Single series
context = torch.randn(512)
predictions = pipe.predict(context, prediction_length=30)
# predictions[0].shape: (1, 21, 30)

# Batch of variable-length series
contexts = [torch.randn(300), torch.randn(400), torch.randn(512)]
predictions = pipe.predict(contexts, prediction_length=30)

# Batch of multivariate series
multi_context = torch.randn(8, 3, 512)
quantiles, mean = pipe.predict_quantiles(
    multi_context,
    prediction_length=24,
    quantile_levels=[0.1, 0.5, 0.9],
)
# quantiles[0].shape: (3, 24, 3)
# mean[0].shape: (3, 24)

CuteChronos2Pipeline now supports:

  • 1-D tensors for a single univariate series
  • 2-D tensors for a batch of univariate series
  • 3-D tensors with shape (batch, n_variates, history_length) for multivariate inference
  • lists mixing univariate 1-D items and multivariate 2-D items

Dictionary-style covariate inputs are still not implemented in Cute. For past_covariates and future_covariates, use the upstream Chronos-2 pipeline for now.

Run Benchmarks

# Quick benchmark with synthetic data
python -m cutechronos.benchmark --device cuda --n-runs 20

# Benchmark with real data
python -m cutechronos.benchmark \
    --data-dir /path/to/csv/data \
    --symbols BTCUSD ETHUSD \
    --context-length 512 \
    --prediction-length 30

Parakeet STT Experiments

CuteDSL also includes a Parakeet ASR experiment harness for sweeping latency, real-time factor, and transcript drift across inference settings.

Install the optional ASR stack:

uv pip install -e ".[parakeet]"

Run a sweep over batch size, timestamp extraction, input mode, and autocast:

python -m cuteparakeet.benchmark \
    --device cuda \
    --audio-dir /home/lee/code/voicetype/test_audio \
    --reference-json /home/lee/code/voicetype/test_audio/ground_truth.json \
    --batch-sizes 1 4 8 12 16 \
    --timestamps off on \
    --input-modes path array \
    --amp-dtypes none bf16 \
    --allow-tf32 \
    --output-json /tmp/parakeet_results.json

Current direction from local experiments on the bundled voicetype samples:

  • batch_size=12 with input_mode=array, bf16, timestamps=off was the best throughput point so far.
  • Preloading audio arrays was slightly faster and slightly more accurate than path-based input on this corpus.
  • Enabling timestamps reduced throughput materially.
  • torch.compile on the Parakeet encoder was a large regression and is not recommended.
  • The first GPU call is materially slower than steady-state because NeMo does decoder graph setup on first use, so benchmark first-call and steady-state separately.

There is also an ONNX export utility:

python -m cuteparakeet.export_onnx \
    --device cuda \
    --batch-size 1 \
    --audio-seconds 4 \
    --output-dir /tmp/parakeet_onnx \
    --simplify

CuteChronos2 also has a TensorRT probe/export path. It exports the model through the pure PyTorch fallback kernels, writes a JSON report with torch._dynamo and ONNX details, and optionally builds a TensorRT engine when tensorrt is installed on the target host:

python -m cutechronos.export_tensorrt \
    --model-path amazon/chronos-2 \
    --output-dir /tmp/cutechronos_trt \
    --batch-size 8 \
    --context-length 512 \
    --max-batch-size 16 \
    --max-context-length 1024 \
    --build-engine \
    --fp16

To turn the JSON report into a short Markdown brief:

python /home/lee/code/dotfiles/trt_export_markdown.py \
    /tmp/cutechronos_trt/report.json \
    --output /tmp/cutechronos_trt/report.md

Cross-Language Wrapper Examples

There is also a wrapper benchmark under examples/chronos_wrappers/ that drives the same CuteChronos2 inference path from:

  • C
  • Go
  • Python

and compares those wrapper timings against upstream Chronos-2 Python.

make -C examples/chronos_wrappers

/home/lee/code/stock/.venv/bin/python examples/chronos_wrappers/benchmark.py \
    --device cuda \
    --runs 5 \
    --warmup 1 \
    --stock-data-dir /home/lee/code/stock/trainingdatadailybinance \
    --symbols BTCUSD ETHUSD SOLUSD

The benchmark includes the toy sequence [2, 4, 6, 8, 12] -> [14, 16, 18] plus a few real stock/crypto CSV series, and reports both absolute MAE and percentage MAPE%.

There is also a tuner at examples/chronos_wrappers/tune_cutechronos.py that sweeps pre-augmentation strategy, context length, quantile, inference ensemble strategy, dtype, and compile mode over rolling holdout windows, ranking configs by MAPE% by default.

For LoRA fine-tuning, there is a frontier runner at examples/lora_frontier.py. It drives the stock repo's Chronos-2 LoRA trainer, measures wall-clock training time, records validation/test MAE%, and emits the Pareto frontier for train_seconds versus val_mae_percent.

/home/lee/code/stock/.venv/bin/python examples/lora_frontier.py \
    --symbols BTCUSD \
    --preaugs percent_change log_returns baseline \
    --context-lengths 128 256 \
    --learning-rates 5e-5 1e-4 \
    --num-steps 50 100 250

To verify an exported LoRA checkpoint can be loaded back through CuteChronos inference, use examples/eval_exported_lora.py:

/home/lee/code/stock/.venv/bin/python examples/eval_exported_lora.py \
    --model-id /path/to/finetuned-ckpt \
    --symbol BTCUSD \
    --preaug percent_change \
    --context-length 128 \
    --prediction-length 24

For bulk hourly-crypto tuning, use examples/crypto_hourly_frontier.py. It stages a broad per-symbol inference sweep first, optionally tries dilation ensembles only on laggards, then launches bounded LoRA frontier runs only where the base model is still weak.

/home/lee/code/stock/.venv/bin/python examples/crypto_hourly_frontier.py \
    --data-root /home/lee/code/stock/trainingdatahourly/crypto \
    --run-dilation-pass \
    --run-lora-pass

To test whether a mixed-symbol crypto LoRA base transfers to harder targets, use examples/mixed_crypto_lora_transfer.py. It trains a LoRA on a pool of hourly crypto symbols, exports a normal finetuned-ckpt, and reports target-symbol validation/test MAPE before and after the mixed training.

/home/lee/code/stock/.venv/bin/python examples/mixed_crypto_lora_transfer.py \
    --max-train-symbols 8 \
    --eval-symbols SOLUSDT ETHUSDT

To test raw multivariate pair forecasting on hourly crypto data, use examples/multivariate_pair_eval.py. It now supports small frontier-style searches over partner pools, context lengths, strides, and cross_learning, so it is useful both for quick pair checks and for finding lower-MAE multivariate setups before porting them into native services.

/home/lee/code/stock/.venv/bin/python examples/multivariate_pair_eval.py \
    --backends cute \
    --target-symbols TAOUSD SOLUSDT BTCUSDT ETHUSDT \
    --partner-symbols ETHUSDT SOLUSDT BTCUSDT BTCUSD ETHUSD SOLUSD \
    --partner-pool-sizes 1 2 \
    --context-lengths 128 256 \
    --strides 12 24 \
    --cross-learning-options off on \
    --windows 64 \
    --warmup-windows 4 \
    --max-multivariate-runs 64

This emits both JSON and CSV summaries plus a Pareto-style frontier ranked by avg_latency_ms versus mean_mape_pct.

Run Tests

# All tests (requires chronos-forecasting for model comparison tests)
pytest cutechronos/tests/ -v

# Tests that don't require the original model
pytest cutechronos/tests/ -v -m "not model_required"

# Individual kernel tests
pytest cutechronos/tests/test_attention.py -v
pytest cutechronos/tests/test_rms_layernorm.py -v
pytest cutechronos/tests/test_rope.py -v

Benchmark Results

On RTX 5090, Chronos-2 base (768 d_model, 12 layers), B=1 L=512:

ImplementationLatency (ms)SpeedupGPU Memory (MB)
Original Chronos2Pipeline30.9baseline248
CuteChronos2Model (eager)24.01.3x247
CuteChronos2Model (torch.compile)1.324.4x237

The compiled mode uses torch.compile(mode="reduce-overhead") which captures CUDA graphs for near-zero kernel launch overhead. This is the recommended mode for production inference with fixed input shapes.

Second Model: CuteZImage

Z-Image Turbo is 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
  • Fused output epilogue: RMSNorm + AdaLN gate + residual add in one launch
  • Complex-valued RoPE kernel: Fused reshape + complex multiply + flatten
  • RMS LayerNorm kernel: Triton-accelerated T5-style normalization
  • PyTorch SDPA attention: Leverages FlashAttention v2 automatically
  • Regional torch.compile: Compiles the repeated refiner/transformer blocks without placing variable-length pipeline orchestration in the initial graph
  • 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 pipeline defaults to ZIMAGE_COMPILE_MODE=regional, which keeps the variable-length orchestration eager and compiles the repeated denoiser blocks. Use regional:max-autotune for a longer compile with additional tuning. Whole- transformer modes remain available explicitly, but regional compilation is the safe production setting because the model returns Python lists and uses custom CUDA ops.

For quality-preserving fresh generations, keep the default accelerated QKV path and automatic SDPA, leave CUTEZIMAGE_USE_TRITON_ROPE=0 and leave the default CUTEZIMAGE_FUSED_RESIDUAL=1; do not enable latent teleportation/Taylor skips. The latter are useful research controls, but the 16-step image sweep shows visible retention loss even at k=1. Use the cache adapter only for its validated offline cache procedure, or use exact replay for an identical cached request.

The output epilogue is enabled by default after qualification on an RTX 3090 Ti. At the production (1, 2048, 3840) shape its isolated operation is 3.86x faster, and a production-width block is 5.8% faster. A 10-block bf16 model probe measured 0.75% relative output RMSE (maximum absolute error 0.018). Disable it for strict baseline numerics:

CUTEZIMAGE_FUSED_RESIDUAL=0 \
  .venv/bin/python -m cutezimage.benchmark --pipeline

There is also a pure stable-diffusion.cpp Z-Image benchmark path for testing the "all C/C++ runtime" direction against the Python backends:

./scripts/setup_external_zimage.sh --build-sdcpp
python scripts/download_zimage_lowvram_weights.py

python -m zimageaccelerated.sdcpp_benchmark \
  --offload-to-cpu \
  --diffusion-fa \
  --vae-tiling \
  --clip-on-cpu

That writes a benchmark JSON file plus the generated image under benchmark_images/sdcpp/, and gives us a concrete baseline before attempting deeper C/CUDA rewrites of Z-Image-specific kernels or cache paths.

Third Model: CuteAnima

Anima 2.9B is an anime/illustration diffusion transformer: a Cosmos Predict2 backbone (40 layers, dim 2048, 16 heads) expanded to 2.9B, with a Qwen3 text encoder and an LLM adapter that cross-attends T5 token embeddings. CuteAnima keeps the reference numerics and attacks the parts around the GEMMs:

  • Meta-device load: the reference randomly initialises 2.9B fp32 parameters on CPU before overwriting them, and separately loads a 3.9 GB base transformer it then discards. CuteAnima instantiates on meta, adopts the bf16 checkpoint with assign=True, and never reads the base transformer - load drops from minutes to ~6-20 s.
  • Batched classifier-free guidance: one batch-2 transformer call per step instead of two batch-1 calls.
  • Fused adaLN modulation: LayerNorm + (1 + scale) * x + shift in one Triton pass, and a fused x + gate * delta residual. Each block runs that pair three times.
  • Prompt-embedding cache: production traffic reuses one negative prompt, so it is encoded once.

Measured on an RTX 3090 Ti at 832x1216, 28 steps, guidance 4.0, bf16, minimum of two runs on a shared desktop GPU (python -m cuteanima.benchmark):

variantlatencyvs sequential CFGPSNR vs reference
sequential CFG (reference loop)24.9 s1.00x-
batched CFG25.3 s0.98x35.0 dB
batched CFG + fused kernels25.7 s0.97x35.7 dB
sequential CFG + torch.compile19.3 s1.29x-
batched CFG + fused + torch.compile19.1 s1.30x31.2 dB

The honest read: at this resolution Anima is GEMM-bound, so the pointwise fusions land inside run-to-run noise and only torch.compile moves the wall clock. The fused kernels are kept because they are the layer that matters at small resolutions and for the compile-free path, and benchmark_kernels.py reports their isolated numbers. PSNR differences are floating-point association, not quality: the images are the same picture from the same seed.

uv pip install -e '.[anima]'
python -m cuteanima.benchmark --variants reference batched fused fused_compile --repeat 2
python -m cuteanima.benchmark_kernels
pytest cuteanima/tests -q

Project Structure

cutedsl/
  cutechronos/
    model.py              # Full CuteChronos2Model (encoder + weight loading)
    pipeline.py           # Drop-in pipeline wrapper
    convert.py            # Model download + conversion CLI
    benchmark.py          # End-to-end benchmarking
    kernels.py            # C++/CUDA extension loader
    modules/
      _fallbacks.py       # Pure PyTorch fallback implementations
      time_attention.py   # Fused TimeSelfAttention module
      group_attention.py  # Fused GroupSelfAttention module
      feedforward.py      # Fused FeedForward module
      output.py           # Fused output head module
    triton_kernels/
      attention.py        # Unscaled tiled attention (FlashAttention-style)
      rms_layernorm.py    # T5-style RMS LayerNorm
      rope.py             # Fused RoPE (inv_freq + cos/sin + apply)
      fused_layernorm_linear.py  # LN + Linear / LN + QKV
      fused_mlp.py        # Linear+ReLU / full MLP fusion
      fused_preprocess.py # NaN-aware preprocessing pipeline
      fused_output.py     # Rearrange + sinh + unscale
    cpp/
      preprocessing.cpp   # C++ preprocessing kernels
      preprocessing.cu    # CUDA preprocessing kernels
    tests/
      test_model.py       # Model equivalence tests
      test_pipeline.py    # Pipeline API tests
      test_attention.py   # Attention kernel tests
      ...
  cutezimage/
    model.py              # CuteZImageTransformer (30 layers + weight loading)
    triton_kernels/
      rms_norm.py         # RMS LayerNorm
      fused_silu_gate_ffn.py  # SiLU + gating + FFN fusion
      fused_adaln_norm.py # AdaLN + RMS norm fusion
      rope_complex.py     # Complex-valued RoPE
    tests/
      test_model.py       # Model component tests
      test_kernels.py     # Triton kernel correctness tests
  cuteanima/
    loader.py             # Meta-device Anima 2.9B load + pinned reference pipeline
    runner.py             # Batched-CFG denoise loop + prompt-embedding cache
    patch.py              # Binds fused block forward onto Cosmos blocks
    benchmark.py          # End-to-end variants + image quality metrics
    benchmark_kernels.py  # Isolated kernel timings
    triton_kernels/
      fused_adaln.py      # adaLN modulation + gated residual
    tests/
      test_kernels.py     # Kernel equivalence vs PyTorch reference

Adding New Models

CuteDSL is designed to accelerate more models beyond Chronos-2. The pattern:

  1. Identify the bottleneck operations in the original model
  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
  5. Benchmark to confirm speedup

License

Apache-2.0