dora-factored-kernel

Fused Triton kernels for the DoRA weight-adaptation operator from Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels (arXiv:2603.22276), packaged for huggingface/kernels via kernels.get_kernel(...).

DoRA's merge path materializes the effective weight (dora_scale / ||W + s·BA||) ⊙ (W + s·BA) over the full [d_out, d_in] tensor. This package fuses the elementwise passes into one memory-bandwidth-bound Triton kernel. Measured 1.24× vs the pure-PyTorch reference on A100 4K × 4K fp32, holding across the 4K → 16K shape ladder.

License status

The kernel bodies here are ported byte-for-byte from sockeye44/dorafactors, which does not currently carry a LICENSE file. Under default copyright, that leaves the upstream reference all-rights-reserved. We're coordinating with the authors at sockeye44/dorafactors#1 to add a permissive license (Apache-2.0 preferred, to align with huggingface/peft and huggingface/kernels); this package will inherit whatever license the upstream authors settle on, with attribution.

Until then, this repo is intended as a preview / coordination artifact, not for production distribution. The Python wrappers and PEFT integration surface (autograd.py, triton_compose_strided.py, integration test) are our own work and will be Apache-2.0 once the upstream reference resolves; the ported kernel bodies (triton_compose.py, triton_backward.py) are the specific pieces awaiting upstream authorization.

Usage

from kernels import get_kernel

dora = get_kernel("remyxai/dora-factored-kernel", revision="v0.1.0", trust_remote_code=True)

W_eff = dora.dora_factored_forward(
    base_weight,   # [d_out, d_in]
    lora_a,        # [r, d_in]
    lora_b,        # [d_out, r]
    scaling,       # float
    dora_scale,    # [d_out]
)  # -> [d_out, d_in]

Auto-dispatches to the fused Triton kernel on CUDA; falls back to the pure-PyTorch reference on CPU or when Triton is unavailable. Numerically equivalent to the reference within fp32 accumulation tolerance.

What it computes

The factored decomposition expands DoRA's squared row norm into three terms that never materialize the dense [d_out, d_in] product:

||W_i + s·(BA)_i||² = ||W_i||² + 2s·<W_i, (BA)_i> + s²·||(BA)_i||²
                     = base_i    + 2s·cross_i        + s²·gram_i

with O(d_out·r + r²) intermediates. The fused compose kernel then applies the magnitude rescale and adds base + s·BA in a single pass over the weight tensor.

The dense path (left) is what HuggingFace PEFT [Mangrulkar et al., 2022] and five other major frameworks surveyed in the paper's Appendix G — torchtune, Unsloth, SWIFT, LLaMA-Factory, Axolotl — use today: materialize a [d_in, d_in] identity matrix, materialize the dense B@A product, allocate a composed-weight temporary, take the row-wise L2 norm. The factored path (right) skips the dense product entirely. At d_in=8192, r=64, bf16 the per-module peak drops from ~640 MB to ~2 MB — 316× less memory, and unchanged under gradient checkpointing. Ratio grows to 634× at d_in=16K. Across hundreds of adapted modules in an 8-32B model, this is what turns an OOM at high rank into a routine training run.

Downstream reach: per Appendix G Table 15 of the paper (Feb 2026 snapshot), five of the six surveyed frameworks route their DoRA norm computation through PEFT — Unsloth falls back to PEFT for the norm step, and SWIFT / LLaMA-Factory / Axolotl carry no custom DoRA code and delegate entirely to PEFT. A PEFT-side improvement lifts all five without any additional integration work. Only torchtune ships an independent copy of the same torch.eye-based algorithm (modules/peft/dora.py) and would need a parallel update.

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 via torch.autograd.grad vs the pure-PyTorch reference under DoRA's detached-norm policy (§4.3).

pytest tests/ -m cuda

End-to-end with a real PEFT merge integration test: Colab — Hub-consumed A100 run.

Provenance

Compose (forward) and two-stage backward kernel bodies are ported byte-for-byte from the paper's authoritative reference implementation at sockeye44/dorafactors — @triton.autotune configs, @triton.heuristics predicates, and stride patterns preserved. Only the Python launch wrappers are renamed to fit the package layout.

A strided compose variant (triton_compose_strided.py) sits alongside the verbatim kernel and accepts explicit strides — needed so callers can pass zero-cost transposed views without materializing copies. This is what got the fused path above 1× vs the reference at production shapes.

PEFT integration

Follows the two-step path PEFT maintainers laid out in sockeye44/dorafactors#1:

1. Ship the kernels on the Hub via the kernels library — this package.
2. Add an option to PEFT to load them — drafted at smellslikeml/peft#18 as an opt-in USE_FACTORED_DORA_KERNEL flag routing DoraLinearVariant.merge_safe / merge_unsafe through kernels.get_kernel(...).

References

• Paper: Scaling DoRA (arXiv:2603.22276)
• Kernel source: sockeye44/dorafactors
• Factored-norm math (fork PR, not yet upstream): smellslikeml/peft#8
• Integration invitation: sockeye44/dorafactors#1