File size: 2,970 Bytes
0387c74 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | #include "smallm_gemv.h"
#include <torch/extension.h>
#include <optional>
namespace {
torch::Tensor smallm_nvfp4_linear(
const torch::Tensor& input,
const torch::Tensor& packed_weight,
const torch::Tensor& weight_block_scales,
const torch::Tensor& weight_tensor_scale,
const std::optional<torch::Tensor>& bias) {
TORCH_CHECK(input.is_cuda(), "small-M GEMV requires a CUDA input");
TORCH_CHECK(
input.scalar_type() == at::kBFloat16,
"small-M GEMV input must be bfloat16");
TORCH_CHECK(
input.dim() >= 1 && input.is_contiguous(),
"small-M GEMV input must be contiguous");
TORCH_CHECK(
packed_weight.is_cuda() && packed_weight.scalar_type() == at::kByte &&
packed_weight.dim() == 2 && packed_weight.is_contiguous(),
"packed weight must be contiguous CUDA uint8 [N,K/2]");
TORCH_CHECK(
weight_block_scales.is_cuda() && weight_block_scales.dim() == 2 &&
weight_block_scales.is_contiguous() &&
weight_block_scales.element_size() == 1,
"weight block scales must be contiguous CUDA byte-sized [padded_N,padded_K/16]");
TORCH_CHECK(
weight_tensor_scale.is_cuda() &&
weight_tensor_scale.scalar_type() == at::kFloat &&
weight_tensor_scale.numel() == 1 &&
weight_tensor_scale.is_contiguous(),
"weight tensor scale must be one contiguous CUDA float32 value");
TORCH_CHECK(
input.device() == packed_weight.device() &&
input.device() == weight_block_scales.device() &&
input.device() == weight_tensor_scale.device(),
"all small-M GEMV tensors must use the same CUDA device");
const int64_t out_features = packed_weight.size(0);
const int64_t in_features = packed_weight.size(1) * 2;
TORCH_CHECK(
input.size(-1) == in_features,
"small-M GEMV expected input width ",
in_features,
" but got ",
input.size(-1));
TORCH_CHECK(
in_features > 0 && in_features % 32 == 0,
"small-M GEMV requires K divisible by 32");
TORCH_CHECK(
out_features > 0,
"small-M GEMV requires positive N");
TORCH_CHECK(
weight_block_scales.size(0) >= out_features &&
weight_block_scales.size(1) >= in_features / 16,
"weight block scale tensor is too small");
if (bias.has_value()) {
const auto& value = *bias;
TORCH_CHECK(
value.is_cuda() && value.scalar_type() == at::kBFloat16 &&
value.dim() == 1 && value.is_contiguous() &&
value.numel() == out_features &&
value.device() == input.device(),
"bias must be contiguous CUDA bfloat16 [N]");
}
return smallm_nvfp4_linear_cuda(
input,
packed_weight,
weight_block_scales,
weight_tensor_scale,
bias);
}
} // namespace
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
module.def(
"linear",
&smallm_nvfp4_linear,
"Fused BF16-activation x packed-NVFP4-weight small-M GEMV");
}
|