Add Sol-Attn Kernel Builder source

8e9f35a verified 1 day ago

1.29 kB

license: apache-2.0
tags:
  - kernels
  - cuda
  - attention
  - triton
  - cute-dsl

Sol-Attn Kernel Builder Source

This repository contains the Hugging Face Kernel Builder packaging for Sol-Attn. The kernel implementation is pinned to NVIDIA's NVlabs/Sana commit 8a26fb0.

The packaged Python sources preserve that implementation. The only source relocation required by Kernel Hub is converting internal sol_attn.* imports to package-relative imports, so the kernel remains loadable under the version-isolated module name assigned by kernels.get_kernel(...).

Published builds are loaded from Efficient-Large-Model/Sol-Attn:

from kernels import get_kernel

kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)

out = kernel.sol_attn(
    q,  # Contiguous BF16 CUDA tensor [batch, tokens, heads, 128].
    k,  # Same shape, dtype, layout, and device as q.
    v,  # Same shape, dtype, layout, and device as q.
    tau=1.0,
    thresh_type="exact",
)

See SOURCE.md for provenance and the verification command.