Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0a231e0a1a | ||
|
|
e9efd52211 | ||
|
|
2b3109b68d | ||
|
|
ebd4d59951 | ||
|
|
a6eb382d34 |
@@ -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[@]}"
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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);
|
||||
Reference in New Issue
Block a user