File size: 422 Bytes
11e2c7c | 1 2 3 4 5 6 7 8 | """VBQ — a 75M-parameter decoder-only language model by Voussoir AI."""
import os
# Flash / mem-efficient SDPA kernels on RDNA3 are gated behind this flag in PyTorch's ROCm build.
# They are ~10x faster than the math fallback here, so enable them before torch is imported.
os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
os.environ.setdefault("PYTORCH_HIP_ALLOC_CONF", "expandable_segments:True")
|