diff --git a/.github/workflows/publish-fastvideo-kernel-hub.yml b/.github/workflows/publish-fastvideo-kernel-hub.yml new file mode 100644 index 000000000..90fcbdbd4 --- /dev/null +++ b/.github/workflows/publish-fastvideo-kernel-hub.yml @@ -0,0 +1,57 @@ +name: Publish FastVideo Kernel to Hugging Face Hub + +on: + push: + branches: + - main + paths: + - ".github/workflows/publish-fastvideo-kernel-hub.yml" + - "fastvideo-kernel/build.toml" + - "fastvideo-kernel/flake.nix" + - "fastvideo-kernel/CARD.md" + - "fastvideo-kernel/example.py" + - "fastvideo-kernel/torch-ext/**" + - "fastvideo-kernel/csrc/**" + - "fastvideo-kernel/include/**" + workflow_dispatch: + inputs: + branch: + description: "Hub branch to upload the build" + default: "" + required: false + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Build kernel via HF Jobs and upload to hao-ai-lab/fastvideo-kernel + uses: huggingface/kernel-builder-job@main + with: + token: ${{ secrets.HF_TOKEN }} + namespace: hao-ai-lab + flavor: cpu-xl + timeout: "21600" + script: | + set +x + export HF_TOKEN="${{ secrets.HF_TOKEN }}" + export GIT_LFS_SKIP_SMUDGE=1 + + git clone "${{ github.server_url }}/${{ github.repository }}" FastVideo + cd FastVideo + git checkout "${{ github.sha }}" + git submodule update --init --recursive fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk + cd fastvideo-kernel + + UPLOAD_BRANCH="${{ github.event.inputs.branch }}" + BRANCH_ARGS=() + if [ -n "${UPLOAD_BRANCH}" ]; then + BRANCH_ARGS=(--branch "${UPLOAD_BRANCH}") + fi + + nix run github:huggingface/kernels#kernel-builder -- build-and-upload \ + --max-jobs 4 \ + --cores 8 \ + --repo-id hao-ai-lab/fastvideo-kernel \ + "${BRANCH_ARGS[@]}" diff --git a/fastvideo-kernel/CARD.md b/fastvideo-kernel/CARD.md new file mode 100644 index 000000000..6fb9718ae --- /dev/null +++ b/fastvideo-kernel/CARD.md @@ -0,0 +1,37 @@ +--- +license: apache-2.0 +--- + +# FastVideo Kernel + +Hub-compatible FastVideo CUDA kernels packaged with Hugging Face `kernel-builder`. + +## Usage + +```python +from kernels import get_kernel + +fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1) +``` + +The module exposes: + +- `sta_fwd` +- `block_sparse_fwd` +- `block_sparse_bwd` +- `rms_norm` +- `layer_norm` +- `int8_quant` +- `int8_gemm` + +These APIs are thin wrappers around the native FastVideo kernels. Higher-level +fallback logic remains in the regular `fastvideo-kernel` Python package. + +## Requirements + +This first Hub build targets CUDA Hopper (`sm_90a`) because the native attention +kernels rely on Hopper/ThunderKittens features. + +## Source + +Source repository: https://github.com/hao-ai-lab/FastVideo/tree/main/fastvideo-kernel diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 00da6c625..d342684c2 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -40,6 +40,25 @@ Notes: ## Installation +### Hugging Face Kernel Hub + +FastVideo kernels are built and published with the Hugging Face +[`kernels`](https://huggingface.co/docs/kernels/index) project. Downstream code +loads the precompiled kernels from +[`hao-ai-lab/fastvideo-kernel`](https://huggingface.co/hao-ai-lab/fastvideo-kernel) +without building locally: + +```python +from kernels import get_kernel + +fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1) +``` + +The Hub package exposes the native CUDA entrypoints (`sta_fwd`, +`block_sparse_fwd`, `block_sparse_bwd`, `rms_norm`, `layer_norm`, `int8_quant`, +and `int8_gemm`). The full FastVideo Python package still provides the +high-level fallback logic used during local development. + ### Standard Installation (Local Development) This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime. diff --git a/fastvideo-kernel/build.toml b/fastvideo-kernel/build.toml new file mode 100644 index 000000000..4cdd3fa3b --- /dev/null +++ b/fastvideo-kernel/build.toml @@ -0,0 +1,50 @@ +[general] +name = "fastvideo_kernel" +version = 1 +license = "Apache-2.0" +backends = ["cuda"] +upstream = "https://github.com/hao-ai-lab/FastVideo.git" +source = "https://github.com/hao-ai-lab/FastVideo.git" + +[general.hub] +repo-id = "hao-ai-lab/fastvideo-kernel" + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", +] + +[kernel.fastvideo_kernel] +backend = "cuda" +depends = ["torch"] +include = [ + "csrc", + "csrc/turbodiffusion", + "include/cutlass/include", + "include/tk/include", + "include/tk/prototype", +] +src = [ + "csrc/attention/st_attn_h100.cu", + "csrc/attention/block_sparse_h100.cu", + "csrc/turbodiffusion/gemm/gemm.cu", + "csrc/turbodiffusion/norm/rmsnorm.cu", + "csrc/turbodiffusion/norm/layernorm.cu", + "csrc/turbodiffusion/quant/quant.cu", +] +cuda-capabilities = ["9.0a"] +cuda-flags = [ + "-DNDEBUG", + "-O3", + "-std=c++20", + "--use_fast_math", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-Xcompiler=-fno-strict-aliasing", + "-Xcompiler=-fPIC", + "-DTORCH_COMPILE", + "-DKITTENS_HOPPER", + "-DFASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION", + "-Xptxas=--warn-on-spills", +] diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu b/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu index 238e90d50..95d32b1f0 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu @@ -71,11 +71,12 @@ void int8_gemm( } +#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION void register_gemm(pybind11::module_ &m) { m.def("gemm_cuda", &int8_gemm); } +#endif - diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu index 6d54bf22c..4c50d9b95 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu @@ -9,10 +9,10 @@ #include "norm/layernorm.hpp" auto layer_norm( - at::Tensor const Input, + at::Tensor const Input, float eps, - std::optional W, - std::optional const B, + std::optional W, + std::optional const B, std::optional Output ) { int64_t const m = Input.size(0); @@ -93,7 +93,8 @@ auto layer_norm( return Output; } +#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION void register_layer_norm(pybind11::module_ &m) { m.def("layer_norm_cuda", &layer_norm); } - +#endif diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu index 17b64e890..85893095d 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu @@ -75,6 +75,8 @@ auto rms_norm( return Output; } +#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION void register_rms_norm(pybind11::module_ &m) { m.def("rms_norm_cuda", &rms_norm); } +#endif diff --git a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu index 9a56e3881..214abe211 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu @@ -70,6 +70,8 @@ auto quant( return std::make_tuple(Output, Output_S); } +#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION void register_quant(pybind11::module_ &m) { m.def("quant_cuda", &quant); } +#endif diff --git a/fastvideo-kernel/example.py b/fastvideo-kernel/example.py new file mode 100644 index 000000000..d452f6834 --- /dev/null +++ b/fastvideo-kernel/example.py @@ -0,0 +1,11 @@ +import torch +from kernels import get_kernel + + +fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1) + +x = torch.randn((16, 1024), device="cuda", dtype=torch.float16) +weight = torch.ones((1024,), device="cuda", dtype=torch.float16) +y = fastvideo_kernel.rms_norm(x, 1e-6, weight) + +print(y.shape) diff --git a/fastvideo-kernel/flake.nix b/fastvideo-kernel/flake.nix new file mode 100644 index 000000000..a476767bc --- /dev/null +++ b/fastvideo-kernel/flake.nix @@ -0,0 +1,18 @@ +{ + description = "Flake for FastVideo Kernel Hub builds"; + + inputs = { + self.submodules = true; + kernel-builder.url = "github:huggingface/kernels"; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/fastvideo-kernel/torch-ext/fastvideo_kernel/__init__.py b/fastvideo-kernel/torch-ext/fastvideo_kernel/__init__.py new file mode 100644 index 000000000..fb6d3f09e --- /dev/null +++ b/fastvideo-kernel/torch-ext/fastvideo_kernel/__init__.py @@ -0,0 +1,122 @@ +from typing import Optional + +import torch + +from ._ops import ops + + +def sta_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + kernel_t_size: int, + kernel_h_size: int, + kernel_w_size: int, + text_length: int, + process_text: bool, + has_text: bool, + kernel_aspect_ratio_flag: int, +) -> torch.Tensor: + return ops.sta_fwd( + q, + k, + v, + out, + kernel_t_size, + kernel_h_size, + kernel_w_size, + text_length, + process_text, + has_text, + kernel_aspect_ratio_flag, + ) + + +def block_sparse_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_block_sparse_index: torch.Tensor, + q2k_block_sparse_num: torch.Tensor, + kv_block_size: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + out, lse = ops.block_sparse_fwd( + q, + k, + v, + q2k_block_sparse_index, + q2k_block_sparse_num, + kv_block_size, + ) + return out, lse + + +def block_sparse_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + l_vec: torch.Tensor, + out_grad: torch.Tensor, + k2q_block_sparse_index: torch.Tensor, + k2q_block_sparse_num: torch.Tensor, + kv_block_size: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_grad, k_grad, v_grad = ops.block_sparse_bwd( + q, + k, + v, + out, + l_vec, + out_grad, + k2q_block_sparse_index, + k2q_block_sparse_num, + kv_block_size, + ) + return q_grad, k_grad, v_grad + + +def rms_norm( + input: torch.Tensor, + eps: float, + weight: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return ops.rms_norm_cuda(input, eps, weight, output) + + +def layer_norm( + input: torch.Tensor, + eps: float, + weight: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return ops.layer_norm_cuda(input, eps, weight, bias, output) + + +def int8_quant(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + quantized, scale = ops.quant_cuda(input) + return quantized, scale + + +def int8_gemm( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, +) -> torch.Tensor: + return ops.gemm_cuda(a, a_scale, b, b_scale, out) + + +__all__ = [ + "sta_fwd", + "block_sparse_fwd", + "block_sparse_bwd", + "rms_norm", + "layer_norm", + "int8_quant", + "int8_gemm", +] diff --git a/fastvideo-kernel/torch-ext/torch_binding.cpp b/fastvideo-kernel/torch-ext/torch_binding.cpp new file mode 100644 index 000000000..edb4a36ec --- /dev/null +++ b/fastvideo-kernel/torch-ext/torch_binding.cpp @@ -0,0 +1,93 @@ +#include + +#include "registration.h" +#include "torch_binding.h" + +torch::Tensor rms_norm_cuda( + torch::Tensor const& input, + double eps, + std::optional const& weight, + std::optional const& output) { + std::optional mutable_output = output; + auto result = rms_norm(input, static_cast(eps), weight, mutable_output); + TORCH_CHECK(result.has_value(), "rms_norm_cuda did not produce an output tensor"); + return result.value(); +} + +torch::Tensor layer_norm_cuda( + torch::Tensor const& input, + double eps, + std::optional const& weight, + std::optional const& bias, + std::optional const& output) { + auto result = layer_norm(input, static_cast(eps), weight, bias, output); + TORCH_CHECK(result.has_value(), "layer_norm_cuda did not produce an output tensor"); + return result.value(); +} + +std::tuple quant_cuda(torch::Tensor const& input) { + std::optional output; + std::optional output_scale; + auto result = quant(input, output, output_scale); + TORCH_CHECK(std::get<0>(result).has_value(), "quant_cuda did not produce an int8 tensor"); + TORCH_CHECK(std::get<1>(result).has_value(), "quant_cuda did not produce a scale tensor"); + return std::make_tuple(std::get<0>(result).value(), std::get<1>(result).value()); +} + +torch::Tensor gemm_cuda( + torch::Tensor const& a, + torch::Tensor const& a_scale, + torch::Tensor const& b, + torch::Tensor const& b_scale, + torch::Tensor c) { + int8_gemm(a, a_scale, b, b_scale, c); + return c; +} + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def( + "sta_fwd(Tensor q, Tensor k, Tensor v, Tensor! out, int kernel_t_size, int kernel_h_size, " + "int kernel_w_size, int text_length, bool process_text, bool has_text, " + "int kernel_aspect_ratio_flag) -> Tensor"); +#if defined(CUDA_KERNEL) + ops.impl("sta_fwd", torch::kCUDA, &sta_forward); +#endif + + ops.def( + "block_sparse_fwd(Tensor q, Tensor k, Tensor v, Tensor q2k_block_sparse_index, " + "Tensor q2k_block_sparse_num, Tensor kv_block_size) -> Tensor[]"); +#if defined(CUDA_KERNEL) + ops.impl("block_sparse_fwd", torch::kCUDA, &block_sparse_attention_forward); +#endif + + ops.def( + "block_sparse_bwd(Tensor q, Tensor k, Tensor v, Tensor out, Tensor l_vec, Tensor out_grad, " + "Tensor k2q_block_sparse_index, Tensor k2q_block_sparse_num, Tensor kv_block_size) -> Tensor[]"); +#if defined(CUDA_KERNEL) + ops.impl("block_sparse_bwd", torch::kCUDA, &block_sparse_attention_backward); +#endif + + ops.def("rms_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor? output=None) -> Tensor"); +#if defined(CUDA_KERNEL) + ops.impl("rms_norm_cuda", torch::kCUDA, &rms_norm_cuda); +#endif + + ops.def( + "layer_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor? bias=None, " + "Tensor? output=None) -> Tensor"); +#if defined(CUDA_KERNEL) + ops.impl("layer_norm_cuda", torch::kCUDA, &layer_norm_cuda); +#endif + + ops.def("quant_cuda(Tensor input) -> (Tensor, Tensor)"); +#if defined(CUDA_KERNEL) + ops.impl("quant_cuda", torch::kCUDA, &quant_cuda); +#endif + + ops.def("gemm_cuda(Tensor a, Tensor a_scale, Tensor b, Tensor b_scale, Tensor! c) -> Tensor"); +#if defined(CUDA_KERNEL) + ops.impl("gemm_cuda", torch::kCUDA, &gemm_cuda); +#endif +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/fastvideo-kernel/torch-ext/torch_binding.h b/fastvideo-kernel/torch-ext/torch_binding.h new file mode 100644 index 000000000..69e9c38c2 --- /dev/null +++ b/fastvideo-kernel/torch-ext/torch_binding.h @@ -0,0 +1,86 @@ +#pragma once + +#include +#include +#include + +#include + +torch::Tensor sta_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor o, + int kernel_t_size, + int kernel_h_size, + int kernel_w_size, + int text_length, + bool process_text, + bool has_text, + int kernel_aspect_ratio_flag); + +std::vector block_sparse_attention_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor q2k_block_sparse_index, + torch::Tensor q2k_block_sparse_num, + torch::Tensor kv_block_size); + +std::vector block_sparse_attention_backward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor o, + torch::Tensor l_vec, + torch::Tensor og, + torch::Tensor k2q_block_sparse_index, + torch::Tensor k2q_block_sparse_num, + torch::Tensor kv_block_size); + +std::optional rms_norm( + at::Tensor const& input, + float eps, + std::optional const& weight, + std::optional& output); + +std::optional layer_norm( + at::Tensor const input, + float eps, + std::optional weight, + std::optional const bias, + std::optional output); + +std::tuple, std::optional> quant( + torch::Tensor const& input, + std::optional& output, + std::optional& output_scale); + +void int8_gemm( + at::Tensor const& a, + at::Tensor const& a_scale, + at::Tensor const& b, + at::Tensor const& b_scale, + torch::Tensor& c); + +torch::Tensor rms_norm_cuda( + torch::Tensor const& input, + double eps, + std::optional const& weight, + std::optional const& output); + +torch::Tensor layer_norm_cuda( + torch::Tensor const& input, + double eps, + std::optional const& weight, + std::optional const& bias, + std::optional const& output); + +std::tuple quant_cuda(torch::Tensor const& input); + +torch::Tensor gemm_cuda( + torch::Tensor const& a, + torch::Tensor const& a_scale, + torch::Tensor const& b, + torch::Tensor const& b_scale, + torch::Tensor c);