cupydiff

High-Performance Discrete Diffusion Language Models in Triton + CuPy
Strictly NO PyTorch dependency. Pure functional architecture. Optimized for 8x NVIDIA RTX PRO 6000 96GB (Blackwell).

Highlights

• Strictly No PyTorch: Built entirely on Triton kernels and CuPy array operations with pure NumPy fallback for CPU testing.
• Pure Functional Design: Stateless mathematical functions, explicit weight dictionaries, and deterministic analytical forward and backward passes.
• Triton GPU Kernels:
  • Fused RMSNorm & AdaLN: High-bandwidth elementwise normalization and adaptive modulation.
  • Fused RoPE: Analytical rotary position embedding rotation.
  • Fused SwiGLU: silu(gate) * up with analytical backward gradients.
  • Bidirectional FlashAttention: Non-causal scaled dot-product attention executing in $O(L)$ SRAM instead of $O(L^2)$ global memory.
  • Cut Cross-Entropy (CCE): Chunked vocabulary projection that avoids materializing $(N, 49152)$ float32 logit tensors, preventing out-of-memory errors on massive context lengths.
• Hybrid Muon + AdamW Optimizer:
  • Muon: Quintic Newton-Schulz iteration (polar orthogonalization) for internal 2D weight matrices with Nesterov momentum.
  • AdamW: First and second moment buffers with decoupled weight decay for 1D vectors, biases, and token embeddings.
  • Cosine Learning Rate Schedule: Warmup + cosine decay.
• Multi-GPU Scalability (8x RTX PRO 6000 96GB Blackwell):
  • Native CuPy NCCL communicator with in-place gradient all-reduce.
  • Multi-process launcher (cupydiff-launch or python -m cupydiff.cli).
• Curriculum Shard Pipeline:
  • Compatible with binary token shards (.bin, uint16) and direct streaming from Hugging Face: vovaRL/slm388-corpus.

Installation

From Source / Local Development

git clone https://github.com/vovaRL/cupydiff.git
cd cupydiff
pip install -e .

With CUDA / Triton Acceleration (GPU)

pip install cupydiff[cuda]

(Installs cupy-cuda12x and triton)

Testing

cupydiff includes a complete standalone test suite that runs on CPU (using NumPy reference kernels) when no GPU is present, and automatically verifies Triton + CuPy GPU kernels when CUDA is available:

python tests/run_all.py

or with pytest:

pytest -v tests

Usage

1. Profiling & Benchmark Mode

Measure exact kernel latencies and throughput (tokens/sec):

# Single GPU
cupydiff --mode bench --bench_steps 20 --batch_size 64 --seq_len 2048

# 8x RTX PRO 6000 96GB
cupydiff --mode bench --nproc_per_node 8 --batch_size 512 --seq_len 2048

2. Pretraining on vovaRL/slm388-corpus

Trains the MDLM model using the curriculum binary shard loader:

cupydiff --mode train \
    --hf_dataset vovaRL/slm388-corpus \
    --data_dir /data/slm388-corpus \
    --batch_size 512 \
    --seq_len 2048 \
    --nproc_per_node 8 \
    --total_steps 85000 \
    --repo_id vovaRL/DiffLM-388M \
    --checkpoint_interval 2500

Hugging Face Jobs (8x RTX PRO 6000 96GB Blackwell)

To run distributed training on Hugging Face Jobs with 8x RTX PRO 6000:

hf jobs run \
    --flavor 8x-rtx-pro-6000 \
    --image nvidia/cuda:12.4.1-devel-ubuntu22.04 \
    --env HF_TOKEN="hf_your_token_here" \
    --command "pip install cupydiff[cuda] && cupydiff --mode train --hf_dataset vovaRL/slm388-corpus --batch_size 512 --seq_len 2048 --nproc_per_node 8"

Building and Pushing to PyPI

pip install build twine
python -m build
twine check dist/*
twine upload dist/*

Architecture Specification

License

Apache License 2.0.