Compare commits

...
Author SHA1 Message Date
SolitaryThinker 0a231e0a1a [docs]: note the Hub build's Hopper-only arch subset is deliberate
Main's kernel wheels went multi-arch (0.3.x); the Hub build.toml stays
sm_90a because the TK attention kernels it packages require Hopper.
2026-07-10 19:31:18 -07:00
SolitaryThinker e9efd52211 [kernel]: declare mutation + return aliases in Hub torch-op schemas
sta_fwd and gemm_cuda return the tensor they write in place, and the
optional output args of rms_norm_cuda/layer_norm_cuda are mutated when
provided — none of which the schemas declared, which silently breaks
torch.compile. Annotate with alias sets (Tensor(a!) ... -> Tensor(a!)).
Schemas validated with torch.library on torch 2.x (parse + register).
2026-07-10 19:31:10 -07:00
SolitaryThinker 2b3109b68d [ci]: SHA-pin kernel-builder-job and serialize Hub publishes
The third-party action receives HF_TOKEN, so pin it to a commit SHA
instead of @main. Add a concurrency group so two merges can't race
multi-hour publish jobs against the same Hub repo.
2026-07-10 19:31:00 -07:00
SolitaryThinker ebd4d59951 [ci]: version-gate the HF Kernel Hub publish like the PyPI flow
Publish to the Hub only when the fastvideo-kernel version bumps
(same trigger + version-extraction as publish-kernel.yml), instead
of on every main push touching csrc/**. workflow_dispatch stays as
the manual override.
2026-07-10 19:30:47 -07:00
Aryan Kumar a6eb382d34 [feat]: add Hugging Face kernel hub packaging 2026-07-10 19:25:34 -07:00
13 changed files with 550 additions and 5 deletions
@@ -0,0 +1,95 @@
name: Publish FastVideo Kernel to Hugging Face Hub
# Version-gated like the PyPI flow (publish-kernel.yml): a push to main only
# publishes when the fastvideo-kernel version actually changes. Use
# workflow_dispatch as the manual override.
on:
push:
branches:
- main
paths:
- "fastvideo-kernel/pyproject.toml"
workflow_dispatch:
inputs:
branch:
description: "Hub branch to upload the build"
default: ""
required: false
# A Hub build takes hours; never race two publishes to the same repo.
concurrency:
group: ${{ github.workflow }}
cancel-in-progress: false
jobs:
check-version-change:
runs-on: ubuntu-latest
outputs:
version-changed: ${{ steps.check-version.outputs.changed }}
new-version: ${{ steps.check-version.outputs.new-version }}
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
fetch-depth: 2
- name: Check if version changed
id: check-version
run: |
cd fastvideo-kernel
# Get current commit's version from pyproject.toml
# Use ^ to match start of line to avoid matching minimum-version
NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
echo "New version: $NEW_VERSION"
# Get previous version from git history
# Note: git show expects path relative to repo root
OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
echo "Old version: $OLD_VERSION"
if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
echo "Version changed from $OLD_VERSION to $NEW_VERSION"
echo "changed=true" >> "$GITHUB_OUTPUT"
echo "new-version=$NEW_VERSION" >> "$GITHUB_OUTPUT"
else
echo "Version did not change"
echo "changed=false" >> "$GITHUB_OUTPUT"
fi
build:
needs: check-version-change
if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Build kernel via HF Jobs and upload to hao-ai-lab/fastvideo-kernel
# Pinned to a commit SHA: this third-party action receives HF_TOKEN.
uses: huggingface/kernel-builder-job@ee590cee1ff01e922310401f18d5554b8da448ac # main, 2026-06-03
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[@]}"
+37
View File
@@ -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
+19
View File
@@ -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.
+52
View File
@@ -0,0 +1,52 @@
[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",
]
# Deliberately Hopper-only (sm_90a): the TK attention kernels here require
# Hopper, unlike the multi-arch PyPI wheels (see .github/workflows/publish-kernel.yml).
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",
]
@@ -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
@@ -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
@@ -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
@@ -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
+11
View File
@@ -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)
+18
View File
@@ -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 = ./.;
};
}
@@ -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",
]
@@ -0,0 +1,99 @@
#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(
// `out` is written in place and returned, so the return must be declared
// as the same alias set — otherwise torch.compile can reorder/eliminate
// the mutation and silently produce wrong results.
"sta_fwd(Tensor q, Tensor k, Tensor v, Tensor(a!) 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(a!)");
#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
// When `output` is provided it is written in place and aliased by the
// return value; the schema must say so for torch.compile correctness.
ops.def("rms_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor(a!)? output=None) -> Tensor(a!)");
#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(a!)? output=None) -> Tensor(a!)");
#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
// `c` is the accumulator written in place and returned (same alias set).
ops.def("gemm_cuda(Tensor a, Tensor a_scale, Tensor b, Tensor b_scale, Tensor(a!) c) -> Tensor(a!)");
#if defined(CUDA_KERNEL)
ops.impl("gemm_cuda", torch::kCUDA, &gemm_cuda);
#endif
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
@@ -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);