| #include "direct_w4a4_m1.h" |
|
|
| #include <torch/extension.h> |
|
|
| #include <optional> |
|
|
| namespace { |
|
|
| torch::Tensor direct_w4a4_m1_linear( |
| const torch::Tensor& activation_qdata, |
| const torch::Tensor& activation_block_scales, |
| const torch::Tensor& activation_tensor_scale, |
| const torch::Tensor& weight_qdata, |
| const torch::Tensor& weight_block_scales, |
| const torch::Tensor& weight_tensor_scale, |
| const std::optional<torch::Tensor>& bias) { |
| TORCH_CHECK( |
| activation_qdata.is_cuda() && |
| activation_qdata.scalar_type() == at::kByte && |
| activation_qdata.dim() == 2 && activation_qdata.is_contiguous(), |
| "activation qdata must be contiguous CUDA uint8 [padded_M,K/2]"); |
| TORCH_CHECK( |
| activation_qdata.size(0) >= 1, |
| "direct W4A4 M=1 requires at least one packed activation row"); |
| TORCH_CHECK( |
| activation_block_scales.is_cuda() && |
| activation_block_scales.dim() == 2 && |
| activation_block_scales.is_contiguous() && |
| activation_block_scales.element_size() == 1, |
| "activation block scales must be contiguous CUDA byte-sized [padded_M,padded_K/16]"); |
| TORCH_CHECK( |
| activation_tensor_scale.is_cuda() && |
| activation_tensor_scale.scalar_type() == at::kFloat && |
| activation_tensor_scale.numel() == 1 && |
| activation_tensor_scale.is_contiguous(), |
| "activation tensor scale must be one contiguous CUDA float32 value"); |
| TORCH_CHECK( |
| weight_qdata.is_cuda() && weight_qdata.scalar_type() == at::kByte && |
| weight_qdata.dim() == 2 && weight_qdata.is_contiguous(), |
| "weight qdata 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( |
| activation_qdata.device() == activation_block_scales.device() && |
| activation_qdata.device() == activation_tensor_scale.device() && |
| activation_qdata.device() == weight_qdata.device() && |
| activation_qdata.device() == weight_block_scales.device() && |
| activation_qdata.device() == weight_tensor_scale.device(), |
| "all direct W4A4 tensors must use the same CUDA device"); |
|
|
| const int64_t in_features = activation_qdata.size(1) * 2; |
| const int64_t out_features = weight_qdata.size(0); |
| TORCH_CHECK( |
| in_features > 0 && in_features % 32 == 0, |
| "direct W4A4 M=1 requires K divisible by 32"); |
| TORCH_CHECK( |
| weight_qdata.size(1) == activation_qdata.size(1), |
| "activation and weight packed K dimensions differ"); |
| TORCH_CHECK(out_features > 0, "direct W4A4 M=1 requires positive N"); |
| TORCH_CHECK( |
| activation_block_scales.size(0) >= 1 && |
| activation_block_scales.size(1) >= in_features / 16, |
| "activation block-scale tensor is too small"); |
| 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() == activation_qdata.device(), |
| "bias must be contiguous CUDA bfloat16 [N]"); |
| } |
| return direct_w4a4_m1_linear_cuda( |
| activation_qdata, |
| activation_block_scales, |
| activation_tensor_scale, |
| weight_qdata, |
| weight_block_scales, |
| weight_tensor_scale, |
| bias); |
| } |
|
|
| } |
|
|
| PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { |
| module.def( |
| "linear", |
| &direct_w4a4_m1_linear, |
| "Direct packed-NVFP4 activation x packed-NVFP4 weight M=1 linear"); |
| } |
|
|
