| |
|
|
| #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 "dequantize_fp4_sfa.cuh" |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| #include "gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cuh" |
| #include "gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cuh" |
| #include "gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh" |
| #include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cuh" |
| #include "gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh" |
| #endif |
| #include "gemm/fp4/sm110_dispatch.cuh" |
| #include "quantize/quantize_fp4_sfa.cuh" |
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| flash_rt::hub::Sm110GemmDispatch flash_rt::hub::sm110_gemm_dispatch = nullptr; |
| flash_rt::hub::Sm110GemmBiasDispatch |
| flash_rt::hub::sm110_gemm_bias_dispatch = nullptr; |
| flash_rt::hub::Sm110GemmBiasResidualDispatch |
| flash_rt::hub::sm110_gemm_bias_residual_dispatch = nullptr; |
| flash_rt::hub::Sm110GemmBiasGeluFp4Dispatch |
| flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch = nullptr; |
| flash_rt::hub::Sm110QuantizeBf16Dispatch |
| flash_rt::hub::sm110_quantize_bf16_dispatch = nullptr; |
|
|
| namespace { |
|
|
| void check_cuda_contiguous(torch::Tensor const& tensor, const char* name) { |
| TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); |
| TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); |
| } |
|
|
| void check_uint8_cuda(torch::Tensor const& tensor, const char* name) { |
| check_cuda_contiguous(tensor, name); |
| TORCH_CHECK(tensor.scalar_type() == torch::kUInt8, |
| name, " must have dtype torch.uint8"); |
| } |
|
|
| void check_fp16_cuda(torch::Tensor const& tensor, const char* name) { |
| check_cuda_contiguous(tensor, name); |
| TORCH_CHECK(tensor.scalar_type() == torch::kFloat16, |
| name, " must have dtype torch.float16"); |
| } |
|
|
| void check_bf16_cuda(torch::Tensor const& tensor, const char* name) { |
| check_cuda_contiguous(tensor, name); |
| TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16, |
| name, " must have dtype torch.bfloat16"); |
| } |
|
|
| 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); |
| } |
|
|
| int64_t swizzled_bytes(int64_t rows, int64_t dim) { |
| TORCH_CHECK(rows > 0 && dim > 0 && dim % 16 == 0, |
| "rows must be positive and dim must be positive/divisible by 16"); |
| const int64_t n_blocks = dim / 16; |
| const int64_t n_row_super = (rows + 127) / 128; |
| const int64_t n_col_super = (n_blocks + 3) / 4; |
| return n_row_super * n_col_super * 512; |
| } |
|
|
| void check_same_device(torch::Tensor const& a, torch::Tensor const& b, |
| const char* a_name, const char* b_name) { |
| TORCH_CHECK(a.get_device() == b.get_device(), |
| a_name, " and ", b_name, " must be on the same CUDA device"); |
| } |
|
|
| struct GemmShape { |
| int64_t m; |
| int64_t n; |
| int64_t k; |
| }; |
|
|
| #if defined(CUDA_KERNEL) |
| cudaDeviceProp const* current_device_properties(torch::Tensor const& anchor) { |
| return at::cuda::getDeviceProperties(anchor.get_device()); |
| } |
|
|
| void require_sm120(torch::Tensor const& anchor, const char* operation) { |
| auto const* props = current_device_properties(anchor); |
| TORCH_CHECK(props->major == 12 && props->minor == 0, |
| operation, " is an SM120 fused epilogue; got SM", |
| props->major, props->minor, |
| ". On SM110 use nvfp4_gemm_bf16 with fp4-fused-ops producers."); |
| } |
| #endif |
|
|
| GemmShape check_fp4_gemm_inputs( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb) { |
| check_uint8_cuda(a_packed, "a_packed"); |
| check_uint8_cuda(b_packed, "b_packed"); |
| check_uint8_cuda(sfa, "sfa"); |
| check_uint8_cuda(sfb, "sfb"); |
| TORCH_CHECK(a_packed.dim() == 2, "a_packed must have shape (M, K / 2)"); |
| TORCH_CHECK(b_packed.dim() == 2, "b_packed must have shape (N, K / 2)"); |
| const int64_t m = a_packed.size(0); |
| const int64_t n = b_packed.size(0); |
| const int64_t k_half = a_packed.size(1); |
| TORCH_CHECK(m > 0 && n > 0 && k_half > 0, "M, N, and K must be positive"); |
| TORCH_CHECK(b_packed.size(1) == k_half, |
| "a_packed and b_packed must have the same K / 2 dimension"); |
| const int64_t k = k_half * 2; |
| TORCH_CHECK(k % 16 == 0, "K must be divisible by 16"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(m, k), |
| "sfa is too small for CUTLASS SFA layout"); |
| TORCH_CHECK(sfb.numel() >= swizzled_bytes(n, k), |
| "sfb is too small for CUTLASS SFB layout"); |
| check_same_device(a_packed, b_packed, "a_packed", "b_packed"); |
| check_same_device(a_packed, sfa, "a_packed", "sfa"); |
| check_same_device(a_packed, sfb, "a_packed", "sfb"); |
| return {m, n, k}; |
| } |
|
|
| int select_sm110_variant(GemmShape const& shape, int64_t requested) { |
| if (requested >= 0) return static_cast<int>(requested); |
| if (shape.n >= 4 * shape.k) return 1; |
| if (shape.n == 3 * shape.k) return 2; |
| return 0; |
| } |
|
|
| } |
|
|
| void fp4_w4a4_gemv_warpsplit_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor& out, |
| double alpha, |
| int64_t warps, |
| int64_t stages) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(shape.m == 1, |
| "warp-split GEMV serves the M=1 decode row only"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (1, N)"); |
| TORCH_CHECK(warps == 2 || warps == 4 || warps == 8, |
| "warps must be 2, 4 or 8"); |
| TORCH_CHECK(stages == 3 || stages == 4 || stages == 6, |
| "stages must be 3, 4 or 6"); |
| TORCH_CHECK(shape.n % 8 == 0, "N must be a multiple of 8"); |
| TORCH_CHECK(shape.k % 64 == 0 && (shape.k / 64) % warps == 0, |
| "K must be a multiple of 64*warps"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| auto const* props = current_device_properties(a_packed); |
| TORCH_CHECK(props->major == 12 && props->minor == 0, |
| "the warp-split GEMV is an SM120 kernel; got SM", |
| props->major, props->minor); |
| #if defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build"); |
| #else |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| const int rc = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_bf16out( |
| a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), |
| checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), |
| static_cast<int>(warps), static_cast<int>(stages), stream); |
| TORCH_CHECK(rc == 0, "fp4_w4a4_gemv_warpsplit_bf16 failed with rc=", rc); |
| #endif |
| #else |
| TORCH_CHECK(false, "fp4-gemm was not built with CUDA support"); |
| #endif |
| } |
|
|
| void fp4_w4a16_linear_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor& out, |
| double alpha, |
| int64_t variant) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (M, N)"); |
| TORCH_CHECK(variant >= -1 && variant <= 2, |
| "variant must be -1(auto), 0(default), 1(widen), or 2(pingpong)"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| auto const* props = current_device_properties(a_packed); |
| TORCH_CHECK((props->major == 11 && props->minor == 0) || |
| (props->major == 12 && props->minor == 0), |
| "nvfp4_gemm_bf16 requires SM110 or SM120; got SM", |
| props->major, props->minor); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| if (props->major == 11) { |
| variant = select_sm110_variant(shape, variant); |
| TORCH_CHECK(flash_rt::hub::sm110_gemm_dispatch != nullptr, |
| "SM110 FP4 GEMM source is not present in this build"); |
| flash_rt::hub::sm110_gemm_dispatch( |
| a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(), |
| static_cast<float>(alpha), variant, stream); |
| } else { |
| #if defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build"); |
| #else |
| if (variant == 1) { |
| flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out_widen( |
| a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream); |
| } else if (variant == 2) { |
| flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out_pingpong( |
| a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream); |
| } else { |
| flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out( |
| a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream); |
| } |
| #endif |
| } |
| #endif |
| } |
|
|
| void nvfp4_gemm_bias_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& bias, |
| torch::Tensor& out) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(bias, "bias"); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n, |
| "bias must have shape (N,)"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (M, N)"); |
| check_same_device(a_packed, bias, "a_packed", "bias"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| auto const* props = current_device_properties(a_packed); |
| TORCH_CHECK(props->major == 11 && props->minor == 0, |
| "nvfp4_gemm_bias_bf16 currently requires SM110; got SM", |
| props->major, props->minor); |
| TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_dispatch != nullptr, |
| "SM110 fused-bias FP4 GEMM source is not present in this build"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| const int rc = flash_rt::hub::sm110_gemm_bias_dispatch( |
| a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), |
| checked_int(shape.n, "N"), checked_int(shape.k, "K"), stream); |
| TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_bf16 failed with rc=", rc); |
| #endif |
| } |
|
|
| void nvfp4_gemm_bias_residual_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& bias, |
| torch::Tensor const& residual, |
| torch::Tensor& out) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(bias, "bias"); |
| check_bf16_cuda(residual, "residual"); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n, |
| "bias must have shape (N,)"); |
| TORCH_CHECK(residual.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "residual must have shape (M, N)"); |
| TORCH_CHECK(out.sizes() == residual.sizes(), "out must match residual"); |
| check_same_device(a_packed, bias, "a_packed", "bias"); |
| check_same_device(a_packed, residual, "a_packed", "residual"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| auto const* props = current_device_properties(a_packed); |
| TORCH_CHECK(props->major == 11 && props->minor == 0, |
| "nvfp4_gemm_bias_residual_bf16 currently requires SM110; got SM", |
| props->major, props->minor); |
| TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_residual_dispatch != nullptr, |
| "SM110 bias-residual FP4 GEMM source is not present in this build"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| const int rc = flash_rt::hub::sm110_gemm_bias_residual_dispatch( |
| a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), residual.data_ptr(), out.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), stream); |
| TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_residual_bf16 failed with rc=", rc); |
| #endif |
| } |
|
|
| void nvfp4_gemm_residual_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& residual, |
| torch::Tensor& out, |
| double alpha) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(residual, "residual"); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(residual.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "residual must have shape (M, N)"); |
| TORCH_CHECK(out.sizes() == residual.sizes(), "out must match residual"); |
| check_same_device(a_packed, residual, "a_packed", "residual"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| require_sm120(a_packed, "nvfp4_gemm_residual_bf16"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| flash_rt::gemm::fp4_w4a16_gemm_residual_sm120_bf16out( |
| a_packed.data_ptr(), b_packed.data_ptr(), residual.data_ptr(), |
| out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(), |
| static_cast<float>(alpha), stream); |
| #endif |
| #endif |
| } |
|
|
| void nvfp4_gemm_bias_gelu_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& bias, |
| torch::Tensor& out, |
| double alpha) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(bias, "bias"); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n, |
| "bias must have shape (N,)"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (M, N)"); |
| check_same_device(a_packed, bias, "a_packed", "bias"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| require_sm120(a_packed, "nvfp4_gemm_bias_gelu_bf16"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_bf16out_sm120( |
| a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), |
| checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| static_cast<float>(alpha), stream); |
| #endif |
| #endif |
| } |
|
|
| void nvfp4_gemm_bias_gelu_nvfp4( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& bias, |
| torch::Tensor& out_packed, |
| torch::Tensor& out_sfa, |
| double alpha) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(bias, "bias"); |
| check_uint8_cuda(out_packed, "out_packed"); |
| check_uint8_cuda(out_sfa, "out_sfa"); |
| TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n, |
| "bias must have shape (N,)"); |
| TORCH_CHECK(shape.n % 2 == 0 && |
| out_packed.sizes() == |
| torch::IntArrayRef({shape.m, shape.n / 2}), |
| "out_packed must have shape (M, N / 2)"); |
| TORCH_CHECK(out_sfa.numel() >= swizzled_bytes(shape.m, shape.n), |
| "out_sfa is too small for output scale layout"); |
| check_same_device(a_packed, bias, "a_packed", "bias"); |
| check_same_device(a_packed, out_packed, "a_packed", "out_packed"); |
| check_same_device(a_packed, out_sfa, "a_packed", "out_sfa"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| auto const* props = current_device_properties(a_packed); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| if (props->major == 11 && props->minor == 0) { |
| TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch != nullptr, |
| "SM110 bias-GELU-FP4 GEMM source is not present in this build"); |
| const int rc = flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch( |
| a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), stream); |
| TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_gelu_nvfp4 failed with rc=", rc); |
| return; |
| } |
| require_sm120(a_packed, "nvfp4_gemm_bias_gelu_nvfp4"); |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_fp4out_sm120( |
| a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), |
| checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), static_cast<float>(alpha), stream); |
| #endif |
| #endif |
| } |
|
|
| void nvfp4_gemm_streamk_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor& out, |
| double alpha) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (M, N)"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| require_sm120(a_packed, "nvfp4_gemm_streamk_bf16"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bf16out_sm120( |
| a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), |
| out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), |
| checked_int(shape.k, "K"), static_cast<float>(alpha), stream); |
| #endif |
| #endif |
| } |
|
|
| void nvfp4_gemm_streamk_bias_bf16( |
| torch::Tensor const& a_packed, |
| torch::Tensor const& b_packed, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb, |
| torch::Tensor const& bias, |
| torch::Tensor& out, |
| double alpha) { |
| auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb); |
| check_bf16_cuda(bias, "bias"); |
| check_bf16_cuda(out, "out"); |
| TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n, |
| "bias must have shape (N,)"); |
| TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}), |
| "out must have shape (M, N)"); |
| check_same_device(a_packed, bias, "a_packed", "bias"); |
| check_same_device(a_packed, out, "a_packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(a_packed.device()); |
| require_sm120(a_packed, "nvfp4_gemm_streamk_bias_bf16"); |
| auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); |
| #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) |
| flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120( |
| a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), |
| bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), |
| checked_int(shape.n, "N"), checked_int(shape.k, "K"), |
| static_cast<float>(alpha), stream); |
| #endif |
| #endif |
| } |
|
|
| void quantize_fp4_sfa_fp16( |
| torch::Tensor const& x, |
| torch::Tensor& packed, |
| torch::Tensor& sfa, |
| bool is_sfb) { |
| check_fp16_cuda(x, "x"); |
| check_uint8_cuda(packed, "packed"); |
| check_uint8_cuda(sfa, "sfa"); |
| TORCH_CHECK(x.dim() == 2, "x must have shape (rows, dim)"); |
| const int64_t rows = x.size(0); |
| const int64_t dim = x.size(1); |
| TORCH_CHECK(dim % 16 == 0, "x.shape[1] must be divisible by 16"); |
| TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}), |
| "packed must have shape (rows, dim / 2)"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim), |
| "sfa is too small for CUTLASS SFA/SFB layout"); |
| check_same_device(x, packed, "x", "packed"); |
| check_same_device(x, sfa, "x", "sfa"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(x.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream(); |
| const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_fp16( |
| x.data_ptr(), packed.data_ptr(), sfa.data_ptr(), |
| checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream); |
| TORCH_CHECK(rc == 0, "quantize_fp4_dynamic_sfa_fp16 failed with rc=", rc); |
| #endif |
| } |
|
|
| void quantize_fp4_sfa_bf16( |
| torch::Tensor const& x, |
| torch::Tensor& packed, |
| torch::Tensor& sfa, |
| bool is_sfb) { |
| check_bf16_cuda(x, "x"); |
| check_uint8_cuda(packed, "packed"); |
| check_uint8_cuda(sfa, "sfa"); |
| TORCH_CHECK(x.dim() == 2, "x must have shape (rows, dim)"); |
| const int64_t rows = x.size(0); |
| const int64_t dim = x.size(1); |
| TORCH_CHECK(dim % 16 == 0, "x.shape[1] must be divisible by 16"); |
| TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}), |
| "packed must have shape (rows, dim / 2)"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim), |
| "sfa is too small for CUTLASS SFA/SFB layout"); |
| check_same_device(x, packed, "x", "packed"); |
| check_same_device(x, sfa, "x", "sfa"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(x.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream(); |
| auto const* props = current_device_properties(x); |
| int rc = 0; |
| if (props->major == 11 && props->minor == 0) { |
| TORCH_CHECK(flash_rt::hub::sm110_quantize_bf16_dispatch != nullptr, |
| "SM110 vectorized BF16 FP4 quantizer is not present in this build"); |
| rc = flash_rt::hub::sm110_quantize_bf16_dispatch( |
| x.data_ptr(), packed.data_ptr(), sfa.data_ptr(), |
| checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream); |
| } else { |
| rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16( |
| x.data_ptr(), packed.data_ptr(), sfa.data_ptr(), |
| checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream); |
| } |
| TORCH_CHECK(rc == 0, "quantize_fp4_dynamic_sfa_bf16 failed with rc=", rc); |
| #endif |
| } |
|
|
| void dequantize_fp4_sfa_fp16( |
| torch::Tensor const& packed, |
| torch::Tensor const& sfa, |
| torch::Tensor& out, |
| bool is_sfb) { |
| check_uint8_cuda(packed, "packed"); |
| check_uint8_cuda(sfa, "sfa"); |
| check_fp16_cuda(out, "out"); |
| TORCH_CHECK(out.dim() == 2, "out must have shape (rows, dim)"); |
| const int64_t rows = out.size(0); |
| const int64_t dim = out.size(1); |
| TORCH_CHECK(dim % 16 == 0, "out.shape[1] must be divisible by 16"); |
| TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}), |
| "packed must have shape (rows, dim / 2)"); |
| TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim), |
| "sfa is too small for CUTLASS SFA layout"); |
| check_same_device(packed, sfa, "packed", "sfa"); |
| check_same_device(packed, out, "packed", "out"); |
| #if defined(CUDA_KERNEL) |
| at::cuda::CUDAGuard device_guard(packed.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(packed.get_device()).stream(); |
| flash_rt::fused_fp4::dequantize_fp4_sfa_fp16( |
| reinterpret_cast<const uint8_t*>(packed.data_ptr()), |
| reinterpret_cast<const uint8_t*>(sfa.data_ptr()), |
| reinterpret_cast<__half*>(out.data_ptr()), |
| checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream); |
| #endif |
| } |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def("nvfp4_gemm_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int variant=-1) -> ()"); |
| ops.def("fp4_w4a16_linear_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int variant=-1) -> ()"); |
| ops.def("fp4_w4a4_gemv_warpsplit_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int warps=4, int stages=4) -> ()"); |
| ops.def("nvfp4_gemm_bias_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out) -> ()"); |
| ops.def("nvfp4_gemm_bias_residual_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor residual, Tensor! out) -> ()"); |
| ops.def("nvfp4_gemm_residual_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor residual, Tensor! out, float alpha=1.0) -> ()"); |
| ops.def("nvfp4_gemm_bias_gelu_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out, float alpha=1.0) -> ()"); |
| ops.def("nvfp4_gemm_bias_gelu_nvfp4(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out_packed, Tensor! out_sfa, float alpha=1.0) -> ()"); |
| ops.def("nvfp4_gemm_streamk_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0) -> ()"); |
| ops.def("nvfp4_gemm_streamk_bias_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out, float alpha=1.0) -> ()"); |
| ops.def("quantize_fp4_sfa_fp16(Tensor x, Tensor! packed, Tensor! sfa, bool is_sfb=False) -> ()"); |
| ops.def("quantize_fp4_sfa_bf16(Tensor x, Tensor! packed, Tensor! sfa, bool is_sfb=False) -> ()"); |
| ops.def("dequantize_fp4_sfa_fp16(Tensor packed, Tensor sfa, Tensor! out, bool is_sfb=False) -> ()"); |
| #if defined(CUDA_KERNEL) |
| ops.impl("nvfp4_gemm_bf16", torch::kCUDA, &fp4_w4a16_linear_bf16); |
| ops.impl("fp4_w4a16_linear_bf16", torch::kCUDA, &fp4_w4a16_linear_bf16); |
| ops.impl("fp4_w4a4_gemv_warpsplit_bf16", torch::kCUDA, &fp4_w4a4_gemv_warpsplit_bf16); |
| ops.impl("nvfp4_gemm_bias_bf16", torch::kCUDA, &nvfp4_gemm_bias_bf16); |
| ops.impl("nvfp4_gemm_bias_residual_bf16", torch::kCUDA, &nvfp4_gemm_bias_residual_bf16); |
| ops.impl("nvfp4_gemm_residual_bf16", torch::kCUDA, &nvfp4_gemm_residual_bf16); |
| ops.impl("nvfp4_gemm_bias_gelu_bf16", torch::kCUDA, &nvfp4_gemm_bias_gelu_bf16); |
| ops.impl("nvfp4_gemm_bias_gelu_nvfp4", torch::kCUDA, &nvfp4_gemm_bias_gelu_nvfp4); |
| ops.impl("nvfp4_gemm_streamk_bf16", torch::kCUDA, &nvfp4_gemm_streamk_bf16); |
| ops.impl("nvfp4_gemm_streamk_bias_bf16", torch::kCUDA, &nvfp4_gemm_streamk_bias_bf16); |
| ops.impl("quantize_fp4_sfa_fp16", torch::kCUDA, &quantize_fp4_sfa_fp16); |
| ops.impl("quantize_fp4_sfa_bf16", torch::kCUDA, &quantize_fp4_sfa_bf16); |
| ops.impl("dequantize_fp4_sfa_fp16", torch::kCUDA, &dequantize_fp4_sfa_fp16); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
