wehub-resource-sync
21 lines
779 B
C++
21 lines
779 B
C++
// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License").
|
|
// You may not use this file except in compliance with the License.
|
|
// To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
//
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
#include <torch/extension.h>
|
|
|
|
TORCH_LIBRARY(longlive_kernels, m)
|
|
{
|
|
m.def("dequantize_kv_cache_fp4(Tensor[] values, Tensor[] scale_factors, Tensor[] amax, int num_heads, int block_token_size, int dtype_code, float e2m1_max, float e4m3_max) -> Tensor");
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.doc() = "LongLive custom CUDA kernels";
|
|
}
|