| |
|
|
| #include <torch/all.h> |
| #include <torch/library.h> |
|
|
| #include <limits> |
|
|
| #if defined(CUDA_KERNEL) |
| #include <ATen/cuda/CUDAContext.h> |
| #include <c10/cuda/CUDAGuard.h> |
| #endif |
|
|
| #include "kernels/nexn2_moe_grouped_w4a16.cuh" |
| #include "kernels/nexn2_w4a16_gemv.cuh" |
| #include "kernels/grouped_w4a4_gemv_sm120.cuh" |
| #include "kernels/quantize_activations_nvfp4.cuh" |
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| namespace { |
|
|
| void check_cuda_contiguous(torch::Tensor const& t, const char* name) { |
| TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); |
| TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); |
| } |
|
|
| void check_bf16(torch::Tensor const& t, const char* name) { |
| check_cuda_contiguous(t, name); |
| TORCH_CHECK(t.scalar_type() == torch::kBFloat16, name, " must be torch.bfloat16"); |
| } |
|
|
| void check_u8(torch::Tensor const& t, const char* name) { |
| check_cuda_contiguous(t, name); |
| TORCH_CHECK(t.scalar_type() == torch::kUInt8, name, " must be torch.uint8"); |
| } |
|
|
| void check_f32(torch::Tensor const& t, const char* name) { |
| check_cuda_contiguous(t, name); |
| TORCH_CHECK(t.scalar_type() == torch::kFloat32, name, " must be torch.float32"); |
| } |
|
|
| void check_i32(torch::Tensor const& t, const char* name) { |
| check_cuda_contiguous(t, name); |
| TORCH_CHECK(t.scalar_type() == torch::kInt32, name, " must be torch.int32"); |
| } |
|
|
| int checked_int(int64_t value, const char* name) { |
| TORCH_CHECK(value > 0 && value <= std::numeric_limits<int>::max(), |
| name, " must fit in positive int"); |
| return static_cast<int>(value); |
| } |
|
|
| long checked_long(int64_t value, const char* name) { |
| TORCH_CHECK(value >= 0, name, " must be non-negative"); |
| return static_cast<long>(value); |
| } |
|
|
| int64_t swizzled_bytes(int64_t rows, int64_t dim) { |
| const int64_t row_supers = (rows + 127) / 128; |
| const int64_t col_supers = ((dim / 16) + 3) / 4; |
| return row_supers * col_supers * 512; |
| } |
|
|
| void check_same_device(torch::Tensor const& reference, |
| torch::Tensor const& value, |
| const char* name) { |
| TORCH_CHECK(reference.device() == value.device(), name, " must be on ", |
| reference.device()); |
| } |
|
|
| } |
|
|
| void w4a16_decode_gemv_bf16(torch::Tensor const& x_bf16, |
| torch::Tensor const& weight_packed, |
| torch::Tensor const& sfb, |
| double alpha, |
| torch::Tensor& out) { |
| check_bf16(x_bf16, "x_bf16"); |
| check_u8(weight_packed, "weight_packed"); |
| check_u8(sfb, "sfb"); |
| check_bf16(out, "out"); |
| TORCH_CHECK(x_bf16.dim() == 1 || (x_bf16.dim() == 2 && x_bf16.size(0) == 1), |
| "x_bf16 must have shape (K,) or (1,K)"); |
| const int64_t k = x_bf16.dim() == 1 ? x_bf16.size(0) : x_bf16.size(1); |
| TORCH_CHECK(weight_packed.dim() == 2 && weight_packed.size(1) == k / 2, |
| "weight_packed must have shape (N,K/2)"); |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({weight_packed.size(0)}), "out shape mismatch"); |
| #if defined(CUDA_KERNEL) |
| c10::cuda::CUDAGuard guard(x_bf16.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(x_bf16.get_device()).stream(); |
| const int rc = flash_rt::kernels::nexn2_w4a16_matvec_bf16( |
| x_bf16.data_ptr(), weight_packed.data_ptr(), sfb.data_ptr(), out.data_ptr(), |
| checked_int(weight_packed.size(0), "N"), checked_int(k, "K"), |
| static_cast<float>(alpha), stream); |
| TORCH_CHECK(rc == 0, "w4a16_decode_gemv_bf16 failed with rc=", rc); |
| #else |
| TORCH_CHECK(false, "grouped-moe-gemv was not built with CUDA support"); |
| #endif |
| } |
|
|
| void grouped_w4a16_gemv_bf16(torch::Tensor const& activations, |
| torch::Tensor const& weight_stack, |
| torch::Tensor const& sfb_stack, |
| torch::Tensor const& alpha_stack, |
| torch::Tensor const& expert_idx, |
| int64_t w_stride, |
| int64_t sfb_stride, |
| torch::Tensor& out) { |
| check_bf16(activations, "activations"); |
| check_u8(weight_stack, "weight_stack"); |
| check_u8(sfb_stack, "sfb_stack"); |
| check_f32(alpha_stack, "alpha_stack"); |
| check_i32(expert_idx, "expert_idx"); |
| check_bf16(out, "out"); |
| TORCH_CHECK(activations.dim() == 2, "activations must have shape (slots,K)"); |
| TORCH_CHECK(weight_stack.dim() >= 2, "weight_stack must be a flat or 3D uint8 stack"); |
| const int64_t slots = activations.size(0); |
| const int64_t k = activations.size(1); |
| const int64_t n = out.size(1); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({slots, n}), "out must have shape (slots,N)"); |
| TORCH_CHECK(expert_idx.sizes() == torch::IntArrayRef({slots}), "expert_idx must have shape (slots,)"); |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(w_stride > 0 && sfb_stride > 0, "w_stride and sfb_stride must be positive byte strides"); |
| #if defined(CUDA_KERNEL) |
| c10::cuda::CUDAGuard guard(activations.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(activations.get_device()).stream(); |
| const int rc = flash_rt::kernels::nexn2_moe_grouped_w4a16_bf16( |
| activations.data_ptr(), weight_stack.data_ptr(), sfb_stack.data_ptr(), |
| alpha_stack.data_ptr(), expert_idx.data_ptr(), out.data_ptr(), |
| checked_int(slots, "slots"), checked_int(n, "N"), checked_int(k, "K"), |
| checked_long(k, "a_stride"), checked_long(w_stride, "w_stride"), |
| checked_long(sfb_stride, "sfb_stride"), stream); |
| TORCH_CHECK(rc == 0, "grouped_w4a16_gemv_bf16 failed with rc=", rc); |
| #else |
| TORCH_CHECK(false, "grouped-moe-gemv was not built with CUDA support"); |
| #endif |
| } |
|
|
| void quantize_activations_nvfp4_bf16(torch::Tensor const& activations, |
| torch::Tensor& packed, |
| torch::Tensor& sfa) { |
| check_bf16(activations, "activations"); |
| check_u8(packed, "packed"); |
| check_u8(sfa, "sfa"); |
| TORCH_CHECK(activations.dim() == 2, "activations must have shape (M,K)"); |
| const int64_t m = activations.size(0); |
| const int64_t k = activations.size(1); |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(packed.sizes() == torch::IntArrayRef({m, k / 2}), |
| "packed must have shape (M,K/2)"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(m, k), |
| "sfa is too small for the CUTLASS SFA layout"); |
| check_same_device(activations, packed, "packed"); |
| check_same_device(activations, sfa, "sfa"); |
| #if defined(CUDA_KERNEL) |
| c10::cuda::CUDAGuard guard(activations.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(activations.get_device()).stream(); |
| const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16( |
| activations.data_ptr(), packed.data_ptr(), sfa.data_ptr(), |
| checked_int(m, "M"), checked_int(k, "K"), false, stream); |
| TORCH_CHECK(rc == 0, "quantize_activations_nvfp4_bf16 failed with rc=", rc); |
| #else |
| TORCH_CHECK(false, "grouped-moe-gemv was not built with CUDA support"); |
| #endif |
| } |
|
|
| void quantize_weights_nvfp4_bf16(torch::Tensor const& weights, |
| torch::Tensor& packed, |
| torch::Tensor& sfb) { |
| check_bf16(weights, "weights"); |
| check_u8(packed, "packed"); |
| check_u8(sfb, "sfb"); |
| TORCH_CHECK(weights.dim() == 2, "weights must have shape (N,K)"); |
| const int64_t n = weights.size(0); |
| const int64_t k = weights.size(1); |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(packed.sizes() == torch::IntArrayRef({n, k / 2}), |
| "packed must have shape (N,K/2)"); |
| TORCH_CHECK(sfb.numel() >= swizzled_bytes(n, k), |
| "sfb is too small for the CUTLASS SFB layout"); |
| check_same_device(weights, packed, "packed"); |
| check_same_device(weights, sfb, "sfb"); |
| #if defined(CUDA_KERNEL) |
| c10::cuda::CUDAGuard guard(weights.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(weights.get_device()).stream(); |
| const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16( |
| weights.data_ptr(), packed.data_ptr(), sfb.data_ptr(), |
| checked_int(n, "N"), checked_int(k, "K"), true, stream); |
| TORCH_CHECK(rc == 0, "quantize_weights_nvfp4_bf16 failed with rc=", rc); |
| #else |
| TORCH_CHECK(false, "grouped-moe-gemv was not built with CUDA support"); |
| #endif |
| } |
|
|
| void grouped_w4a4_gemv_bf16(torch::Tensor const& activations_packed, |
| torch::Tensor const& weight_stack, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb_stack, |
| torch::Tensor const& alpha_stack, |
| torch::Tensor const& expert_idx, |
| torch::Tensor& out) { |
| check_u8(activations_packed, "activations_packed"); |
| check_u8(weight_stack, "weight_stack"); |
| check_u8(sfa, "sfa"); |
| check_u8(sfb_stack, "sfb_stack"); |
| check_f32(alpha_stack, "alpha_stack"); |
| check_i32(expert_idx, "expert_idx"); |
| check_bf16(out, "out"); |
| TORCH_CHECK(activations_packed.dim() == 2, |
| "activations_packed must have shape (M,K/2)"); |
| TORCH_CHECK(weight_stack.dim() == 3, |
| "weight_stack must have shape (E,N,K/2)"); |
| TORCH_CHECK(expert_idx.dim() == 2, |
| "expert_idx must have shape (M,top_k)"); |
| const int64_t m = activations_packed.size(0); |
| const int64_t k_half = activations_packed.size(1); |
| const int64_t k = k_half * 2; |
| const int64_t experts = weight_stack.size(0); |
| const int64_t n = weight_stack.size(1); |
| const int64_t top_k = expert_idx.size(1); |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(weight_stack.size(2) == k_half, |
| "weight_stack K does not match activations_packed"); |
| TORCH_CHECK(expert_idx.size(0) == m, |
| "expert_idx first dimension must equal M"); |
| TORCH_CHECK(alpha_stack.numel() == experts, |
| "alpha_stack must have shape (E,)"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(m, k), |
| "sfa is too small for the CUTLASS SFA layout"); |
| TORCH_CHECK(sfb_stack.dim() == 2 && sfb_stack.size(0) == experts && |
| sfb_stack.size(1) >= swizzled_bytes(n, k), |
| "sfb_stack must have shape (E,sfb_bytes) with sufficient storage"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({m, top_k, n}), |
| "out must have shape (M,top_k,N)"); |
| check_same_device(activations_packed, weight_stack, "weight_stack"); |
| check_same_device(activations_packed, sfa, "sfa"); |
| check_same_device(activations_packed, sfb_stack, "sfb_stack"); |
| check_same_device(activations_packed, alpha_stack, "alpha_stack"); |
| check_same_device(activations_packed, expert_idx, "expert_idx"); |
| check_same_device(activations_packed, out, "out"); |
| #if defined(CUDA_KERNEL) |
| c10::cuda::CUDAGuard guard(activations_packed.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(activations_packed.get_device()).stream(); |
| const int rc = flash_rt::gemm::grouped_w4a4_gemv_sm120_bf16( |
| activations_packed.data_ptr(), weight_stack.data_ptr(), out.data_ptr(), |
| sfa.data_ptr(), sfb_stack.data_ptr(), alpha_stack.data_ptr(), |
| expert_idx.data_ptr(), checked_int(m, "M"), checked_int(top_k, "top_k"), |
| checked_int(n, "N"), checked_int(k, "K"), |
| checked_long(weight_stack.size(1) * weight_stack.size(2), "w_stride"), |
| checked_long(sfb_stack.size(1), "sfb_stride"), stream); |
| TORCH_CHECK(rc != 110, |
| "grouped_w4a4_gemv_bf16 requires SM120/SM121 block-scaled " |
| "MMA; use grouped_w4a16_gemv_bf16 on SM110"); |
| TORCH_CHECK(rc == 0, "grouped_w4a4_gemv_bf16 failed with rc=", rc); |
| #else |
| TORCH_CHECK(false, "grouped-moe-gemv was not built with CUDA support"); |
| #endif |
| } |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def("w4a16_decode_gemv_bf16(Tensor x_bf16, Tensor weight_packed, Tensor sfb, float alpha, Tensor! out) -> ()"); |
| ops.def("grouped_w4a16_gemv_bf16(Tensor activations, Tensor weight_stack, Tensor sfb_stack, Tensor alpha_stack, Tensor expert_idx, int w_stride, int sfb_stride, Tensor! out) -> ()"); |
| ops.def("quantize_activations_nvfp4_bf16(Tensor activations, Tensor! packed, Tensor! sfa) -> ()"); |
| ops.def("quantize_weights_nvfp4_bf16(Tensor weights, Tensor! packed, Tensor! sfb) -> ()"); |
| ops.def("grouped_w4a4_gemv_bf16(Tensor activations_packed, Tensor weight_stack, Tensor sfa, Tensor sfb_stack, Tensor alpha_stack, Tensor expert_idx, Tensor! out) -> ()"); |
| #if defined(CUDA_KERNEL) |
| ops.impl("w4a16_decode_gemv_bf16", torch::kCUDA, &w4a16_decode_gemv_bf16); |
| ops.impl("grouped_w4a16_gemv_bf16", torch::kCUDA, &grouped_w4a16_gemv_bf16); |
| ops.impl("quantize_activations_nvfp4_bf16", torch::kCUDA, &quantize_activations_nvfp4_bf16); |
| ops.impl("quantize_weights_nvfp4_bf16", torch::kCUDA, &quantize_weights_nvfp4_bf16); |
| ops.impl("grouped_w4a4_gemv_bf16", torch::kCUDA, &grouped_w4a4_gemv_bf16); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
