Spaces:
Running
Fused Triton kernels for DoRA — 1.24× over the PyTorch reference on A100
Howdy kernels-community! 🤠
Sharing progress on a kernels-library package for the DoRA weight-merge path — pure Triton, torch-universal build, ported byte-for-byte from Scaling DoRA (arXiv:2603.22276) and its authoritative reference at https://github.com/sockeye44/dorafactors.
Source is published under https://huggingface.co/remyxai/dora-factored-kernel with the build.toml + build/torch-universal/ layout that get_kernel() expects, but currently as repo_type="model" — kernel-repository creation on the remyxai org is pending approval.
Why now: PEFT maintainers laid out the two-step integration path in sockeye44/dorafactors#1 — (1) ship the kernels on the Hub via the kernels library, (2) add an option to PEFT to load them. This package is step 1.
Perf: measured 1.24× vs the pure-PyTorch reference on A100 4K × 4K fp32 (median of 20), holding at 1.22–1.29× across the 4K → 16K shape ladder. The strided compose variant is what landed the win — the wrapper's transposes were eating the fusion at first pass, and it's what got the ratio above 1×.
Correctness: 60 parity tests across {fp16, bf16, fp32} × 6 tile-boundary shapes × 3 LoRA scalings, exercising both EVEN and masked-tail branches of the autotune configs. Forward + backward parity under DoRA's detached-norm policy (§4.3). Plus a PEFT merge_and_unload() end-to-end integration test — 3/3 green on A100.
Provenance: the compose (forward) and two-stage backward kernel bodies are kept byte-identical to the paper's source — @triton.autotune configs, @triton.heuristics, and stride patterns preserved. Only the Python launch wrappers were renamed to fit the package layout.
Step 2 is drafted as a PEFT PR with an opt-in USE_FACTORED_DORA_KERNEL flag, ready to file upstream once this lives at repo_type="kernel" (and, eventually, under kernels-community/). Happy to iterate on shape, naming, or the promotion path.
Access granted, thanks! The kernel is now published as repo_type="kernel" at https://huggingface.co/kernels/remyxai/dora-factored-kernel and consumable via kernels.get_kernel("remyxai/dora-factored-kernel", revision="v0.1.0", trust_remote_code=True).
End-to-end validated on A100 via this Colab — the notebook loads the kernel from Hub, runs a PEFT merge_and_unload() integration test (3/3 pass), hand-checks forward + backward parity, and confirms the 1.24× perf gate at 4K × 4K fp32, holding at 1.21–1.28× up to 16K × 8K.
License status: as flagged in the repo README, the upstream reference at sockeye44/dorafactors doesn't currently carry a LICENSE, so the ported kernel bodies (triton_compose.py, triton_backward.py) are provisionally labeled license: other (license_name: pending-upstream-authorization). We're coordinating with the paper authors at sockeye44/dorafactors#1 to add a permissive license (Apache-2.0 preferred, aligning with huggingface/peft and huggingface/kernels). Once that lands, this package flips to the same, with NOTICE attribution to the upstream repo.
Migration to kernels-community/ waits on that step.