[feat]: add Hugging Face kernel hub packaging

.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[@]}"

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

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.
 

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",
+]

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
 
 
   
 
-

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<at::Tensor const> W,
-  std::optional<at::Tensor const> const B,
+  std::optional<at::Tensor> W,
+  std::optional<at::Tensor> const B,
   std::optional<at::Tensor> 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

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

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

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)

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 = ./.;
+    };
+}

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",
+]

fastvideo-kernel/torch-ext/torch_binding.cpp

@@ -0,0 +1,93 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+torch::Tensor rms_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& output) {
+  std::optional<at::Tensor> mutable_output = output;
+  auto result = rms_norm(input, static_cast<float>(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<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& bias,
+    std::optional<torch::Tensor> const& output) {
+  auto result = layer_norm(input, static_cast<float>(eps), weight, bias, output);
+  TORCH_CHECK(result.has_value(), "layer_norm_cuda did not produce an output tensor");
+  return result.value();
+}
+
+std::tuple<torch::Tensor, torch::Tensor> quant_cuda(torch::Tensor const& input) {
+  std::optional<torch::Tensor> output;
+  std::optional<torch::Tensor> 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)

fastvideo-kernel/torch-ext/torch_binding.h

@@ -0,0 +1,86 @@
+#pragma once
+
+#include <optional>
+#include <tuple>
+#include <vector>
+
+#include <torch/torch.h>
+
+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<torch::Tensor> 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<torch::Tensor> 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<at::Tensor> rms_norm(
+    at::Tensor const& input,
+    float eps,
+    std::optional<at::Tensor> const& weight,
+    std::optional<at::Tensor>& output);
+
+std::optional<at::Tensor> layer_norm(
+    at::Tensor const input,
+    float eps,
+    std::optional<at::Tensor> weight,
+    std::optional<at::Tensor> const bias,
+    std::optional<at::Tensor> output);
+
+std::tuple<std::optional<torch::Tensor>, std::optional<torch::Tensor>> quant(
+    torch::Tensor const& input,
+    std::optional<torch::Tensor>& output,
+    std::optional<torch::Tensor>& 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<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& output);
+
+torch::Tensor layer_norm_cuda(
+    torch::Tensor const& input,
+    double eps,
+    std::optional<torch::Tensor> const& weight,
+    std::optional<torch::Tensor> const& bias,
+    std::optional<torch::Tensor> const& output);
+
+std::tuple<torch::Tensor, torch::Tensor> 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);