Skip to content

GitLab

"vllm/vscode:/vscode.git/clone" did not exist on "ecb4f822091a64b5084b3a4aff326906487a363f"
Unverified Commit 7c080dd3 authored by mikaylagawarecki's avatar mikaylagawarecki Committed by GitHub
Browse files

[4/n] Migrate FP4/W4A8 CUTLASS kernels to torch stable ABI (#37503) 


Signed-off-by: default avatarMikayla Gawarecki <mikaylagawarecki@gmail.com>
parent 0dd25a44
Show whitespace changes
Inline Side-by-side
Showing with 1017 additions and 639 deletions
CMakeLists.txt
View file @ 7c080dd3
...@@ -340,8 +340,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") ...@@ -340,8 +340,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_EXT_SRC list(APPEND VLLM_EXT_SRC
"csrc/quantization/awq/gemm_kernels.cu" "csrc/quantization/awq/gemm_kernels.cu"
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
"csrc/cutlass_extensions/common.cpp") "csrc/cutlass_extensions/common.cpp")
set_gencode_flags_for_srcs( set_gencode_flags_for_srcs(
...@@ -489,59 +487,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") ...@@ -489,59 +487,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
" in CUDA target architectures") " in CUDA target architectures")
endif() endif()
# The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
# CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
"csrc/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
"csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
else()
message(STATUS "Not building NVFP4 as no compatible archs were found.")
# clear FP4_ARCHS
set(FP4_ARCHS)
endif()
# FP4 Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
"csrc/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
else()
message(STATUS "Not building NVFP4 as no compatible archs were found.")
# clear FP4_ARCHS
set(FP4_ARCHS)
endif()
# CUTLASS MLA Archs and flags # CUTLASS MLA Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}") cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
...@@ -681,34 +626,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") ...@@ -681,34 +626,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif() endif()
endif() endif()
# Only build W4A8 kernels if we are building for something compatible with sm90a
cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
set(SRCS
"csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
"csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
"csrc/quantization/cutlass_w4a8/w4a8_utils.cu"
)
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${W4A8_ARCHS}")
list(APPEND VLLM_EXT_SRC "${SRCS}")
message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
AND W4A8_ARCHS)
message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
"not >= 12.0, we recommend upgrading to CUDA 12.0 or "
"later if you intend on running w4a16 quantized models on "
"Hopper.")
else()
message(STATUS "Not building W4A8 kernels as no compatible archs "
"found in CUDA target architectures")
endif()
endif()
# Hadacore kernels # Hadacore kernels
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
...@@ -760,7 +677,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") ...@@ -760,7 +677,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
set(VLLM_STABLE_EXT_SRC set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp" "csrc/libtorch_stable/torch_bindings.cpp"
"csrc/cutlass_extensions/common.cpp" "csrc/cutlass_extensions/common.cpp"
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu") "csrc/cuda_utils_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA") if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_STABLE_EXT_SRC list(APPEND VLLM_STABLE_EXT_SRC
...@@ -978,6 +898,96 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") ...@@ -978,6 +898,96 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif() endif()
endif() endif()
#
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
#
# The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
# CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
else()
message(STATUS "Not building NVFP4 as no compatible archs were found.")
# clear FP4_ARCHS
set(FP4_ARCHS)
endif()
# FP4 Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
else()
message(STATUS "Not building NVFP4 as no compatible archs were found.")
# clear FP4_ARCHS
set(FP4_ARCHS)
endif()
#
# W4A8 kernels (moved from _C to _C_stable_libtorch)
#
# Only build W4A8 kernels if we are building for something compatible with sm90a
cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
set(SRCS
"csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
"csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
"csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu"
)
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${W4A8_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
AND W4A8_ARCHS)
message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
"not >= 12.0, we recommend upgrading to CUDA 12.0 or "
"later if you intend on running w4a16 quantized models on "
"Hopper.")
else()
message(STATUS "Not building W4A8 kernels as no compatible archs "
"found in CUDA target architectures")
endif()
endif()
message(STATUS "Enabling C_stable extension.") message(STATUS "Enabling C_stable extension.")
define_extension_target( define_extension_target(
_C_stable_libtorch _C_stable_libtorch
......
csrc/cuda_vec_utils.cuh
View file @ 7c080dd3
...@@ -3,8 +3,8 @@ ...@@ -3,8 +3,8 @@
#pragma once #pragma once
#include <c10/util/BFloat16.h> #include <torch/headeronly/util/BFloat16.h>
#include <c10/util/Half.h> #include <torch/headeronly/util/Half.h>
#include <cassert> #include <cassert>
#ifdef USE_ROCM #ifdef USE_ROCM
......
csrc/cutlass_extensions/cute_utils.cuh
View file @ 7c080dd3
#pragma once #pragma once
#include <cute/tensor.hpp> #include <cute/tensor.hpp>
#include <torch/all.h>
namespace cute { namespace cute {
//////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////
......
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
View file @ 7c080dd3
...@@ -189,9 +189,9 @@ struct Sm90RowOrScalarBroadcastArray { ...@@ -189,9 +189,9 @@ struct Sm90RowOrScalarBroadcastArray {
} }
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
Tensor tGS_gRow_flt = filter_zeros(tGS_gRow); cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
Tensor tGS_sRow_flt = filter_zeros(tGS_sRow); cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride())); cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
for (int i = 0; i < size(tGS_gRow_flt); ++i) { for (int i = 0; i < size(tGS_gRow_flt); ++i) {
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) { if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
...@@ -211,8 +211,8 @@ struct Sm90RowOrScalarBroadcastArray { ...@@ -211,8 +211,8 @@ struct Sm90RowOrScalarBroadcastArray {
begin_loop(int epi_m, int epi_n) { begin_loop(int epi_m, int epi_n) {
if (epi_m == 0) { // Assumes M-major subtile loop if (epi_m == 0) { // Assumes M-major subtile loop
if (!params.row_broadcast) return; // Do not issue LDS when row is scalar if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n)); cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
Tensor tSR_rRow_flt = filter_zeros(tSR_rRow); cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
copy(tSR_sRow_flt, tSR_rRow_flt); copy(tSR_sRow_flt, tSR_rRow_flt);
} }
} }
...@@ -241,9 +241,9 @@ struct Sm90RowOrScalarBroadcastArray { ...@@ -241,9 +241,9 @@ struct Sm90RowOrScalarBroadcastArray {
auto [m, n, k, l] = args.tile_coord_mnkl; auto [m, n, k, l] = args.tile_coord_mnkl;
using ThreadCount = decltype(size(args.tiled_copy)); using ThreadCount = decltype(size(args.tiled_copy));
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow); cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N) cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem), cute::Tensor sRow = make_tensor(make_smem_ptr(smem),
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N) make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
//// G2S: Gmem to Smem //// G2S: Gmem to Smem
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{}, auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
...@@ -251,16 +251,16 @@ struct Sm90RowOrScalarBroadcastArray { ...@@ -251,16 +251,16 @@ struct Sm90RowOrScalarBroadcastArray {
Stride<_0, _1>>{}, Stride<_0, _1>>{},
Layout<_1>{}); Layout<_1>{});
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx); auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
Tensor tGS_gRow = thr_g2s.partition_S(gRow); cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
Tensor tGS_sRow = thr_g2s.partition_D(sRow); cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
//// G2S: Coord //// G2S: Coord
auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}))); auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
Tensor tGS_cRow = thr_g2s.partition_S(cRow); cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
//// S2R: Smem to Reg //// S2R: Smem to Reg
Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx); cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N) cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>( return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
tGS_gRow, tGS_gRow,
...@@ -389,7 +389,7 @@ struct Sm90ColOrScalarBroadcastArray { ...@@ -389,7 +389,7 @@ struct Sm90ColOrScalarBroadcastArray {
CUTLASS_DEVICE void CUTLASS_DEVICE void
begin() { begin() {
Tensor pred = make_tensor<bool>(shape(tCgCol)); cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
CUTLASS_PRAGMA_UNROLL CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) { for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol(i)) < m; pred(i) = get<0>(tCcCol(i)) < m;
...@@ -409,7 +409,7 @@ struct Sm90ColOrScalarBroadcastArray { ...@@ -409,7 +409,7 @@ struct Sm90ColOrScalarBroadcastArray {
CUTLASS_DEVICE Array<Element, FragmentSize> CUTLASS_DEVICE Array<Element, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) { visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
Array<Element, FragmentSize> frg_col; Array<Element, FragmentSize> frg_col;
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n); cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) { for (int i = 0; i < FragmentSize; ++i) {
...@@ -431,16 +431,16 @@ struct Sm90ColOrScalarBroadcastArray { ...@@ -431,16 +431,16 @@ struct Sm90ColOrScalarBroadcastArray {
auto [M, N, K, L] = args.problem_shape_mnkl; auto [M, N, K, L] = args.problem_shape_mnkl;
auto [m, n, k, l] = args.tile_coord_mnkl; auto [m, n, k, l] = args.tile_coord_mnkl;
Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol); cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
// Generate an identity tensor matching the shape of the global tensor and // Generate an identity tensor matching the shape of the global tensor and
// partition the same way, this will be used to generate the predicate // partition the same way, this will be used to generate the predicate
// tensor for loading // tensor for loading
Tensor cCol = make_identity_tensor(mCol.shape()); cute::Tensor cCol = make_identity_tensor(mCol.shape());
Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
return ConsumerStoreCallbacks( return ConsumerStoreCallbacks(
......
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
View file @ 7c080dd3
...@@ -186,9 +186,9 @@ struct Sm90RowOrScalarBroadcast { ...@@ -186,9 +186,9 @@ struct Sm90RowOrScalarBroadcast {
} }
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
Tensor tGS_gRow_flt = filter_zeros(tGS_gRow); cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
Tensor tGS_sRow_flt = filter_zeros(tGS_sRow); cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride())); cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
for (int i = 0; i < size(tGS_gRow_flt); ++i) { for (int i = 0; i < size(tGS_gRow_flt); ++i) {
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) { if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
...@@ -208,8 +208,8 @@ struct Sm90RowOrScalarBroadcast { ...@@ -208,8 +208,8 @@ struct Sm90RowOrScalarBroadcast {
begin_loop(int epi_m, int epi_n) { begin_loop(int epi_m, int epi_n) {
if (epi_m == 0) { // Assumes M-major subtile loop if (epi_m == 0) { // Assumes M-major subtile loop
if (!params.row_broadcast) return; // Do not issue LDS when row is scalar if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n)); cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
Tensor tSR_rRow_flt = filter_zeros(tSR_rRow); cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
copy(tSR_sRow_flt, tSR_rRow_flt); copy(tSR_sRow_flt, tSR_rRow_flt);
} }
} }
...@@ -238,9 +238,9 @@ struct Sm90RowOrScalarBroadcast { ...@@ -238,9 +238,9 @@ struct Sm90RowOrScalarBroadcast {
auto [m, n, k, l] = args.tile_coord_mnkl; auto [m, n, k, l] = args.tile_coord_mnkl;
using ThreadCount = decltype(size(args.tiled_copy)); using ThreadCount = decltype(size(args.tiled_copy));
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow); cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N) cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem), cute::Tensor sRow = make_tensor(make_smem_ptr(smem),
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N) make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
//// G2S: Gmem to Smem //// G2S: Gmem to Smem
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{}, auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
...@@ -248,16 +248,16 @@ struct Sm90RowOrScalarBroadcast { ...@@ -248,16 +248,16 @@ struct Sm90RowOrScalarBroadcast {
Stride<_0, _1>>{}, Stride<_0, _1>>{},
Layout<_1>{}); Layout<_1>{});
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx); auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
Tensor tGS_gRow = thr_g2s.partition_S(gRow); cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
Tensor tGS_sRow = thr_g2s.partition_D(sRow); cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
//// G2S: Coord //// G2S: Coord
auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}))); auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
Tensor tGS_cRow = thr_g2s.partition_S(cRow); cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
//// S2R: Smem to Reg //// S2R: Smem to Reg
Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx); cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N) cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>( return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
tGS_gRow, tGS_gRow,
...@@ -382,7 +382,7 @@ struct Sm90ColOrScalarBroadcast { ...@@ -382,7 +382,7 @@ struct Sm90ColOrScalarBroadcast {
CUTLASS_DEVICE void CUTLASS_DEVICE void
begin() { begin() {
Tensor pred = make_tensor<bool>(shape(tCgCol)); cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
CUTLASS_PRAGMA_UNROLL CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) { for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol(i)) < m; pred(i) = get<0>(tCcCol(i)) < m;
...@@ -402,7 +402,7 @@ struct Sm90ColOrScalarBroadcast { ...@@ -402,7 +402,7 @@ struct Sm90ColOrScalarBroadcast {
CUTLASS_DEVICE Array<Element, FragmentSize> CUTLASS_DEVICE Array<Element, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) { visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
Array<Element, FragmentSize> frg_col; Array<Element, FragmentSize> frg_col;
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n); cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) { for (int i = 0; i < FragmentSize; ++i) {
...@@ -422,16 +422,16 @@ struct Sm90ColOrScalarBroadcast { ...@@ -422,16 +422,16 @@ struct Sm90ColOrScalarBroadcast {
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) { get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
auto [M, N, K, L] = args.problem_shape_mnkl; auto [M, N, K, L] = args.problem_shape_mnkl;
Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol); cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
// Generate an identity tensor matching the shape of the global tensor and // Generate an identity tensor matching the shape of the global tensor and
// partition the same way, this will be used to generate the predicate // partition the same way, this will be used to generate the predicate
// tensor for loading // tensor for loading
Tensor cCol = make_identity_tensor(mCol.shape()); cute::Tensor cCol = make_identity_tensor(mCol.shape());
Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
return ConsumerStoreCallbacks( return ConsumerStoreCallbacks(
......
csrc/cutlass_extensions/torch_utils.hpp
View file @ 7c080dd3
#pragma once #pragma once
#include <torch/all.h> // This header is shared between _C (unstable ABI, used by machete) and
// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION
// is defined only for the stable target, so we switch includes and types
// accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor.
#ifdef TORCH_TARGET_VERSION
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/BFloat16.h>
#include <torch/headeronly/util/Half.h>
#include <torch/headeronly/util/shim_utils.h> // for STD_TORCH_CHECK
using TorchTensor = torch::stable::Tensor;
#define TORCH_UTILS_CHECK STD_TORCH_CHECK
#else
#include <torch/all.h>
using TorchTensor = torch::Tensor;
#define TORCH_UTILS_CHECK TORCH_CHECK
#endif
#include "cute/layout.hpp" #include "cute/layout.hpp"
#include "cutlass/layout/matrix.h" #include "cutlass/layout/matrix.h"
...@@ -55,16 +70,16 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) { ...@@ -55,16 +70,16 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) {
// If `tensor.dim() < rank(Stride{})`, the shape is padded with 1s and the extra // If `tensor.dim() < rank(Stride{})`, the shape is padded with 1s and the extra
// strides are set to be 0 or 1. // strides are set to be 0 or 1.
template <typename Stride> template <typename Stride>
static inline auto make_cute_layout(torch::Tensor const& tensor, static inline auto make_cute_layout(TorchTensor const& tensor,
std::string_view name = "tensor") { std::string_view name = "tensor") {
TORCH_CHECK(tensor.dim() <= rank(Stride{})); TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{}));
auto stride = cute::transform_with_idx( auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele,
Stride{}, [&](auto const& stride_ele, auto const& idx) { auto const& idx) {
using StrideEle = std::decay_t<decltype(stride_ele)>; using StrideEle = std::decay_t<decltype(stride_ele)>;
if (idx < tensor.dim()) { if (idx < tensor.dim()) {
if constexpr (cute::is_static_v<StrideEle>) { if constexpr (cute::is_static_v<StrideEle>) {
TORCH_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ",
name, ".stride(", idx, ") to be ", StrideEle::value); name, ".stride(", idx, ") to be ", StrideEle::value);
return StrideEle{}; return StrideEle{};
} else { } else {
...@@ -97,7 +112,7 @@ static inline auto make_cute_layout(torch::Tensor const& tensor, ...@@ -97,7 +112,7 @@ static inline auto make_cute_layout(torch::Tensor const& tensor,
template <typename Stride> template <typename Stride>
static inline auto maybe_make_cute_layout( static inline auto maybe_make_cute_layout(
std::optional<torch::Tensor> const& tensor, std::optional<TorchTensor> const& tensor,
std::string_view name = "tensor") { std::string_view name = "tensor") {
using Layout = decltype(make_cute_layout<Stride>(*tensor)); using Layout = decltype(make_cute_layout<Stride>(*tensor));
...@@ -121,12 +136,12 @@ template <typename T> ...@@ -121,12 +136,12 @@ template <typename T>
using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type; using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type;
template <> template <>
struct equivalent_cutlass_type<c10::Half> { struct equivalent_cutlass_type<torch::headeronly::Half> {
using type = cutlass::half_t; using type = cutlass::half_t;
}; };
template <> template <>
struct equivalent_cutlass_type<c10::BFloat16> { struct equivalent_cutlass_type<torch::headeronly::BFloat16> {
using type = cutlass::bfloat16_t; using type = cutlass::bfloat16_t;
}; };
...@@ -134,8 +149,8 @@ struct equivalent_cutlass_type<c10::BFloat16> { ...@@ -134,8 +149,8 @@ struct equivalent_cutlass_type<c10::BFloat16> {
// equivalent_scalar_t (basically inverse of equivalent_cutlass_type) // equivalent_scalar_t (basically inverse of equivalent_cutlass_type)
// //
// Return a `c10::CppTypeToScalarType<T>` compatible type, i.e. get the C++ from // Return a `torch::headeronly::CppTypeToScalarType<T>` compatible type, i.e.
// c10 that is equivalent to T, e.g.: `cutlass::half_t -> c10::Half` // get the C++ type equivalent to T, e.g.: `cutlass::half_t -> Half`
template <typename T> template <typename T>
struct equivalent_scalar_type { struct equivalent_scalar_type {
using type = T; using type = T;
...@@ -146,15 +161,15 @@ using equivalent_scalar_type_t = typename equivalent_scalar_type<T>::type; ...@@ -146,15 +161,15 @@ using equivalent_scalar_type_t = typename equivalent_scalar_type<T>::type;
template <> template <>
struct equivalent_scalar_type<cutlass::half_t> { struct equivalent_scalar_type<cutlass::half_t> {
using type = c10::Half; using type = torch::headeronly::Half;
}; };
template <> template <>
struct equivalent_scalar_type<cutlass::bfloat16_t> { struct equivalent_scalar_type<cutlass::bfloat16_t> {
using type = c10::BFloat16; using type = torch::headeronly::BFloat16;
}; };
// get equivalent c10::ScalarType tag from compile time type // get equivalent torch::headeronly::ScalarType tag from compile time type
template <typename T> template <typename T>
static inline constexpr c10::ScalarType equivalent_scalar_type_v = static inline constexpr torch::headeronly::ScalarType equivalent_scalar_type_v =
c10::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value; torch::headeronly::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;
\ No newline at end of file
csrc/libtorch_stable/dispatch_utils.h
View file @ 7c080dd3
...@@ -49,6 +49,15 @@ ...@@ -49,6 +49,15 @@
THO_DISPATCH_SWITCH(TYPE, NAME, \ THO_DISPATCH_SWITCH(TYPE, NAME, \
VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
// Half types dispatch (Half + BFloat16)
#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \
THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
THO_DISPATCH_SWITCH(TYPE, NAME, \
VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
// Boolean dispatch // Boolean dispatch
#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ #define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
if (expr) { \ if (expr) { \
......
csrc/libtorch_stable/ops.h
View file @ 7c080dd3
...@@ -84,4 +84,54 @@ void get_cutlass_batched_moe_mm_data( ...@@ -84,4 +84,54 @@ void get_cutlass_batched_moe_mm_data(
const torch::stable::Tensor& expert_num_tokens, const torch::stable::Tensor& expert_num_tokens,
const int64_t num_local_experts, const int64_t padded_m, const int64_t n, const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
const int64_t k); const int64_t k);
// FP4/NVFP4 ops
bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
torch::stable::Tensor const& A,
torch::stable::Tensor const& B,
torch::stable::Tensor const& A_sf,
torch::stable::Tensor const& B_sf,
torch::stable::Tensor const& alpha);
void cutlass_fp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& alphas,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets);
std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout);
void scaled_fp4_quant_out(torch::stable::Tensor const& input,
torch::stable::Tensor const& input_scale,
bool is_sf_swizzled_layout,
torch::stable::Tensor& output,
torch::stable::Tensor& output_scale);
void scaled_fp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_global_scale,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts);
void silu_and_mul_scaled_fp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_global_scale,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts);
void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
torch::stable::Tensor& output_block_scale,
torch::stable::Tensor& input,
torch::stable::Tensor& input_global_scale);
#endif #endif
csrc/quantization/cutlass_w4a8/get_group_starts.cuh → csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh
View file @ 7c080dd3
...@@ -2,10 +2,9 @@ ...@@ -2,10 +2,9 @@
#pragma once #pragma once
#include <cuda.h> #include <cuda.h>
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include <c10/cuda/CUDAStream.h> #include "libtorch_stable/torch_utils.h"
#include "core/scalar_type.hpp"
#include "cutlass/bfloat16.h" #include "cutlass/bfloat16.h"
#include "cutlass/float8.h" #include "cutlass/float8.h"
...@@ -41,7 +40,7 @@ __global__ void get_group_gemm_starts( ...@@ -41,7 +40,7 @@ __global__ void get_group_gemm_starts(
} }
#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \ #define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
else if (out_tensors.dtype() == TENSOR_C_TYPE) { \ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \ get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \
cutlass::Array<cutlass::float_e4m3_t, 8>> \ cutlass::Array<cutlass::float_e4m3_t, 8>> \
<<<1, num_experts, 0, stream>>>( \ <<<1, num_experts, 0, stream>>>( \
...@@ -66,23 +65,34 @@ __global__ void get_group_gemm_starts( ...@@ -66,23 +65,34 @@ __global__ void get_group_gemm_starts(
namespace { namespace {
void run_get_group_gemm_starts( void run_get_group_gemm_starts(
torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs, torch::stable::Tensor const& expert_offsets, torch::stable::Tensor& a_ptrs,
torch::Tensor& b_ptrs, torch::Tensor& out_ptrs, torch::stable::Tensor& b_ptrs, torch::stable::Tensor& out_ptrs,
torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs, torch::stable::Tensor& a_scales_ptrs, torch::stable::Tensor& b_scales_ptrs,
torch::Tensor& b_group_scales_ptrs, torch::Tensor const& a_tensors, torch::stable::Tensor& b_group_scales_ptrs,
torch::Tensor const& b_tensors, torch::Tensor& out_tensors, torch::stable::Tensor const& a_tensors,
torch::Tensor const& a_scales, torch::Tensor const& b_scales, torch::stable::Tensor const& b_tensors, torch::stable::Tensor& out_tensors,
torch::Tensor const& b_group_scales, const int64_t b_group_size) { torch::stable::Tensor const& a_scales,
TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn); torch::stable::Tensor const& b_scales,
TORCH_CHECK(b_tensors.dtype() == torch::kInt32); // int4 8x packed into int32 torch::stable::Tensor const& b_group_scales, const int64_t b_group_size) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32); STD_TORCH_CHECK(a_tensors.scalar_type() ==
TORCH_CHECK(b_scales.dtype() == torch::kFloat32); torch::headeronly::ScalarType::Float8_e4m3fn);
TORCH_CHECK(b_group_scales.dtype() == STD_TORCH_CHECK(
torch::kFloat8_e4m3fn); // the underlying torch type is e4m3 b_tensors.scalar_type() ==
TORCH_CHECK(out_tensors.dtype() == torch::headeronly::ScalarType::Int); // int4 8x packed into int32
torch::kBFloat16); // only support bf16 for now STD_TORCH_CHECK(a_scales.scalar_type() ==
torch::headeronly::ScalarType::Float);
STD_TORCH_CHECK(b_scales.scalar_type() ==
torch::headeronly::ScalarType::Float);
STD_TORCH_CHECK(
b_group_scales.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn); // the underlying torch
// type is e4m3
STD_TORCH_CHECK(
out_tensors.scalar_type() ==
torch::headeronly::ScalarType::BFloat16); // only support bf16 for now
// expect int64_t to avoid overflow during offset calculations // expect int64_t to avoid overflow during offset calculations
TORCH_CHECK(expert_offsets.dtype() == torch::kInt64); STD_TORCH_CHECK(expert_offsets.scalar_type() ==
torch::headeronly::ScalarType::Long);
int num_experts = static_cast<int>(expert_offsets.size(0)); int num_experts = static_cast<int>(expert_offsets.size(0));
// logical k, n // logical k, n
...@@ -90,14 +100,15 @@ void run_get_group_gemm_starts( ...@@ -90,14 +100,15 @@ void run_get_group_gemm_starts(
int64_t k = a_tensors.size(1); int64_t k = a_tensors.size(1);
int64_t scale_k = cutlass::ceil_div(k, b_group_size); int64_t scale_k = cutlass::ceil_div(k, b_group_size);
auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index());
if (false) { if (false) {
} }
__CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t) __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::BFloat16,
__CALL_GET_STARTS_KERNEL(torch::kFloat16, half) cutlass::bfloat16_t)
__CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::Half, half)
else { else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
} }
} }
......
csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu → csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu
View file @ 7c080dd3
...@@ -14,13 +14,12 @@ ...@@ -14,13 +14,12 @@
#include "cutlass/util/mixed_dtype_utils.hpp" #include "cutlass/util/mixed_dtype_utils.hpp"
// vllm includes // vllm includes
#include <ATen/cuda/CUDAContext.h> #include <torch/csrc/stable/library.h>
#include <c10/cuda/CUDAGuard.h> #include <torch/csrc/stable/tensor.h>
#include <torch/all.h> #include "libtorch_stable/torch_utils.h"
#include "cutlass_extensions/torch_utils.hpp" #include "cutlass_extensions/torch_utils.hpp"
#include "cutlass_extensions/common.hpp" #include "cutlass_extensions/common.hpp"
#include "core/registration.h"
#include "get_group_starts.cuh" #include "get_group_starts.cuh"
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
#include "w4a8_utils.cuh" #include "w4a8_utils.cuh"
...@@ -168,31 +167,40 @@ struct W4A8GroupedGemmKernel { ...@@ -168,31 +167,40 @@ struct W4A8GroupedGemmKernel {
static_assert(sizeof(LayoutB_Reordered) % sizeof(int32_t) == 0, static_assert(sizeof(LayoutB_Reordered) % sizeof(int32_t) == 0,
"LayoutB_Reordered size must be divisible by 4 bytes"); "LayoutB_Reordered size must be divisible by 4 bytes");
static void grouped_mm( static void grouped_mm(torch::stable::Tensor& out_tensors,
torch::Tensor& out_tensors, const torch::Tensor& a_tensors, const torch::stable::Tensor& a_tensors,
const torch::Tensor& b_tensors, const torch::Tensor& a_scales, const torch::stable::Tensor& b_tensors,
const torch::Tensor& b_scales, const torch::Tensor& b_group_scales, const torch::stable::Tensor& a_scales,
const int64_t b_group_size, const torch::Tensor& expert_offsets, const torch::stable::Tensor& b_scales,
const torch::Tensor& problem_sizes_torch, const torch::Tensor& a_strides, const torch::stable::Tensor& b_group_scales,
const torch::Tensor& b_strides, const torch::Tensor& c_strides, const int64_t b_group_size,
const torch::Tensor& group_scale_strides) { const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& problem_sizes_torch,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides,
const torch::stable::Tensor& group_scale_strides) {
auto device = a_tensors.device(); auto device = a_tensors.device();
auto device_id = device.index(); auto device_id = device.index();
const at::cuda::OptionalCUDAGuard device_guard(device); const torch::stable::accelerator::DeviceGuard device_guard(device_id);
auto stream = at::cuda::getCurrentCUDAStream(device_id); auto stream = get_current_cuda_stream(device_id);
int num_experts = static_cast<int>(expert_offsets.size(0)); int num_experts = static_cast<int>(expert_offsets.size(0));
int n = static_cast<int>(b_tensors.size(1)); int n = static_cast<int>(b_tensors.size(1));
int k = static_cast<int>(b_tensors.size(2)) * PackFactor; int k = static_cast<int>(b_tensors.size(2)) * PackFactor;
auto options_int = torch::stable::Tensor a_ptrs = torch::stable::empty(
torch::TensorOptions().dtype(torch::kInt64).device(device); num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
torch::Tensor a_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor b_ptrs = torch::stable::empty(
torch::Tensor b_ptrs = torch::empty(num_experts, options_int); num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
torch::Tensor out_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor out_ptrs = torch::stable::empty(
torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int); num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor a_scales_ptrs = torch::stable::empty(
torch::Tensor b_group_scales_ptrs = torch::empty(num_experts, options_int); num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
torch::stable::Tensor b_scales_ptrs = torch::stable::empty(
num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
torch::stable::Tensor b_group_scales_ptrs = torch::stable::empty(
num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
// get the correct offsets to pass to gemm // get the correct offsets to pass to gemm
run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs, run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
...@@ -247,9 +255,9 @@ struct W4A8GroupedGemmKernel { ...@@ -247,9 +255,9 @@ struct W4A8GroupedGemmKernel {
// Allocate workspace // Allocate workspace
size_t workspace_size = GemmShuffled::get_workspace_size(arguments); size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
torch::Tensor workspace = torch::stable::Tensor workspace = torch::stable::empty(
torch::empty(workspace_size, workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
torch::TensorOptions().dtype(torch::kU8).device(device)); device);
// Run GEMM // Run GEMM
GemmShuffled gemm; GemmShuffled gemm;
...@@ -294,14 +302,20 @@ using Kernel_256x128_2x1x1_Coop = ...@@ -294,14 +302,20 @@ using Kernel_256x128_2x1x1_Coop =
using Kernel_128x256_2x1x1_Coop = using Kernel_128x256_2x1x1_Coop =
W4A8GroupedGemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>, Coop, CoopEpi>; W4A8GroupedGemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>, Coop, CoopEpi>;
void mm_dispatch( void mm_dispatch(torch::stable::Tensor& out_tensors,
torch::Tensor& out_tensors, const torch::Tensor& a_tensors, const torch::stable::Tensor& a_tensors,
const torch::Tensor& b_tensors, const torch::Tensor& a_scales, const torch::stable::Tensor& b_tensors,
const torch::Tensor& b_scales, const torch::Tensor& b_group_scales, const torch::stable::Tensor& a_scales,
const int64_t b_group_size, const torch::Tensor& expert_offsets, const torch::stable::Tensor& b_scales,
const torch::Tensor& problem_sizes, const torch::Tensor& a_strides, const torch::stable::Tensor& b_group_scales,
const torch::Tensor& b_strides, const torch::Tensor& c_strides, const int64_t b_group_size,
const torch::Tensor& group_scale_strides, const std::string& schedule) { const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides,
const torch::stable::Tensor& group_scale_strides,
const std::string& schedule) {
if (schedule == "Kernel_128x16_1x1x1_Coop") { if (schedule == "Kernel_128x16_1x1x1_Coop") {
Kernel_128x16_1x1x1_Coop::grouped_mm( Kernel_128x16_1x1x1_Coop::grouped_mm(
out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales, out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
...@@ -358,18 +372,23 @@ void mm_dispatch( ...@@ -358,18 +372,23 @@ void mm_dispatch(
b_group_size, expert_offsets, problem_sizes, a_strides, b_strides, b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
c_strides, group_scale_strides); c_strides, group_scale_strides);
} else { } else {
TORCH_CHECK(false, STD_TORCH_CHECK(false,
"cutlass_w4a8_moe_mm: unknown schedule string: ", schedule); "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
} }
} }
void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors, void mm(torch::stable::Tensor& out_tensors,
const torch::Tensor& b_tensors, const torch::Tensor& a_scales, const torch::stable::Tensor& a_tensors,
const torch::Tensor& b_scales, const torch::Tensor& b_group_scales, const torch::stable::Tensor& b_tensors,
const int64_t b_group_size, const torch::Tensor& expert_offsets, const torch::stable::Tensor& a_scales,
const torch::Tensor& problem_sizes, const torch::Tensor& a_strides, const torch::stable::Tensor& b_scales,
const torch::Tensor& b_strides, const torch::Tensor& c_strides, const torch::stable::Tensor& b_group_scales, const int64_t b_group_size,
const torch::Tensor& group_scale_strides, const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides,
const torch::stable::Tensor& group_scale_strides,
std::optional<std::string> maybe_schedule) { std::optional<std::string> maybe_schedule) {
// user has specified a schedule // user has specified a schedule
if (maybe_schedule) { if (maybe_schedule) {
...@@ -406,26 +425,27 @@ void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors, ...@@ -406,26 +425,27 @@ void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
a_strides, b_strides, c_strides, group_scale_strides, schedule); a_strides, b_strides, c_strides, group_scale_strides, schedule);
} }
std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b( std::tuple<torch::stable::Tensor, torch::stable::Tensor>
torch::Tensor const& b_tensors) { encode_and_reorder_int4b(torch::stable::Tensor const& b_tensors) {
TORCH_CHECK(b_tensors.dtype() == torch::kInt32); STD_TORCH_CHECK(b_tensors.scalar_type() ==
TORCH_CHECK(b_tensors.dim() == 3); // (experts, n, k) torch::headeronly::ScalarType::Int);
TORCH_CHECK(b_tensors.is_contiguous()); STD_TORCH_CHECK(b_tensors.dim() == 3); // (experts, n, k)
TORCH_CHECK(b_tensors.is_cuda()); STD_TORCH_CHECK(b_tensors.is_contiguous());
STD_TORCH_CHECK(b_tensors.is_cuda());
int n = static_cast<int>(b_tensors.size(1)); int n = static_cast<int>(b_tensors.size(1));
int k = static_cast<int>(b_tensors.size(2)) * PackFactor; // logical k int k = static_cast<int>(b_tensors.size(2)) * PackFactor; // logical k
// CUTLASS reorder_tensor requires k % 256 == 0 and n % 16 == 0. // CUTLASS reorder_tensor requires k % 256 == 0 and n % 16 == 0.
// These misalignments cause silent OOB unless run under Compute Sanitizer. // These misalignments cause silent OOB unless run under Compute Sanitizer.
TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256"); STD_TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
TORCH_CHECK(n % 16 == 0, "n must be divisible by 16"); STD_TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
// we will store the layout to an int32 tensor; // we will store the layout to an int32 tensor;
// this is the number of elements we need per layout // this is the number of elements we need per layout
constexpr size_t layout_width = sizeof(LayoutB_Reordered) / sizeof(int32_t); constexpr size_t layout_width = sizeof(LayoutB_Reordered) / sizeof(int32_t);
torch::Tensor b_tensors_packed = torch::empty_like(b_tensors); torch::stable::Tensor b_tensors_packed = torch::stable::empty_like(b_tensors);
int num_experts = static_cast<int>(b_tensors.size(0)); int num_experts = static_cast<int>(b_tensors.size(0));
auto b_ptr = static_cast<QuantType const*>(b_tensors.const_data_ptr()); auto b_ptr = static_cast<QuantType const*>(b_tensors.const_data_ptr());
...@@ -435,7 +455,7 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b( ...@@ -435,7 +455,7 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
size_t num_int4_elems = 1ull * num_experts * n * k; size_t num_int4_elems = 1ull * num_experts * n * k;
bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(b_ptr, b_packed_ptr, bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(b_ptr, b_packed_ptr,
num_int4_elems); num_int4_elems);
TORCH_CHECK(ok, "unified_encode_int4b failed"); STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
// construct the layout once; assumes each expert has the same layout // construct the layout once; assumes each expert has the same layout
using LayoutType = LayoutB_Reordered; using LayoutType = LayoutB_Reordered;
...@@ -456,28 +476,28 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b( ...@@ -456,28 +476,28 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
} }
// save the packed layout to torch tensor so we can re-use it // save the packed layout to torch tensor so we can re-use it
auto cpu_opts = torch::stable::Tensor layout_cpu = torch::stable::empty(
torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU); {num_experts, layout_width}, torch::headeronly::ScalarType::Int,
torch::Tensor layout_cpu = std::nullopt, torch::stable::Device(torch::stable::DeviceType::CPU));
torch::empty({num_experts, layout_width}, cpu_opts);
int32_t* layout_data = layout_cpu.data_ptr<int32_t>(); int32_t* layout_data = layout_cpu.mutable_data_ptr<int32_t>();
for (int i = 0; i < num_experts; ++i) { for (int i = 0; i < num_experts; ++i) {
std::memcpy(layout_data + i * layout_width, // dst (int32*) std::memcpy(layout_data + i * layout_width, // dst (int32*)
&layout_B_reordered, // src (LayoutType*) &layout_B_reordered, // src (LayoutType*)
sizeof(LayoutType)); // number of bytes sizeof(LayoutType)); // number of bytes
} }
torch::Tensor packed_layout = torch::stable::Tensor packed_layout =
layout_cpu.to(b_tensors.device(), /*non_blocking=*/false); torch::stable::to(layout_cpu, b_tensors.device(),
/*non_blocking=*/false);
return {b_tensors_packed, packed_layout}; return {b_tensors_packed, packed_layout};
} }
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("cutlass_w4a8_moe_mm", &mm); m.impl("cutlass_w4a8_moe_mm", TORCH_BOX(&mm));
m.impl("cutlass_encode_and_reorder_int4b_grouped", &encode_and_reorder_int4b); m.impl("cutlass_encode_and_reorder_int4b_grouped",
TORCH_BOX(&encode_and_reorder_int4b));
} }
} // namespace vllm::cutlass_w4a8_moe } // namespace vllm::cutlass_w4a8_moe
/////////////////////////////////////////////////////////////////////////////////////////////////
\ No newline at end of file
csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu → csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu
View file @ 7c080dd3
...@@ -3,14 +3,12 @@ ...@@ -3,14 +3,12 @@
// https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu // https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu
// //
#include <ATen/cuda/CUDAContext.h> #include <torch/csrc/stable/library.h>
#include <c10/cuda/CUDAGuard.h> #include <torch/csrc/stable/tensor.h>
#include <torch/all.h> #include "libtorch_stable/torch_utils.h"
#include "cutlass_extensions/torch_utils.hpp" #include "cutlass_extensions/torch_utils.hpp"
#include "w4a8_utils.cuh" #include "w4a8_utils.cuh"
#include "core/registration.h"
#include "cutlass/cutlass.h" #include "cutlass/cutlass.h"
#include <limits> #include <limits>
...@@ -161,31 +159,31 @@ struct W4A8GemmKernel { ...@@ -161,31 +159,31 @@ struct W4A8GemmKernel {
using StrideD = typename GemmKernelShuffled::StrideD; using StrideD = typename GemmKernelShuffled::StrideD;
using StrideS = typename CollectiveMainloopShuffled::StrideScale; using StrideS = typename CollectiveMainloopShuffled::StrideScale;
static torch::Tensor mm(torch::Tensor const& A, static torch::stable::Tensor mm(
torch::Tensor const& B, // already packed torch::stable::Tensor const& A,
torch::Tensor const& group_scales, // already packed torch::stable::Tensor const& B, // already packed
int64_t group_size, torch::stable::Tensor const& group_scales, // already packed
torch::Tensor const& channel_scales, int64_t group_size, torch::stable::Tensor const& channel_scales,
torch::Tensor const& token_scales, torch::stable::Tensor const& token_scales,
std::optional<at::ScalarType> const& maybe_out_type) { std::optional<torch::headeronly::ScalarType> const& maybe_out_type) {
// TODO: param validation // TODO: param validation
int m = A.size(0); int m = A.size(0);
int k = A.size(1); int k = A.size(1);
int n = B.size(1); int n = B.size(1);
// safely cast group_size to int // safely cast group_size to int
TORCH_CHECK(group_size > 0 && group_size <= std::numeric_limits<int>::max(), STD_TORCH_CHECK(
group_size > 0 && group_size <= std::numeric_limits<int>::max(),
"group_size out of supported range for int: ", group_size); "group_size out of supported range for int: ", group_size);
int const group_size_int = static_cast<int>(group_size); int const group_size_int = static_cast<int>(group_size);
// Allocate output // Allocate output
const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); const torch::stable::accelerator::DeviceGuard device_guard(
A.get_device_index());
auto device = A.device(); auto device = A.device();
auto stream = at::cuda::getCurrentCUDAStream(device.index()); auto stream = get_current_cuda_stream(device.index());
torch::Tensor D = torch::stable::Tensor D = torch::stable::empty(
torch::empty({m, n}, torch::TensorOptions() {m, n}, equivalent_scalar_type_v<ElementD>, std::nullopt, device);
.dtype(equivalent_scalar_type_v<ElementD>)
.device(device));
// prepare arg pointers // prepare arg pointers
auto A_ptr = static_cast<MmaType const*>(A.const_data_ptr()); auto A_ptr = static_cast<MmaType const*>(A.const_data_ptr());
auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr()); auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
...@@ -237,9 +235,9 @@ struct W4A8GemmKernel { ...@@ -237,9 +235,9 @@ struct W4A8GemmKernel {
// Workspace // Workspace
size_t workspace_size = GemmShuffled::get_workspace_size(arguments); size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
torch::Tensor workspace = torch::stable::Tensor workspace = torch::stable::empty(
torch::empty(workspace_size, workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
torch::TensorOptions().dtype(torch::kU8).device(device)); device);
// Run GEMM // Run GEMM
GemmShuffled gemm; GemmShuffled gemm;
...@@ -269,13 +267,13 @@ using Kernel_128x64_1x1x1 = W4A8GemmKernel<Shape<_128, _64>, Shape<_1, _1, _1>>; ...@@ -269,13 +267,13 @@ using Kernel_128x64_1x1x1 = W4A8GemmKernel<Shape<_128, _64>, Shape<_1, _1, _1>>;
using Kernel_128x32_1x1x1 = W4A8GemmKernel<Shape<_128, _32>, Shape<_1, _1, _1>>; using Kernel_128x32_1x1x1 = W4A8GemmKernel<Shape<_128, _32>, Shape<_1, _1, _1>>;
using Kernel_128x16_1x1x1 = W4A8GemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>>; using Kernel_128x16_1x1x1 = W4A8GemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>>;
torch::Tensor mm_dispatch(torch::Tensor const& A, torch::stable::Tensor mm_dispatch(
torch::Tensor const& B, // already packed torch::stable::Tensor const& A,
torch::Tensor const& group_scales, // already packed torch::stable::Tensor const& B, // already packed
int64_t group_size, torch::stable::Tensor const& group_scales, // already packed
torch::Tensor const& channel_scales, int64_t group_size, torch::stable::Tensor const& channel_scales,
torch::Tensor const& token_scales, torch::stable::Tensor const& token_scales,
std::optional<at::ScalarType> const& maybe_out_type, std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
const std::string& schedule) { const std::string& schedule) {
if (schedule == "256x128_1x1x1") { if (schedule == "256x128_1x1x1") {
return Kernel_256x128_1x1x1::mm(A, B, group_scales, group_size, return Kernel_256x128_1x1x1::mm(A, B, group_scales, group_size,
...@@ -318,16 +316,17 @@ torch::Tensor mm_dispatch(torch::Tensor const& A, ...@@ -318,16 +316,17 @@ torch::Tensor mm_dispatch(torch::Tensor const& A,
channel_scales, token_scales, channel_scales, token_scales,
maybe_out_type); maybe_out_type);
} }
TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule); STD_TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
return {}; return {};
} }
torch::Tensor mm(torch::Tensor const& A, torch::stable::Tensor mm(
torch::Tensor const& B, // already packed torch::stable::Tensor const& A,
torch::Tensor const& group_scales, // already packed torch::stable::Tensor const& B, // already packed
int64_t group_size, torch::Tensor const& channel_scales, torch::stable::Tensor const& group_scales, // already packed
torch::Tensor const& token_scales, int64_t group_size, torch::stable::Tensor const& channel_scales,
std::optional<at::ScalarType> const& maybe_out_type, torch::stable::Tensor const& token_scales,
std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
std::optional<std::string> maybe_schedule) { std::optional<std::string> maybe_schedule) {
// requested a specific schedule // requested a specific schedule
if (maybe_schedule) { if (maybe_schedule) {
...@@ -378,14 +377,15 @@ torch::Tensor mm(torch::Tensor const& A, ...@@ -378,14 +377,15 @@ torch::Tensor mm(torch::Tensor const& A,
// ---------------------------------------------------------------------------- // ----------------------------------------------------------------------------
// Pre-processing utils // Pre-processing utils
// ---------------------------------------------------------------------------- // ----------------------------------------------------------------------------
torch::Tensor pack_scale_fp8(torch::Tensor const& scales) { torch::stable::Tensor pack_scale_fp8(torch::stable::Tensor const& scales) {
TORCH_CHECK(scales.dtype() == torch::kFloat8_e4m3fn); STD_TORCH_CHECK(scales.scalar_type() ==
TORCH_CHECK(scales.is_contiguous()); torch::headeronly::ScalarType::Float8_e4m3fn);
TORCH_CHECK(scales.is_cuda()); STD_TORCH_CHECK(scales.is_contiguous());
STD_TORCH_CHECK(scales.is_cuda());
auto packed_scales = torch::empty(
{scales.numel() * ScalePackSize}, auto packed_scales =
torch::TensorOptions().dtype(scales.dtype()).device(scales.device())); torch::stable::empty({scales.numel() * ScalePackSize},
scales.scalar_type(), std::nullopt, scales.device());
auto scales_ptr = static_cast<MmaType const*>(scales.const_data_ptr()); auto scales_ptr = static_cast<MmaType const*>(scales.const_data_ptr());
auto packed_scales_ptr = auto packed_scales_ptr =
static_cast<cutlass::Array<ElementScale, ScalePackSize>*>( static_cast<cutlass::Array<ElementScale, ScalePackSize>*>(
...@@ -396,15 +396,16 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) { ...@@ -396,15 +396,16 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
return packed_scales; return packed_scales;
} }
torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) { torch::stable::Tensor encode_and_reorder_int4b(torch::stable::Tensor const& B) {
TORCH_CHECK(B.dtype() == torch::kInt32); STD_TORCH_CHECK(B.scalar_type() == torch::headeronly::ScalarType::Int);
TORCH_CHECK(B.dim() == 2); STD_TORCH_CHECK(B.dim() == 2);
torch::Tensor B_packed = torch::empty_like(B); torch::stable::Tensor B_packed = torch::stable::empty_like(B);
int k = B.size(0) * PackFactor; // logical k int k = B.size(0) * PackFactor; // logical k
int n = B.size(1); int n = B.size(1);
TORCH_CHECK((n * k) % 32 == 0, "need multiples of 32 int4s for 16B chunks"); STD_TORCH_CHECK((n * k) % 32 == 0,
"need multiples of 32 int4s for 16B chunks");
auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr()); auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
auto B_packed_ptr = static_cast<QuantType*>(B_packed.data_ptr()); auto B_packed_ptr = static_cast<QuantType*>(B_packed.data_ptr());
...@@ -415,16 +416,17 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) { ...@@ -415,16 +416,17 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(B_ptr, B_packed_ptr, bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(B_ptr, B_packed_ptr,
n * k); n * k);
TORCH_CHECK(ok, "unified_encode_int4b failed"); STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered); cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
return B_packed; return B_packed;
} }
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("cutlass_w4a8_mm", &mm); m.impl("cutlass_w4a8_mm", TORCH_BOX(&mm));
m.impl("cutlass_pack_scale_fp8", &pack_scale_fp8); m.impl("cutlass_pack_scale_fp8", TORCH_BOX(&pack_scale_fp8));
m.impl("cutlass_encode_and_reorder_int4b", &encode_and_reorder_int4b); m.impl("cutlass_encode_and_reorder_int4b",
TORCH_BOX(&encode_and_reorder_int4b));
} }
} // namespace vllm::cutlass_w4a8 } // namespace vllm::cutlass_w4a8
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu → csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
View file @ 7c080dd3
...@@ -14,16 +14,15 @@ ...@@ -14,16 +14,15 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include <cuda_runtime_api.h> #include <cuda_runtime_api.h>
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp8.h> #include <cuda_fp8.h>
#include "dispatch_utils.h"
#include "cuda_utils.h" #include "cuda_utils.h"
#include "launch_bounds_utils.h" #include "launch_bounds_utils.h"
...@@ -118,16 +117,18 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) ...@@ -118,16 +117,18 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
} // namespace vllm } // namespace vllm
void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d] void silu_and_mul_nvfp4_quant_sm1xxa(
torch::Tensor& output_sf, torch::stable::Tensor& output, // [..., d]
torch::Tensor& input, // [..., 2 * d] torch::stable::Tensor& output_sf,
torch::Tensor& input_sf) { torch::stable::Tensor& input, // [..., 2 * d]
torch::stable::Tensor& input_sf) {
int32_t m = input.size(0); int32_t m = input.size(0);
int32_t n = input.size(1) / 2; int32_t n = input.size(1) / 2;
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16."); STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
TORCH_CHECK(input.scalar_type() == at::ScalarType::Half || STD_TORCH_CHECK(
input.scalar_type() == at::ScalarType::BFloat16, input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Unsupported input data type for quantize_to_fp4."); "Unsupported input data type for quantize_to_fp4.");
int multiProcessorCount = int multiProcessorCount =
...@@ -136,8 +137,9 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d] ...@@ -136,8 +137,9 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr()); auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr()); auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
auto output_ptr = static_cast<int64_t*>(output.data_ptr()); auto output_ptr = static_cast<int64_t*>(output.data_ptr());
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); const torch::stable::accelerator::DeviceGuard device_guard(
auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); input.get_device_index());
auto stream = get_current_cuda_stream(input.get_device_index());
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512)); dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
int const numBlocksPerSM = int const numBlocksPerSM =
vllm_runtime_blocks_per_sm(static_cast<int>(block.x)); vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
...@@ -149,7 +151,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d] ...@@ -149,7 +151,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
dim3 grid(grid_x, grid_y); dim3 grid(grid_x, grid_y);
VLLM_DISPATCH_HALF_TYPES( VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] { input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr()); auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
......
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
View file @ 7c080dd3
...@@ -14,14 +14,12 @@ ...@@ -14,14 +14,12 @@
* limitations under the License. * limitations under the License.
*/ */
#include "core/registration.h" #include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include <torch/all.h>
#include <cutlass/arch/arch.h> #include <cutlass/arch/arch.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include "cutlass_extensions/common.hpp" #include "cutlass_extensions/common.hpp"
#include "cute/tensor.hpp" #include "cute/tensor.hpp"
...@@ -122,7 +120,7 @@ __global__ void __get_group_gemm_starts( ...@@ -122,7 +120,7 @@ __global__ void __get_group_gemm_starts(
#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(ELEMENT_AB_TYPE, SF_TYPE, \ #define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(ELEMENT_AB_TYPE, SF_TYPE, \
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \ TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
LayoutSFB, ScaleConfig) \ LayoutSFB, ScaleConfig) \
else if (out_tensors.dtype() == TENSOR_C_TYPE) { \ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
__get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, \ __get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, \
LayoutSFA, LayoutSFB, ScaleConfig> \ LayoutSFA, LayoutSFB, ScaleConfig> \
<<<1, num_experts, 0, stream>>>( \ <<<1, num_experts, 0, stream>>>( \
...@@ -150,26 +148,36 @@ __global__ void __get_group_gemm_starts( ...@@ -150,26 +148,36 @@ __global__ void __get_group_gemm_starts(
} }
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig> template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
void run_get_group_gemm_starts( void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts,
const torch::Tensor& a_starts, const torch::Tensor& b_starts, const torch::stable::Tensor& b_starts,
const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts, const torch::stable::Tensor& out_starts,
const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts, const torch::stable::Tensor& a_scales_starts,
const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb, const torch::stable::Tensor& b_scales_starts,
const torch::Tensor& a_strides, const torch::Tensor& b_strides, const torch::stable::Tensor& alpha_starts,
const torch::Tensor& c_strides, int64_t a_stride_val, int64_t b_stride_val, const torch::stable::Tensor& layout_sfa,
const torch::stable::Tensor& layout_sfb,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides,
int64_t a_stride_val, int64_t b_stride_val,
int64_t c_stride_val, int64_t c_stride_val,
/*these are used for their base addresses*/ /*these are used for their base addresses*/
torch::Tensor const& a_tensors, torch::Tensor const& b_tensors, torch::stable::Tensor const& a_tensors,
torch::Tensor const& out_tensors, torch::Tensor const& a_scales, torch::stable::Tensor const& b_tensors,
torch::Tensor const& b_scales, torch::Tensor const& alphas, torch::stable::Tensor const& out_tensors,
torch::Tensor const& expert_offsets, torch::Tensor const& sf_offsets, torch::stable::Tensor const& a_scales,
torch::Tensor const& problem_sizes, int M, int N, int K) { torch::stable::Tensor const& b_scales,
torch::stable::Tensor const& alphas,
torch::stable::Tensor const& expert_offsets,
torch::stable::Tensor const& sf_offsets,
torch::stable::Tensor const& problem_sizes,
int M, int N, int K) {
int num_experts = (int)expert_offsets.size(0); int num_experts = (int)expert_offsets.size(0);
auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index());
TORCH_CHECK(out_tensors.size(1) == N, STD_TORCH_CHECK(out_tensors.size(1) == N,
"Output tensor shape doesn't match expected shape"); "Output tensor shape doesn't match expected shape");
TORCH_CHECK(K / 2 == b_tensors.size(2), STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
"b_tensors(dim = 2) and a_tensors(dim = 1) trailing" "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
" dimension must match"); " dimension must match");
if (false) { if (false) {
...@@ -177,23 +185,27 @@ void run_get_group_gemm_starts( ...@@ -177,23 +185,27 @@ void run_get_group_gemm_starts(
//(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB,
// ScaleConfig) // ScaleConfig)
__CALL_GET_STARTS_KERNEL_BLOCKSCALE( __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kBFloat16, cutlass::float_e2m1_t, cutlass::float_ue4m3_t,
cutlass::bfloat16_t, LayoutSFA, LayoutSFB, ScaleConfig) torch::headeronly::ScalarType::BFloat16, cutlass::bfloat16_t, LayoutSFA,
LayoutSFB, ScaleConfig)
__CALL_GET_STARTS_KERNEL_BLOCKSCALE(cutlass::float_e2m1_t, __CALL_GET_STARTS_KERNEL_BLOCKSCALE(cutlass::float_e2m1_t,
cutlass::float_ue4m3_t, torch::kFloat16, cutlass::float_ue4m3_t,
half, LayoutSFA, LayoutSFB, ScaleConfig) torch::headeronly::ScalarType::Half, half,
LayoutSFA, LayoutSFB, ScaleConfig)
else { else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
} }
} }
template <typename OutType> template <typename OutType>
void run_fp4_blockwise_scaled_group_mm_sm100( void run_fp4_blockwise_scaled_group_mm_sm100(
torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b, torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales, const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::Tensor& alphas, const torch::Tensor& problem_sizes, const torch::stable::Tensor& b_blockscales,
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M, const torch::stable::Tensor& alphas,
int N, int K) { const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
using ProblemShape = using ProblemShape =
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>; cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t; using ElementType = cutlass::float_e2m1_t;
...@@ -272,20 +284,40 @@ void run_fp4_blockwise_scaled_group_mm_sm100( ...@@ -272,20 +284,40 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
int num_experts = static_cast<int>(expert_offsets.size(0)); int num_experts = static_cast<int>(expert_offsets.size(0));
auto options_int =
torch::TensorOptions().dtype(torch::kInt64).device(a.device()); torch::stable::Tensor a_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor a_ptrs = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor b_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor b_ptrs =
torch::Tensor out_ptrs = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor out_ptrs =
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int); std::nullopt, a.device());
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int); torch::stable::Tensor a_scales_ptrs =
torch::Tensor a_strides1 = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor b_strides1 = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor c_strides1 = torch::empty(num_experts, options_int); torch::stable::Tensor b_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor alpha_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor layout_sfa = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor layout_sfb = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor a_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor c_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>( run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
...@@ -308,7 +340,7 @@ void run_fp4_blockwise_scaled_group_mm_sm100( ...@@ -308,7 +340,7 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler; typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM; scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device(); hw_info.device_id = a.get_device_index();
static std::unordered_map<int, int> cached_sm_counts; static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] = cached_sm_counts[hw_info.device_id] =
...@@ -350,32 +382,35 @@ void run_fp4_blockwise_scaled_group_mm_sm100( ...@@ -350,32 +382,35 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
scheduler}; scheduler};
size_t workspace_size = Gemm::get_workspace_size(args); size_t workspace_size = Gemm::get_workspace_size(args);
auto const workspace_options = auto workspace =
torch::TensorOptions().dtype(torch::kUInt8).device(a.device()); torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
auto workspace = torch::empty(workspace_size, workspace_options); std::nullopt, a.device());
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
auto can_implement_status = gemm_op.can_implement(args); auto can_implement_status = gemm_op.can_implement(args);
TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, STD_TORCH_CHECK(
can_implement_status == cutlass::Status::kSuccess,
"Failed to implement GEMM: status=", (int)can_implement_status); "Failed to implement GEMM: status=", (int)can_implement_status);
// Run the GEMM // Run the GEMM
auto status = gemm_op.initialize(args, workspace.data_ptr()); auto status = gemm_op.initialize(args, workspace.data_ptr());
TORCH_CHECK(status == cutlass::Status::kSuccess, STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to initialize GEMM: status=", (int)status, "Failed to initialize GEMM: status=", (int)status,
" workspace_size=", workspace_size, " num_experts=", num_experts, " workspace_size=", workspace_size,
" M=", M, " N=", N, " K=", K); " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
status = gemm_op.run(args, workspace.data_ptr(), stream); status = gemm_op.run(args, workspace.data_ptr(), stream);
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
} }
void run_fp4_blockwise_scaled_group_mm_sm120( void run_fp4_blockwise_scaled_group_mm_sm120(
torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b, torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales, const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::Tensor& alphas, const torch::Tensor& problem_sizes, const torch::stable::Tensor& b_blockscales,
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M, const torch::stable::Tensor& alphas,
int N, int K) { const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
using ProblemShape = using ProblemShape =
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>; cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t; using ElementType = cutlass::float_e2m1_t;
...@@ -446,20 +481,40 @@ void run_fp4_blockwise_scaled_group_mm_sm120( ...@@ -446,20 +481,40 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
int num_experts = static_cast<int>(expert_offsets.size(0)); int num_experts = static_cast<int>(expert_offsets.size(0));
auto options_int =
torch::TensorOptions().dtype(torch::kInt64).device(a.device()); torch::stable::Tensor a_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor a_ptrs = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor b_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor b_ptrs =
torch::Tensor out_ptrs = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int); torch::stable::Tensor out_ptrs =
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int); std::nullopt, a.device());
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int); torch::stable::Tensor a_scales_ptrs =
torch::Tensor a_strides1 = torch::empty(num_experts, options_int); torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
torch::Tensor b_strides1 = torch::empty(num_experts, options_int); std::nullopt, a.device());
torch::Tensor c_strides1 = torch::empty(num_experts, options_int); torch::stable::Tensor b_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor alpha_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor layout_sfa = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor layout_sfb = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor a_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor c_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>( run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
...@@ -480,7 +535,7 @@ void run_fp4_blockwise_scaled_group_mm_sm120( ...@@ -480,7 +535,7 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions; using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler; typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM; scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device(); hw_info.device_id = a.get_device_index();
static std::unordered_map<int, int> cached_sm_counts; static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] = cached_sm_counts[hw_info.device_id] =
...@@ -523,33 +578,36 @@ void run_fp4_blockwise_scaled_group_mm_sm120( ...@@ -523,33 +578,36 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
scheduler}; scheduler};
size_t workspace_size = Gemm::get_workspace_size(args); size_t workspace_size = Gemm::get_workspace_size(args);
auto const workspace_options = auto workspace =
torch::TensorOptions().dtype(torch::kUInt8).device(a.device()); torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
auto workspace = torch::empty(workspace_size, workspace_options); std::nullopt, a.device());
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
auto can_implement_status = gemm_op.can_implement(args); auto can_implement_status = gemm_op.can_implement(args);
TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, STD_TORCH_CHECK(
can_implement_status == cutlass::Status::kSuccess,
"Failed to implement GEMM: status=", (int)can_implement_status); "Failed to implement GEMM: status=", (int)can_implement_status);
// Run the GEMM // Run the GEMM
auto status = gemm_op.initialize(args, workspace.data_ptr()); auto status = gemm_op.initialize(args, workspace.data_ptr());
TORCH_CHECK(status == cutlass::Status::kSuccess, STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to initialize GEMM: status=", (int)status, "Failed to initialize GEMM: status=", (int)status,
" workspace_size=", workspace_size, " num_experts=", num_experts, " workspace_size=", workspace_size,
" M=", M, " N=", N, " K=", K); " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
status = gemm_op.run(args, workspace.data_ptr(), stream); status = gemm_op.run(args, workspace.data_ptr(), stream);
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
} }
template <typename OutType> template <typename OutType>
void run_fp4_blockwise_scaled_group_mm( void run_fp4_blockwise_scaled_group_mm(
torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b, torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales, const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::Tensor& alphas, const torch::Tensor& problem_sizes, const torch::stable::Tensor& b_blockscales,
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M, const torch::stable::Tensor& alphas,
int N, int K) { const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
int32_t version_num = get_sm_version_num(); int32_t version_num = get_sm_version_num();
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
if (version_num >= 120 && version_num < 130) { if (version_num >= 120 && version_num < 130) {
...@@ -567,7 +625,7 @@ void run_fp4_blockwise_scaled_group_mm( ...@@ -567,7 +625,7 @@ void run_fp4_blockwise_scaled_group_mm(
return; return;
} }
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED( STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, false,
"No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ", "No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ",
version_num, ". Required capability: 100 or 120"); version_num, ". Required capability: 100 or 120");
...@@ -575,26 +633,31 @@ void run_fp4_blockwise_scaled_group_mm( ...@@ -575,26 +633,31 @@ void run_fp4_blockwise_scaled_group_mm(
#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \ #if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
(defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120) (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte; constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn; constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
#endif #endif
#define CHECK_TYPE(x, st, m) \ #define CHECK_TYPE(x, st, m) \
TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m) STD_TORCH_CHECK(x.scalar_type() == st, \
": Inconsistency of torch::stable::Tensor type:", m)
#define CHECK_TH_CUDA(x, m) \ #define CHECK_TH_CUDA(x, m) \
TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.") STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
#define CHECK_CONTIGUOUS(x, m) \ #define CHECK_CONTIGUOUS(x, m) \
TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.") STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
#define CHECK_INPUT(x, st, m) \ #define CHECK_INPUT(x, st, m) \
CHECK_TH_CUDA(x, m); \ CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m); \ CHECK_CONTIGUOUS(x, m); \
CHECK_TYPE(x, st, m) CHECK_TYPE(x, st, m)
void cutlass_fp4_group_mm( void cutlass_fp4_group_mm(torch::stable::Tensor& output,
torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b, const torch::stable::Tensor& a,
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales, const torch::stable::Tensor& b,
const torch::Tensor& alphas, const torch::Tensor& problem_sizes, const torch::stable::Tensor& a_blockscale,
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) { const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& alphas,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets) {
#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \ #if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
(defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120) (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
// Input validation // Input validation
...@@ -602,22 +665,26 @@ void cutlass_fp4_group_mm( ...@@ -602,22 +665,26 @@ void cutlass_fp4_group_mm(
CHECK_INPUT(b, FLOAT4_E2M1X2, "b"); CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale"); CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale");
CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales"); CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales");
CHECK_INPUT(alphas, at::ScalarType::Float, "alphas"); CHECK_INPUT(alphas, torch::headeronly::ScalarType::Float, "alphas");
TORCH_CHECK(a_blockscale.dim() == 2, STD_TORCH_CHECK(
a_blockscale.dim() == 2,
"expected a_blockscale to be of shape [num_experts, rounded_m," "expected a_blockscale to be of shape [num_experts, rounded_m,"
" k // group_size], observed rank: ", " k // group_size], observed rank: ",
a_blockscale.dim()) a_blockscale.dim())
TORCH_CHECK(b_blockscales.dim() == 3, STD_TORCH_CHECK(b_blockscales.dim() == 3,
"expected b_blockscale to be of shape: " "expected b_blockscale to be of shape: "
" [num_experts, n, k // group_size], observed rank: ", " [num_experts, n, k // group_size], observed rank: ",
b_blockscales.dim()) b_blockscales.dim())
TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be a 2D tensor"); STD_TORCH_CHECK(problem_sizes.dim() == 2,
TORCH_CHECK(problem_sizes.size(1) == 3, "problem_sizes must be a 2D tensor");
STD_TORCH_CHECK(problem_sizes.size(1) == 3,
"problem_sizes must have the shape (num_experts, 3)"); "problem_sizes must have the shape (num_experts, 3)");
TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0), STD_TORCH_CHECK(
problem_sizes.size(0) == expert_offsets.size(0),
"Number of experts in problem_sizes must match expert_offsets"); "Number of experts in problem_sizes must match expert_offsets");
TORCH_CHECK(problem_sizes.dtype() == torch::kInt32, STD_TORCH_CHECK(
problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
"problem_sizes must be int32."); "problem_sizes must be int32.");
int M = static_cast<int>(a.size(0)); int M = static_cast<int>(a.size(0));
...@@ -625,7 +692,7 @@ void cutlass_fp4_group_mm( ...@@ -625,7 +692,7 @@ void cutlass_fp4_group_mm(
int E = static_cast<int>(b.size(0)); int E = static_cast<int>(b.size(0));
int K = static_cast<int>(2 * b.size(2)); int K = static_cast<int>(2 * b.size(2));
if (output.scalar_type() == torch::kBFloat16) { if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
run_fp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>( run_fp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes, output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
expert_offsets, sf_offsets, M, N, K); expert_offsets, sf_offsets, M, N, K);
...@@ -633,7 +700,7 @@ void cutlass_fp4_group_mm( ...@@ -633,7 +700,7 @@ void cutlass_fp4_group_mm(
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
int32_t version_num = get_sm_version_num(); int32_t version_num = get_sm_version_num();
if (version_num >= 120 && version_num < 130) { if (version_num >= 120 && version_num < 130) {
TORCH_CHECK_NOT_IMPLEMENTED( STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ", false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ",
output.scalar_type()); output.scalar_type());
} }
...@@ -643,7 +710,7 @@ void cutlass_fp4_group_mm( ...@@ -643,7 +710,7 @@ void cutlass_fp4_group_mm(
expert_offsets, sf_offsets, M, N, K); expert_offsets, sf_offsets, M, N, K);
} }
#else #else
TORCH_CHECK_NOT_IMPLEMENTED( STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, false,
"No compiled cutlass_fp4_group_mm kernel, vLLM must " "No compiled cutlass_fp4_group_mm kernel, vLLM must "
"be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 " "be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 "
...@@ -651,6 +718,6 @@ void cutlass_fp4_group_mm( ...@@ -651,6 +718,6 @@ void cutlass_fp4_group_mm(
#endif #endif
} }
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("cutlass_fp4_group_mm", &cutlass_fp4_group_mm); m.impl("cutlass_fp4_group_mm", TORCH_BOX(&cutlass_fp4_group_mm));
} }
csrc/quantization/fp4/nvfp4_experts_quant.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
View file @ 7c080dd3
...@@ -14,16 +14,15 @@ ...@@ -14,16 +14,15 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include <cuda_runtime_api.h> #include <cuda_runtime_api.h>
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp8.h> #include <cuda_fp8.h>
#include "dispatch_utils.h"
#include "cuda_utils.h" #include "cuda_utils.h"
#include "nvfp4_utils.cuh" #include "nvfp4_utils.cuh"
...@@ -327,25 +326,28 @@ void quant_impl(void* output, void* output_scale, void* input, ...@@ -327,25 +326,28 @@ void quant_impl(void* output, void* output_scale, void* input,
} // namespace vllm } // namespace vllm
/*Quantization entry for fp4 experts quantization*/ /*Quantization entry for fp4 experts quantization*/
#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor") #define CHECK_TH_CUDA(x, m) \
STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x, m) \ #define CHECK_CONTIGUOUS(x, m) \
TORCH_CHECK(x.is_contiguous(), m, "must be contiguous") STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
#define CHECK_INPUT(x, m) \ #define CHECK_INPUT(x, m) \
CHECK_TH_CUDA(x, m); \ CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m); CHECK_CONTIGUOUS(x, m);
constexpr auto HALF = at::ScalarType::Half; constexpr auto HALF = torch::headeronly::ScalarType::Half;
constexpr auto BF16 = at::ScalarType::BFloat16; constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
constexpr auto FLOAT = at::ScalarType::Float; constexpr auto FLOAT = torch::headeronly::ScalarType::Float;
constexpr auto INT = at::ScalarType::Int; constexpr auto INT = torch::headeronly::ScalarType::Int;
constexpr auto UINT8 = at::ScalarType::Byte; constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
// Common validation for fp4 experts quantization entry points. // Common validation for fp4 experts quantization entry points.
static void validate_fp4_experts_quant_inputs( static void validate_fp4_experts_quant_inputs(
torch::Tensor const& output, torch::Tensor const& output_scale, torch::stable::Tensor const& output,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& output_scale,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input,
torch::Tensor const& output_scale_offset_by_experts, int64_t m_topk, torch::stable::Tensor const& input_global_scale,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
int64_t k) { int64_t k) {
CHECK_INPUT(output, "output"); CHECK_INPUT(output, "output");
CHECK_INPUT(output_scale, "output_scale"); CHECK_INPUT(output_scale, "output_scale");
...@@ -354,41 +356,42 @@ static void validate_fp4_experts_quant_inputs( ...@@ -354,41 +356,42 @@ static void validate_fp4_experts_quant_inputs(
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts"); CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts"); CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
TORCH_CHECK(output.dim() == 2); STD_TORCH_CHECK(output.dim() == 2);
TORCH_CHECK(output_scale.dim() == 2); STD_TORCH_CHECK(output_scale.dim() == 2);
TORCH_CHECK(input.dim() == 2); STD_TORCH_CHECK(input.dim() == 2);
TORCH_CHECK(input_global_scale.dim() == 1); STD_TORCH_CHECK(input_global_scale.dim() == 1);
TORCH_CHECK(input_offset_by_experts.dim() == 1); STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
TORCH_CHECK(output_scale_offset_by_experts.dim() == 1); STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16); STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
TORCH_CHECK(input_global_scale.scalar_type() == FLOAT); STD_TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
TORCH_CHECK(input_offset_by_experts.scalar_type() == INT); STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT); STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
// output is uint8 (two nvfp4 values are packed into one uint8) // output is uint8 (two nvfp4 values are packed into one uint8)
// output_scale is int32 (four fp8 values are packed into one int32) // output_scale is int32 (four fp8 values are packed into one int32)
TORCH_CHECK(output.scalar_type() == UINT8); STD_TORCH_CHECK(output.scalar_type() == UINT8);
TORCH_CHECK(output_scale.scalar_type() == INT); STD_TORCH_CHECK(output_scale.scalar_type() == INT);
const int BLOCK_SIZE = 16; const int BLOCK_SIZE = 16;
TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); STD_TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
auto n_experts = input_global_scale.size(0); auto n_experts = input_global_scale.size(0);
TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1); STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1); STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
TORCH_CHECK(output.size(0) == m_topk); STD_TORCH_CHECK(output.size(0) == m_topk);
TORCH_CHECK(output.size(1) == k / 2); STD_TORCH_CHECK(output.size(1) == k / 2);
int scales_k = k / BLOCK_SIZE; int scales_k = k / BLOCK_SIZE;
// 4 means the swizzle requirement by nvidia nvfp4. // 4 means the swizzle requirement by nvidia nvfp4.
int padded_k = (scales_k + (4 - 1)) / 4 * 4; int padded_k = (scales_k + (4 - 1)) / 4 * 4;
// 4 means 4 fp8 values are packed into one int32 // 4 means 4 fp8 values are packed into one int32
TORCH_CHECK(output_scale.size(1) * 4 == padded_k); STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
} }
void scaled_fp4_experts_quant_sm1xxa( void scaled_fp4_experts_quant_sm1xxa(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts) { torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts) {
auto m_topk = input.size(0); auto m_topk = input.size(0);
auto k = input.size(1); auto k = input.size(1);
...@@ -397,11 +400,11 @@ void scaled_fp4_experts_quant_sm1xxa( ...@@ -397,11 +400,11 @@ void scaled_fp4_experts_quant_sm1xxa(
output_scale_offset_by_experts, m_topk, k); output_scale_offset_by_experts, m_topk, k);
auto n_experts = input_global_scale.size(0); auto n_experts = input_global_scale.size(0);
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); const torch::stable::accelerator::DeviceGuard device_guard(
const cudaStream_t stream = input.get_device_index());
at::cuda::getCurrentCUDAStream(input.get_device()); const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_DISPATCH_HALF_TYPES( VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "nvfp4_experts_quant_kernel", [&] { input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>( vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
...@@ -413,14 +416,15 @@ void scaled_fp4_experts_quant_sm1xxa( ...@@ -413,14 +416,15 @@ void scaled_fp4_experts_quant_sm1xxa(
} }
void silu_and_mul_scaled_fp4_experts_quant_sm1xxa( void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts) { torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts) {
auto m_topk = input.size(0); auto m_topk = input.size(0);
// Input has gate || up layout, so k = input.size(1) / 2 // Input has gate || up layout, so k = input.size(1) / 2
auto k_times_2 = input.size(1); auto k_times_2 = input.size(1);
TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)"); STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
auto k = k_times_2 / 2; auto k = k_times_2 / 2;
validate_fp4_experts_quant_inputs(output, output_scale, input, validate_fp4_experts_quant_inputs(output, output_scale, input,
...@@ -428,11 +432,11 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa( ...@@ -428,11 +432,11 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
output_scale_offset_by_experts, m_topk, k); output_scale_offset_by_experts, m_topk, k);
auto n_experts = input_global_scale.size(0); auto n_experts = input_global_scale.size(0);
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); const torch::stable::accelerator::DeviceGuard device_guard(
const cudaStream_t stream = input.get_device_index());
at::cuda::getCurrentCUDAStream(input.get_device()); const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_DISPATCH_HALF_TYPES( VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] { input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>( vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
......
csrc/quantization/fp4/nvfp4_quant_entry.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu
View file @ 7c080dd3
...@@ -14,44 +14,48 @@ ...@@ -14,44 +14,48 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "cutlass_extensions/common.hpp" #include "cutlass_extensions/common.hpp"
#include "nvfp4_utils.cuh" #include "nvfp4_utils.cuh"
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
torch::Tensor const& input, torch::stable::Tensor const& input,
torch::Tensor const& output_sf, torch::stable::Tensor const& output_sf,
torch::Tensor const& input_sf, torch::stable::Tensor const& input_sf,
bool is_sf_swizzled_layout); bool is_sf_swizzled_layout);
#endif #endif
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
void scaled_fp4_experts_quant_sm1xxa( void scaled_fp4_experts_quant_sm1xxa(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts); torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts);
#endif #endif
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, void silu_and_mul_nvfp4_quant_sm1xxa(torch::stable::Tensor& output,
torch::Tensor& output_sf, torch::stable::Tensor& output_sf,
torch::Tensor& input, torch::stable::Tensor& input,
torch::Tensor& input_sf); torch::stable::Tensor& input_sf);
#endif #endif
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
void silu_and_mul_scaled_fp4_experts_quant_sm1xxa( void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts); torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts);
#endif #endif
static bool nvfp4_quant_sm_supported() { static bool nvfp4_quant_sm_supported() {
...@@ -65,43 +69,44 @@ static bool nvfp4_quant_sm_supported() { ...@@ -65,43 +69,44 @@ static bool nvfp4_quant_sm_supported() {
return false; return false;
} }
void scaled_fp4_quant_out(torch::Tensor const& input, void scaled_fp4_quant_out(torch::stable::Tensor const& input,
torch::Tensor const& input_sf, torch::stable::Tensor const& input_sf,
bool is_sf_swizzled_layout, torch::Tensor& output, bool is_sf_swizzled_layout,
torch::Tensor& output_sf) { torch::stable::Tensor& output,
torch::stable::Tensor& output_sf) {
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
TORCH_CHECK(nvfp4_quant_sm_supported(), STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
"No compiled nvfp4 quantization kernel for SM ", "No compiled nvfp4 quantization kernel for SM ",
get_sm_version_num(), get_sm_version_num(),
". Recompile with the appropriate CUDA arch."); ". Recompile with the appropriate CUDA arch.");
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf, return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
is_sf_swizzled_layout); is_sf_swizzled_layout);
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel"); STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
"No compiled nvfp4 quantization kernel");
} }
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func( std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
torch::Tensor const& input, torch::Tensor const& input_sf, torch::stable::Tensor const& input, torch::stable::Tensor const& input_sf,
bool is_sf_swizzled_layout) { bool is_sf_swizzled_layout) {
int64_t n = input.size(-1); int64_t n = input.size(-1);
int64_t m = input.numel() / n; int64_t m = input.numel() / n;
auto device = input.device(); auto device = input.device();
// Two fp4 values packed into a uint8 // Two fp4 values packed into a uint8
auto output = torch::empty( auto output = torch::stable::empty(
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8)); {m, n / 2}, torch::headeronly::ScalarType::Byte, std::nullopt, device);
torch::Tensor output_sf; torch::stable::Tensor output_sf;
if (is_sf_swizzled_layout) { if (is_sf_swizzled_layout) {
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n); auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
output_sf = torch::empty( output_sf = torch::stable::empty(
{sf_m, sf_n}, {sf_m, sf_n}, torch::headeronly::ScalarType::Int, std::nullopt, device);
torch::TensorOptions().device(device).dtype(torch::kInt32));
} else { } else {
output_sf = torch::empty( output_sf = torch::stable::empty({m, n / CVT_FP4_SF_VEC_SIZE},
{m, n / CVT_FP4_SF_VEC_SIZE}, torch::headeronly::ScalarType::Byte,
torch::TensorOptions().device(device).dtype(torch::kUInt8)); std::nullopt, device);
} }
scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output, scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
...@@ -110,13 +115,14 @@ std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func( ...@@ -110,13 +115,14 @@ std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
} }
void scaled_fp4_experts_quant( void scaled_fp4_experts_quant(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts) { torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts) {
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
TORCH_CHECK(nvfp4_quant_sm_supported(), STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
"No compiled nvfp4 experts quantization kernel for SM ", "No compiled nvfp4 experts quantization kernel for SM ",
get_sm_version_num(), get_sm_version_num(),
". Recompile with the appropriate CUDA arch."); ". Recompile with the appropriate CUDA arch.");
...@@ -124,32 +130,35 @@ void scaled_fp4_experts_quant( ...@@ -124,32 +130,35 @@ void scaled_fp4_experts_quant(
output, output_scale, input, input_global_scale, input_offset_by_experts, output, output_scale, input, input_global_scale, input_offset_by_experts,
output_scale_offset_by_experts); output_scale_offset_by_experts);
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED(false, STD_TORCH_CHECK_NOT_IMPLEMENTED(
"No compiled nvfp4 experts quantization kernel"); false, "No compiled nvfp4 experts quantization kernel");
} }
void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf, void silu_and_mul_nvfp4_quant(torch::stable::Tensor& output,
torch::Tensor& input, torch::Tensor& input_sf) { torch::stable::Tensor& output_sf,
torch::stable::Tensor& input,
torch::stable::Tensor& input_sf) {
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
TORCH_CHECK(nvfp4_quant_sm_supported(), STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
"No compiled silu_and_mul nvfp4 quantization kernel for SM ", "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
get_sm_version_num(), get_sm_version_num(),
". Recompile with the appropriate CUDA arch."); ". Recompile with the appropriate CUDA arch.");
return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf); return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED( STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, "No compiled silu_and_mul nvfp4 quantization kernel"); false, "No compiled silu_and_mul nvfp4 quantization kernel");
} }
void silu_and_mul_scaled_fp4_experts_quant( void silu_and_mul_scaled_fp4_experts_quant(
torch::Tensor& output, torch::Tensor& output_scale, torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale, torch::stable::Tensor const& input,
torch::Tensor const& input_offset_by_experts, torch::stable::Tensor const& input_global_scale,
torch::Tensor const& output_scale_offset_by_experts) { torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts) {
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
TORCH_CHECK(nvfp4_quant_sm_supported(), STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
"No compiled silu_and_mul nvfp4 experts quantization kernel " "No compiled silu_and_mul nvfp4 experts quantization kernel "
"for SM ", "for SM ",
get_sm_version_num(), get_sm_version_num(),
...@@ -158,6 +167,6 @@ void silu_and_mul_scaled_fp4_experts_quant( ...@@ -158,6 +167,6 @@ void silu_and_mul_scaled_fp4_experts_quant(
output, output_scale, input, input_global_scale, input_offset_by_experts, output, output_scale, input, input_global_scale, input_offset_by_experts,
output_scale_offset_by_experts); output_scale_offset_by_experts);
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED( STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, "No compiled silu_and_mul nvfp4 experts quantization kernel"); false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
} }
csrc/quantization/fp4/nvfp4_quant_kernels.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
View file @ 7c080dd3
...@@ -14,16 +14,16 @@ ...@@ -14,16 +14,16 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include <cuda_runtime_api.h> #include <cuda_runtime_api.h>
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp8.h> #include <cuda_fp8.h>
#include "dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "cuda_utils.h" #include "cuda_utils.h"
#include "launch_bounds_utils.h" #include "launch_bounds_utils.h"
...@@ -173,17 +173,18 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) ...@@ -173,17 +173,18 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
} // namespace vllm } // namespace vllm
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
torch::Tensor const& input, torch::stable::Tensor const& input,
torch::Tensor const& output_sf, torch::stable::Tensor const& output_sf,
torch::Tensor const& input_sf, torch::stable::Tensor const& input_sf,
bool is_sf_swizzled_layout) { bool is_sf_swizzled_layout) {
int32_t m = input.size(0); int32_t m = input.size(0);
int32_t n = input.size(1); int32_t n = input.size(1);
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16."); STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
TORCH_CHECK(input.scalar_type() == at::ScalarType::Half || STD_TORCH_CHECK(
input.scalar_type() == at::ScalarType::BFloat16, input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Unsupported input data type for quantize_to_fp4."); "Unsupported input data type for quantize_to_fp4.");
int multiProcessorCount = int multiProcessorCount =
...@@ -192,8 +193,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, ...@@ -192,8 +193,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr()); auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr()); auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
auto output_ptr = static_cast<int64_t*>(output.data_ptr()); auto output_ptr = static_cast<int64_t*>(output.data_ptr());
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); const torch::stable::accelerator::DeviceGuard device_guard(
auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); input.get_device_index());
auto stream = get_current_cuda_stream(input.get_device_index());
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE); int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
...@@ -213,10 +215,10 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, ...@@ -213,10 +215,10 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
dim3 grid(grid_x, grid_y); dim3 grid(grid_x, grid_y);
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] { VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "nvfp4_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr()); auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
// NOTE: We don't support e8m0 scales at this moment.
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>( vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
m, n, num_padded_cols, input_ptr, input_sf_ptr, m, n, num_padded_cols, input_ptr, input_sf_ptr,
reinterpret_cast<uint32_t*>(output_ptr), reinterpret_cast<uint32_t*>(output_ptr),
...@@ -229,13 +231,13 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, ...@@ -229,13 +231,13 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
dim3 grid(grid_x, grid_y); dim3 grid(grid_x, grid_y);
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] { VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "nvfp4_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr()); auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
// NOTE: We don't support e8m0 scales at this moment.
vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false> vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
<<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, num_packed_cols, <<<grid, block, 0, stream>>>(
input_ptr, input_sf_ptr, m, n, sf_n_unpadded, num_packed_cols, input_ptr, input_sf_ptr,
reinterpret_cast<uint32_t*>(output_ptr), reinterpret_cast<uint32_t*>(output_ptr),
reinterpret_cast<uint32_t*>(sf_out)); reinterpret_cast<uint32_t*>(sf_out));
}); });
......
csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu
View file @ 7c080dd3
...@@ -14,32 +14,39 @@ ...@@ -14,32 +14,39 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include <c10/cuda/CUDAGuard.h>
#include "libtorch_stable/torch_utils.h"
#include "cutlass_extensions/common.hpp" #include "cutlass_extensions/common.hpp"
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A, void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
torch::Tensor const& B, torch::stable::Tensor const& A,
torch::Tensor const& A_sf, torch::stable::Tensor const& B,
torch::Tensor const& B_sf, torch::stable::Tensor const& A_sf,
torch::Tensor const& alpha); torch::stable::Tensor const& B_sf,
torch::stable::Tensor const& alpha);
#endif #endif
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A, void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
torch::Tensor const& B, torch::stable::Tensor const& A,
torch::Tensor const& A_sf, torch::stable::Tensor const& B,
torch::Tensor const& B_sf, torch::stable::Tensor const& A_sf,
torch::Tensor const& alpha); torch::stable::Tensor const& B_sf,
torch::stable::Tensor const& alpha);
#endif #endif
void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A, void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
const torch::Tensor& B, const torch::Tensor& A_sf, const torch::stable::Tensor& A,
const torch::Tensor& B_sf, const torch::stable::Tensor& B,
const torch::Tensor& alpha) { const torch::stable::Tensor& A_sf,
// Make sure we’re on A’s device. const torch::stable::Tensor& B_sf,
const c10::cuda::OptionalCUDAGuard device_guard(device_of(A)); const torch::stable::Tensor& alpha) {
// Make sure we're on A's device.
const torch::stable::accelerator::DeviceGuard device_guard(
A.get_device_index());
const int32_t sm = get_sm_version_num(); const int32_t sm = get_sm_version_num();
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100 #if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
...@@ -56,7 +63,8 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A, ...@@ -56,7 +63,8 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
} }
#endif #endif
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 mm kernel for SM ", sm, STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, "No compiled nvfp4 mm kernel for SM ", sm,
". Recompile with CUDA >= 12.8 and CC >= 100."); ". Recompile with CUDA >= 12.8 and CC >= 100.");
} }
......
csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu → csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu
View file @ 7c080dd3
...@@ -14,10 +14,9 @@ ...@@ -14,10 +14,9 @@
* limitations under the License. * limitations under the License.
*/ */
#include <torch/all.h> #include <torch/csrc/stable/tensor.h>
#include <ATen/cuda/CUDAContext.h> #include "libtorch_stable/torch_utils.h"
#include <c10/cuda/CUDAGuard.h>
#include "cutlass_extensions/common.hpp" #include "cutlass_extensions/common.hpp"
...@@ -127,8 +126,9 @@ struct Fp4GemmSm100 { ...@@ -127,8 +126,9 @@ struct Fp4GemmSm100 {
template <typename Config> template <typename Config>
typename Config::Gemm::Arguments args_from_options( typename Config::Gemm::Arguments args_from_options(
at::Tensor& D, at::Tensor const& A, at::Tensor const& B, torch::stable::Tensor& D, torch::stable::Tensor const& A,
at::Tensor const& A_sf, at::Tensor const& B_sf, at::Tensor const& alpha, torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf,
torch::stable::Tensor const& B_sf, torch::stable::Tensor const& alpha,
int64_t M, int64_t N, int64_t K) { int64_t M, int64_t N, int64_t K) {
using ElementA = typename Config::Gemm::ElementA; using ElementA = typename Config::Gemm::ElementA;
using ElementB = typename Config::Gemm::ElementB; using ElementB = typename Config::Gemm::ElementB;
...@@ -174,19 +174,20 @@ typename Config::Gemm::Arguments args_from_options( ...@@ -174,19 +174,20 @@ typename Config::Gemm::Arguments args_from_options(
} }
template <typename Config> template <typename Config>
void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B, void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A,
at::Tensor const& A_sf, at::Tensor const& B_sf, torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf,
at::Tensor const& alpha, int64_t m, int64_t n, int64_t k, torch::stable::Tensor const& B_sf,
cudaStream_t stream) { torch::stable::Tensor const& alpha, int64_t m, int64_t n,
int64_t k, cudaStream_t stream) {
typename Config::Gemm gemm; typename Config::Gemm gemm;
auto arguments = auto arguments =
args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k); args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k);
size_t workspace_size = Config::Gemm::get_workspace_size(arguments); size_t workspace_size = Config::Gemm::get_workspace_size(arguments);
auto const workspace_options = auto workspace =
torch::TensorOptions().dtype(torch::kUInt8).device(A.device()); torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
auto workspace = torch::empty(workspace_size, workspace_options); std::nullopt, A.device());
CUTLASS_CHECK(gemm.can_implement(arguments)); CUTLASS_CHECK(gemm.can_implement(arguments));
...@@ -197,12 +198,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B, ...@@ -197,12 +198,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
// Dispatch function to select appropriate config based on M // Dispatch function to select appropriate config based on M
template <typename OutType> template <typename OutType>
void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A, void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
torch::Tensor const& B, torch::stable::Tensor const& A,
torch::Tensor const& A_sf, torch::stable::Tensor const& B,
torch::Tensor const& B_sf, torch::stable::Tensor const& A_sf,
torch::Tensor const& alpha, int64_t m, int64_t n, torch::stable::Tensor const& B_sf,
int64_t k, cudaStream_t stream) { torch::stable::Tensor const& alpha, int64_t m,
int64_t n, int64_t k, cudaStream_t stream) {
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m)); uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
if (mp2 <= 16) { if (mp2 <= 16) {
...@@ -222,61 +224,65 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A, ...@@ -222,61 +224,65 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
#else #else
template <typename OutType> template <typename OutType>
void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A, void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
torch::Tensor const& B, torch::stable::Tensor const& A,
torch::Tensor const& A_sf, torch::stable::Tensor const& B,
torch::Tensor const& B_sf, torch::stable::Tensor const& A_sf,
torch::Tensor const& alpha, int64_t m, int64_t n, torch::stable::Tensor const& B_sf,
int64_t k, cudaStream_t stream) { torch::stable::Tensor const& alpha, int64_t m,
TORCH_CHECK(false, int64_t n, int64_t k, cudaStream_t stream) {
STD_TORCH_CHECK(false,
"Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to " "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
"a CUTLASS 3.8 source directory to enable support."); "a CUTLASS 3.8 source directory to enable support.");
} }
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) #endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
#define CHECK_TYPE(x, st, m) \ #define CHECK_TYPE(x, st, m) \
TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m) STD_TORCH_CHECK(x.scalar_type() == st, \
": Inconsistency of torch::stable::Tensor type:", m)
#define CHECK_TH_CUDA(x, m) \ #define CHECK_TH_CUDA(x, m) \
TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor") STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x, m) \ #define CHECK_CONTIGUOUS(x, m) \
TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous") STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
#define CHECK_INPUT(x, st, m) \ #define CHECK_INPUT(x, st, m) \
CHECK_TH_CUDA(x, m); \ CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m); \ CHECK_CONTIGUOUS(x, m); \
CHECK_TYPE(x, st, m) CHECK_TYPE(x, st, m)
constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte; constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn; constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A, void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
torch::Tensor const& B, torch::stable::Tensor const& A,
torch::Tensor const& A_sf, torch::stable::Tensor const& B,
torch::Tensor const& B_sf, torch::stable::Tensor const& A_sf,
torch::Tensor const& alpha) { torch::stable::Tensor const& B_sf,
torch::stable::Tensor const& alpha) {
CHECK_INPUT(A, FLOAT4_E2M1X2, "a"); CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
CHECK_INPUT(B, FLOAT4_E2M1X2, "b"); CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
CHECK_INPUT(A_sf, SF_DTYPE, "scale_a"); CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
CHECK_INPUT(B_sf, SF_DTYPE, "scale_b"); CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
CHECK_INPUT(alpha, at::ScalarType::Float, "alpha"); CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha");
TORCH_CHECK(A.dim() == 2, "a must be a matrix"); STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix");
TORCH_CHECK(B.dim() == 2, "b must be a matrix"); STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix");
TORCH_CHECK(A.sizes()[1] == B.sizes()[1], STD_TORCH_CHECK(A.size(1) == B.size(1),
"a and b shapes cannot be multiplied (", A.sizes()[0], "x", "a and b shapes cannot be multiplied (", A.size(0), "x",
A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")"); A.size(1), " and ", B.size(0), "x", B.size(1), ")");
auto const m = A.sizes()[0]; auto const m = A.size(0);
auto const n = B.sizes()[0]; auto const n = B.size(0);
auto const k = A.sizes()[1] * 2; auto const k = A.size(1) * 2;
constexpr int alignment = 32; constexpr int alignment = 32;
TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment, STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ",
", but got a shape: (", A.sizes()[0], "x", A.sizes()[1], alignment, ", but got a shape: (", A.size(0), "x", A.size(1),
"), k: ", k, "."); "), k: ", k, ".");
TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment, STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ",
", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ")."); alignment, ", but got b shape: (", B.size(0), "x", B.size(1),
").");
auto round_up = [](int x, int y) { return (x + y - 1) / y * y; }; auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
int rounded_m = round_up(m, 128); int rounded_m = round_up(m, 128);
...@@ -285,33 +291,34 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A, ...@@ -285,33 +291,34 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
// integer. // integer.
int rounded_k = round_up(k / 16, 4); int rounded_k = round_up(k / 16, 4);
TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix"); STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix"); STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1], STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1),
"scale_a and scale_b shapes cannot be multiplied (", "scale_a and scale_b shapes cannot be multiplied (",
A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0], A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x",
"x", B_sf.sizes()[1], ")"); B_sf.size(1), ")");
TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k, STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
"scale_a must be padded and swizzled to a shape (", rounded_m, "scale_a must be padded and swizzled to a shape (", rounded_m,
"x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x", "x", rounded_k, "), but got a shape (", A_sf.size(0), "x",
A_sf.sizes()[1], ")"); A_sf.size(1), ")");
TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k, STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
"scale_b must be padded and swizzled to a shape (", rounded_n, "scale_b must be padded and swizzled to a shape (", rounded_n,
"x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x", "x", rounded_k, "), but got a shape (", B_sf.size(0), "x",
B_sf.sizes()[1], ")"); B_sf.size(1), ")");
auto out_dtype = D.dtype(); auto out_dtype = D.scalar_type();
const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); const torch::stable::accelerator::DeviceGuard device_guard(
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device()); A.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
if (out_dtype == at::ScalarType::Half) { if (out_dtype == torch::headeronly::ScalarType::Half) {
cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n,
k, stream); k, stream);
} else if (out_dtype == at::ScalarType::BFloat16) { } else if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha,
m, n, k, stream); m, n, k, stream);
} else { } else {
TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype, STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (",
")"); out_dtype, ")");
} }
} }
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment