Kernel learning: MixFFN and layer benchmarks

This repository collects FFN kernel learning work: Triton MixFFN (low-rank SVD + W8A8), flash decoding attention kernels, and layer-level experiments under speed/kernel/layer_bench/.

See speed/kernel/layer_bench/README.md for directory layout, providers, and benchmark options.

Requirements

CUDA 12.x toolchain, Python 3.10, and submodules (fasthad, gemm_int8). Verified install:

git submodule update --init --recursive
conda create -n kernel_learn python=3.10 -y
conda activate kernel_learn
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
# --force-reinstall
cd speed/third_party/fasthad
pip install -e . --no-build-isolation

# Apply CUDA-graph patch, then build gemm_int8
cd ../../kernel/layer_bench
bash support_gemm_int8_graph.sh

cd ../../third_party/gemm_int8
# cd ../gemm_int8
# # cd speed/third_party/gemm_int8
export CUDA_HOME=/usr/local/cuda
echo $CUDA_HOME
pip install cmake ninja

./build.sh
pip install -e . --no-build-isolation

cd ../../..
pip install pytest matplotlib
pip install -e . --no-build-isolation

Quick start (layer benchmark)

export ROOT_DIR=$(pwd)
export CUDA_VISIBLE_DEVICES=0
bash speed/kernel/layer_bench/bench_layer_do_exp.sh

Key code locations:

• Experiments: speed/kernel/layer_bench/do_exp/ (run_bench.py, run_autotune.py)
• MixFFN kernels: speed/kernel/layer_bench/ffn_kernel/
• Flash decoding: speed/kernel/layer_bench/fd_kernel/
• Example driver: speed/kernel/layer_bench/bench_layer_do_exp.sh

Todo

Task	Status
add rotation	⏩
change to CUDA kernel	⏩ 📑
tile strategy ablataion	⏩
block wise quantization	⏩
more quantization dtypes	⏩ 📑
intergating inference engine	⏩ 📑

mv speed/third_party/fasthad ./idk/fasthad.backup mv speed/third_party/gemm_int8 ./idk/gemm_int8.backup

git submodule add https://github.com/Dao-AILab/fast-hadamard-transform.git speed/third_party/fasthad git submodule add https://github.com/IST-DASLab/gemm-int8.git speed/third_party/gemm_int8

git add .gitmodules speed/third_party/fasthad speed/third_party/gemm_int8 git commit -m "Add third-party submodules"

Acknowledgments

Baseline from QuaRot, WSVD, gemm-int8 and related low-precision LLM kernel work.