A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 18:30 UTC
[4/n] Migrate FP4/W4A8 CUTLASS kernels to torch stable ABI (#37503)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
mikaylagawarecki committed
Mar 31, 2026 at 13:21 UTC
7c080dd3c5b794f31906aff92f4cb829cde4986a
27 files changed
+1215
-1026
CMakeLists.txt
+94
-84
@@ -340,8 +340,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
340
341
list(APPEND VLLM_EXT_SRC
342
"csrc/quantization/awq/gemm_kernels.cu"
343
- "csrc/quantization/fp4/nvfp4_quant_entry.cu"
344
- "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
343
"csrc/cutlass_extensions/common.cpp")
344
345
set_gencode_flags_for_srcs(
@@ -489,59 +487,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
487
" in CUDA target architectures")
488
endif()
489
492
- # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
493
- # CUDA 12.8 or later
494
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
495
- cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
496
- else()
497
- cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
498
- endif()
499
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
500
- set(SRCS
501
- "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
502
- "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
503
- "csrc/quantization/fp4/nvfp4_experts_quant.cu"
504
- "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
505
- "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
506
- set_gencode_flags_for_srcs(
507
- SRCS "${SRCS}"
508
- CUDA_ARCHS "${FP4_ARCHS}")
509
- list(APPEND VLLM_EXT_SRC "${SRCS}")
510
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
511
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
512
- message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
513
- else()
514
- message(STATUS "Not building NVFP4 as no compatible archs were found.")
515
- # clear FP4_ARCHS
516
- set(FP4_ARCHS)
517
- endif()
518
-
519
- # FP4 Archs and flags
520
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
521
- cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
522
- else()
523
- cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
524
- endif()
525
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
526
- set(SRCS
527
- "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
528
- "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
529
- "csrc/quantization/fp4/nvfp4_experts_quant.cu"
530
- "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
531
- "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
532
- set_gencode_flags_for_srcs(
533
- SRCS "${SRCS}"
534
- CUDA_ARCHS "${FP4_ARCHS}")
535
- list(APPEND VLLM_EXT_SRC "${SRCS}")
536
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
537
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
538
- message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
539
- else()
540
- message(STATUS "Not building NVFP4 as no compatible archs were found.")
541
- # clear FP4_ARCHS
542
- set(FP4_ARCHS)
543
- endif()
544
-
490
# CUTLASS MLA Archs and flags
491
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
492
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
@@ -681,34 +626,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
626
endif()
627
endif()
628
684
- # Only build W4A8 kernels if we are building for something compatible with sm90a
685
- cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
686
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
687
- set(SRCS
688
- "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
689
- "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
690
- "csrc/quantization/cutlass_w4a8/w4a8_utils.cu"
691
- )
692
-
693
- set_gencode_flags_for_srcs(
694
- SRCS "${SRCS}"
695
- CUDA_ARCHS "${W4A8_ARCHS}")
696
-
697
- list(APPEND VLLM_EXT_SRC "${SRCS}")
698
-
699
- message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
700
- else()
701
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
702
- AND W4A8_ARCHS)
703
- message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
704
- "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
705
- "later if you intend on running w4a16 quantized models on "
706
- "Hopper.")
707
- else()
708
- message(STATUS "Not building W4A8 kernels as no compatible archs "
709
- "found in CUDA target architectures")
710
- endif()
711
- endif()
629
630
# Hadacore kernels
631
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
@@ -760,7 +677,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
677
set(VLLM_STABLE_EXT_SRC
678
"csrc/libtorch_stable/torch_bindings.cpp"
679
"csrc/cutlass_extensions/common.cpp"
763
- "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
680
+ "csrc/cuda_utils_kernels.cu"
681
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
682
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
683
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
684
685
if(VLLM_GPU_LANG STREQUAL "CUDA")
686
list(APPEND VLLM_STABLE_EXT_SRC
@@ -978,6 +898,96 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
898
endif()
899
endif()
900
901
+ #
902
+ # FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
903
+ #
904
+
905
+ # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
906
+ # CUDA 12.8 or later
907
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
908
+ cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
909
+ else()
910
+ cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
911
+ endif()
912
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
913
+ set(SRCS
914
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
915
+ "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
916
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
917
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
918
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
919
+ set_gencode_flags_for_srcs(
920
+ SRCS "${SRCS}"
921
+ CUDA_ARCHS "${FP4_ARCHS}")
922
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
923
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
924
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
925
+ message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
926
+ else()
927
+ message(STATUS "Not building NVFP4 as no compatible archs were found.")
928
+ # clear FP4_ARCHS
929
+ set(FP4_ARCHS)
930
+ endif()
931
+
932
+ # FP4 Archs and flags
933
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
934
+ cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
935
+ else()
936
+ cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
937
+ endif()
938
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
939
+ set(SRCS
940
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
941
+ "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
942
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
943
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
944
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
945
+ set_gencode_flags_for_srcs(
946
+ SRCS "${SRCS}"
947
+ CUDA_ARCHS "${FP4_ARCHS}")
948
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
949
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
950
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
951
+ message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
952
+ else()
953
+ message(STATUS "Not building NVFP4 as no compatible archs were found.")
954
+ # clear FP4_ARCHS
955
+ set(FP4_ARCHS)
956
+ endif()
957
+
958
+ #
959
+ # W4A8 kernels (moved from _C to _C_stable_libtorch)
960
+ #
961
+
962
+ # Only build W4A8 kernels if we are building for something compatible with sm90a
963
+ cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
964
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
965
+ set(SRCS
966
+ "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
967
+ "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
968
+ "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu"
969
+ )
970
+
971
+ set_gencode_flags_for_srcs(
972
+ SRCS "${SRCS}"
973
+ CUDA_ARCHS "${W4A8_ARCHS}")
974
+
975
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
976
+
977
+ message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
978
+ else()
979
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
980
+ AND W4A8_ARCHS)
981
+ message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
982
+ "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
983
+ "later if you intend on running w4a16 quantized models on "
984
+ "Hopper.")
985
+ else()
986
+ message(STATUS "Not building W4A8 kernels as no compatible archs "
987
+ "found in CUDA target architectures")
988
+ endif()
989
+ endif()
990
+
991
message(STATUS "Enabling C_stable extension.")
992
define_extension_target(
993
_C_stable_libtorch
csrc/cuda_vec_utils.cuh
+2
-2
@@ -3,8 +3,8 @@
3
4
#pragma once
5
6
-#include <c10/util/BFloat16.h>
7
-#include <c10/util/Half.h>
6
+#include <torch/headeronly/util/BFloat16.h>
7
+#include <torch/headeronly/util/Half.h>
8
#include <cassert>
9
10
#ifdef USE_ROCM
csrc/cutlass_extensions/cute_utils.cuh
-1
@@ -1,7 +1,6 @@
1
#pragma once
2
3
#include <cute/tensor.hpp>
4
-#include <torch/all.h>
4
namespace cute {
5
6
////////////////////////////////////////////////////////////////////
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
+20
-20
@@ -189,9 +189,9 @@ struct Sm90RowOrScalarBroadcastArray {
189
}
190
191
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
192
- Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
193
- Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
194
- Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
192
+ cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
193
+ cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
194
+ cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
195
196
for (int i = 0; i < size(tGS_gRow_flt); ++i) {
197
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
@@ -211,8 +211,8 @@ struct Sm90RowOrScalarBroadcastArray {
211
begin_loop(int epi_m, int epi_n) {
212
if (epi_m == 0) { // Assumes M-major subtile loop
213
if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
214
- Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
215
- Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
214
+ cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
215
+ cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
216
copy(tSR_sRow_flt, tSR_rRow_flt);
217
}
218
}
@@ -241,9 +241,9 @@ struct Sm90RowOrScalarBroadcastArray {
241
auto [m, n, k, l] = args.tile_coord_mnkl;
242
using ThreadCount = decltype(size(args.tiled_copy));
243
244
- Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
245
- Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
246
- Tensor sRow = make_tensor(make_smem_ptr(smem),
244
+ cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
245
+ cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
246
+ cute::Tensor sRow = make_tensor(make_smem_ptr(smem),
247
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
248
//// G2S: Gmem to Smem
249
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
@@ -251,16 +251,16 @@ struct Sm90RowOrScalarBroadcastArray {
251
Stride<_0, _1>>{},
252
Layout<_1>{});
253
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
254
- Tensor tGS_gRow = thr_g2s.partition_S(gRow);
255
- Tensor tGS_sRow = thr_g2s.partition_D(sRow);
254
+ cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
255
+ cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
256
257
//// G2S: Coord
258
auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
259
- Tensor tGS_cRow = thr_g2s.partition_S(cRow);
259
+ cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
260
261
//// S2R: Smem to Reg
262
- Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
263
- Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
262
+ cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
263
+ cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
264
265
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>(
266
tGS_gRow,
@@ -389,7 +389,7 @@ struct Sm90ColOrScalarBroadcastArray {
389
390
CUTLASS_DEVICE void
391
begin() {
392
- Tensor pred = make_tensor<bool>(shape(tCgCol));
392
+ cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
393
CUTLASS_PRAGMA_UNROLL
394
for (int i = 0; i < size(pred); ++i) {
395
pred(i) = get<0>(tCcCol(i)) < m;
@@ -409,7 +409,7 @@ struct Sm90ColOrScalarBroadcastArray {
409
CUTLASS_DEVICE Array<Element, FragmentSize>
410
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
411
Array<Element, FragmentSize> frg_col;
412
- Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
412
+ cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
413
414
CUTLASS_PRAGMA_UNROLL
415
for (int i = 0; i < FragmentSize; ++i) {
@@ -431,16 +431,16 @@ struct Sm90ColOrScalarBroadcastArray {
431
auto [M, N, K, L] = args.problem_shape_mnkl;
432
auto [m, n, k, l] = args.tile_coord_mnkl;
433
434
- Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
435
- Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
434
+ cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
435
+ cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
436
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
437
- Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
437
+ cute::Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
438
439
// Generate an identity tensor matching the shape of the global tensor and
440
// partition the same way, this will be used to generate the predicate
441
// tensor for loading
442
- Tensor cCol = make_identity_tensor(mCol.shape());
443
- Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
442
+ cute::Tensor cCol = make_identity_tensor(mCol.shape());
443
+ cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
444
cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
445
446
return ConsumerStoreCallbacks(
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
+20
-20
@@ -186,9 +186,9 @@ struct Sm90RowOrScalarBroadcast {
186
}
187
188
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
189
- Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
190
- Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
191
- Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
189
+ cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
190
+ cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
191
+ cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
192
193
for (int i = 0; i < size(tGS_gRow_flt); ++i) {
194
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
@@ -208,8 +208,8 @@ struct Sm90RowOrScalarBroadcast {
208
begin_loop(int epi_m, int epi_n) {
209
if (epi_m == 0) { // Assumes M-major subtile loop
210
if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
211
- Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
212
- Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
211
+ cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
212
+ cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
213
copy(tSR_sRow_flt, tSR_rRow_flt);
214
}
215
}
@@ -238,9 +238,9 @@ struct Sm90RowOrScalarBroadcast {
238
auto [m, n, k, l] = args.tile_coord_mnkl;
239
using ThreadCount = decltype(size(args.tiled_copy));
240
241
- Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
242
- Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
243
- Tensor sRow = make_tensor(make_smem_ptr(smem),
241
+ cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
242
+ cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
243
+ cute::Tensor sRow = make_tensor(make_smem_ptr(smem),
244
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
245
//// G2S: Gmem to Smem
246
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
@@ -248,16 +248,16 @@ struct Sm90RowOrScalarBroadcast {
248
Stride<_0, _1>>{},
249
Layout<_1>{});
250
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
251
- Tensor tGS_gRow = thr_g2s.partition_S(gRow);
252
- Tensor tGS_sRow = thr_g2s.partition_D(sRow);
251
+ cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
252
+ cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
253
254
//// G2S: Coord
255
auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
256
- Tensor tGS_cRow = thr_g2s.partition_S(cRow);
256
+ cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
257
258
//// S2R: Smem to Reg
259
- Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
260
- Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
259
+ cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
260
+ cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
261
262
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>(
263
tGS_gRow,
@@ -382,7 +382,7 @@ struct Sm90ColOrScalarBroadcast {
382
383
CUTLASS_DEVICE void
384
begin() {
385
- Tensor pred = make_tensor<bool>(shape(tCgCol));
385
+ cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
386
CUTLASS_PRAGMA_UNROLL
387
for (int i = 0; i < size(pred); ++i) {
388
pred(i) = get<0>(tCcCol(i)) < m;
@@ -402,7 +402,7 @@ struct Sm90ColOrScalarBroadcast {
402
CUTLASS_DEVICE Array<Element, FragmentSize>
403
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
404
Array<Element, FragmentSize> frg_col;
405
- Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
405
+ cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
406
407
CUTLASS_PRAGMA_UNROLL
408
for (int i = 0; i < FragmentSize; ++i) {
@@ -422,16 +422,16 @@ struct Sm90ColOrScalarBroadcast {
422
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
423
424
auto [M, N, K, L] = args.problem_shape_mnkl;
425
- Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
426
- Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
425
+ cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
426
+ cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
427
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
428
- Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
428
+ cute::Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
429
430
// Generate an identity tensor matching the shape of the global tensor and
431
// partition the same way, this will be used to generate the predicate
432
// tensor for loading
433
- Tensor cCol = make_identity_tensor(mCol.shape());
434
- Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
433
+ cute::Tensor cCol = make_identity_tensor(mCol.shape());
434
+ cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
435
cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
436
437
return ConsumerStoreCallbacks(
csrc/cutlass_extensions/torch_utils.hpp
+52
-37
@@ -1,6 +1,21 @@
1
#pragma once
2
3
-#include <torch/all.h>
3
+// This header is shared between _C (unstable ABI, used by machete) and
4
+// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION
5
+// is defined only for the stable target, so we switch includes and types
6
+// accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor.
7
+#ifdef TORCH_TARGET_VERSION
8
+ #include <torch/csrc/stable/tensor.h>
9
+ #include <torch/headeronly/util/BFloat16.h>
10
+ #include <torch/headeronly/util/Half.h>
11
+ #include <torch/headeronly/util/shim_utils.h> // for STD_TORCH_CHECK
12
+using TorchTensor = torch::stable::Tensor;
13
+ #define TORCH_UTILS_CHECK STD_TORCH_CHECK
14
+#else
15
+ #include <torch/all.h>
16
+using TorchTensor = torch::Tensor;
17
+ #define TORCH_UTILS_CHECK TORCH_CHECK
18
+#endif
19
20
#include "cute/layout.hpp"
21
#include "cutlass/layout/matrix.h"
@@ -55,35 +70,35 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) {
70
// If `tensor.dim() < rank(Stride{})`, the shape is padded with 1s and the extra
71
// strides are set to be 0 or 1.
72
template <typename Stride>
58
-static inline auto make_cute_layout(torch::Tensor const& tensor,
73
+static inline auto make_cute_layout(TorchTensor const& tensor,
74
std::string_view name = "tensor") {
60
- TORCH_CHECK(tensor.dim() <= rank(Stride{}));
61
- auto stride = cute::transform_with_idx(
62
- Stride{}, [&](auto const& stride_ele, auto const& idx) {
63
- using StrideEle = std::decay_t<decltype(stride_ele)>;
64
-
65
- if (idx < tensor.dim()) {
66
- if constexpr (cute::is_static_v<StrideEle>) {
67
- TORCH_CHECK(StrideEle::value == tensor.stride(idx), "Expected ",
68
- name, ".stride(", idx, ") to be ", StrideEle::value);
69
- return StrideEle{};
70
- } else {
71
- if (tensor.size(idx) == 1) {
72
- // use 0 stride for dim with size 1, this is easier for
73
- // cute/cutlass to optimize (helps the TMA code flatten dims)
74
- return StrideEle{0};
75
- } else {
76
- return tensor.stride(idx);
77
- }
78
- }
75
+ TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{}));
76
+ auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele,
77
+ auto const& idx) {
78
+ using StrideEle = std::decay_t<decltype(stride_ele)>;
79
+
80
+ if (idx < tensor.dim()) {
81
+ if constexpr (cute::is_static_v<StrideEle>) {
82
+ TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ",
83
+ name, ".stride(", idx, ") to be ", StrideEle::value);
84
+ return StrideEle{};
85
+ } else {
86
+ if (tensor.size(idx) == 1) {
87
+ // use 0 stride for dim with size 1, this is easier for
88
+ // cute/cutlass to optimize (helps the TMA code flatten dims)
89
+ return StrideEle{0};
90
} else {
80
- // Extra strides are assumed to be 0 or 1
81
- if constexpr (cute::is_static_v<StrideEle>) {
82
- static_assert(StrideEle::value == 0 || StrideEle::value == 1);
83
- }
84
- return StrideEle{};
91
+ return tensor.stride(idx);
92
}
86
- });
93
+ }
94
+ } else {
95
+ // Extra strides are assumed to be 0 or 1
96
+ if constexpr (cute::is_static_v<StrideEle>) {
97
+ static_assert(StrideEle::value == 0 || StrideEle::value == 1);
98
+ }
99
+ return StrideEle{};
100
+ }
101
+ });
102
103
auto shape = cute::make_shape_from_idx<rank(Stride{})>([&](auto const& idx) {
104
if (idx < tensor.dim())
@@ -97,7 +112,7 @@ static inline auto make_cute_layout(torch::Tensor const& tensor,
112
113
template <typename Stride>
114
static inline auto maybe_make_cute_layout(
100
- std::optional<torch::Tensor> const& tensor,
115
+ std::optional<TorchTensor> const& tensor,
116
std::string_view name = "tensor") {
117
using Layout = decltype(make_cute_layout<Stride>(*tensor));
118
@@ -121,12 +136,12 @@ template <typename T>
136
using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type;
137
138
template <>
124
-struct equivalent_cutlass_type<c10::Half> {
139
+struct equivalent_cutlass_type<torch::headeronly::Half> {
140
using type = cutlass::half_t;
141
};
142
143
template <>
129
-struct equivalent_cutlass_type<c10::BFloat16> {
144
+struct equivalent_cutlass_type<torch::headeronly::BFloat16> {
145
using type = cutlass::bfloat16_t;
146
};
147
@@ -134,8 +149,8 @@ struct equivalent_cutlass_type<c10::BFloat16> {
149
// equivalent_scalar_t (basically inverse of equivalent_cutlass_type)
150
//
151
137
-// Return a `c10::CppTypeToScalarType<T>` compatible type, i.e. get the C++ from
138
-// c10 that is equivalent to T, e.g.: `cutlass::half_t -> c10::Half`
152
+// Return a `torch::headeronly::CppTypeToScalarType<T>` compatible type, i.e.
153
+// get the C++ type equivalent to T, e.g.: `cutlass::half_t -> Half`
154
template <typename T>
155
struct equivalent_scalar_type {
156
using type = T;
@@ -146,15 +161,15 @@ using equivalent_scalar_type_t = typename equivalent_scalar_type<T>::type;
161
162
template <>
163
struct equivalent_scalar_type<cutlass::half_t> {
149
- using type = c10::Half;
164
+ using type = torch::headeronly::Half;
165
};
166
167
template <>
168
struct equivalent_scalar_type<cutlass::bfloat16_t> {
154
- using type = c10::BFloat16;
169
+ using type = torch::headeronly::BFloat16;
170
};
171
157
-// get equivalent c10::ScalarType tag from compile time type
172
+// get equivalent torch::headeronly::ScalarType tag from compile time type
173
template <typename T>
159
-static inline constexpr c10::ScalarType equivalent_scalar_type_v =
160
- c10::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;
\ No newline at end of file
174
+static inline constexpr torch::headeronly::ScalarType equivalent_scalar_type_v =
175
+ torch::headeronly::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;
csrc/libtorch_stable/dispatch_utils.h
+9
@@ -49,6 +49,15 @@
49
THO_DISPATCH_SWITCH(TYPE, NAME, \
50
VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
51
52
+// Half types dispatch (Half + BFloat16)
53
+#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \
54
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
55
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
56
+
57
+#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
58
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
59
+ VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
60
+
61
// Boolean dispatch
62
#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
63
if (expr) { \
csrc/libtorch_stable/ops.h
+50
@@ -84,4 +84,54 @@ void get_cutlass_batched_moe_mm_data(
84
const torch::stable::Tensor& expert_num_tokens,
85
const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
86
const int64_t k);
87
+
88
+// FP4/NVFP4 ops
89
+bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
90
+
91
+void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
92
+ torch::stable::Tensor const& A,
93
+ torch::stable::Tensor const& B,
94
+ torch::stable::Tensor const& A_sf,
95
+ torch::stable::Tensor const& B_sf,
96
+ torch::stable::Tensor const& alpha);
97
+
98
+void cutlass_fp4_group_mm(torch::stable::Tensor& output,
99
+ const torch::stable::Tensor& a,
100
+ const torch::stable::Tensor& b,
101
+ const torch::stable::Tensor& a_blockscale,
102
+ const torch::stable::Tensor& b_blockscales,
103
+ const torch::stable::Tensor& alphas,
104
+ const torch::stable::Tensor& problem_sizes,
105
+ const torch::stable::Tensor& expert_offsets,
106
+ const torch::stable::Tensor& sf_offsets);
107
+
108
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
109
+ torch::stable::Tensor const& input,
110
+ torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout);
111
+
112
+void scaled_fp4_quant_out(torch::stable::Tensor const& input,
113
+ torch::stable::Tensor const& input_scale,
114
+ bool is_sf_swizzled_layout,
115
+ torch::stable::Tensor& output,
116
+ torch::stable::Tensor& output_scale);
117
+
118
+void scaled_fp4_experts_quant(
119
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
120
+ torch::stable::Tensor const& input,
121
+ torch::stable::Tensor const& input_global_scale,
122
+ torch::stable::Tensor const& input_offset_by_experts,
123
+ torch::stable::Tensor const& output_scale_offset_by_experts);
124
+
125
+void silu_and_mul_scaled_fp4_experts_quant(
126
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
127
+ torch::stable::Tensor const& input,
128
+ torch::stable::Tensor const& input_global_scale,
129
+ torch::stable::Tensor const& input_offset_by_experts,
130
+ torch::stable::Tensor const& output_scale_offset_by_experts);
131
+
132
+void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
133
+ torch::stable::Tensor& output_block_scale,
134
+ torch::stable::Tensor& input,
135
+ torch::stable::Tensor& input_global_scale);
136
+
137
#endif
csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh
renamed
+36
-25
@@ -2,10 +2,9 @@
2
#pragma once
3
4
#include <cuda.h>
5
-#include <torch/all.h>
6
-#include <c10/cuda/CUDAStream.h>
5
+#include <torch/csrc/stable/tensor.h>
6
+#include "libtorch_stable/torch_utils.h"
7
8
-#include "core/scalar_type.hpp"
8
#include "cutlass/bfloat16.h"
9
#include "cutlass/float8.h"
10
@@ -41,7 +40,7 @@ __global__ void get_group_gemm_starts(
40
}
41
42
#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
44
- else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
43
+ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
44
get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \
45
cutlass::Array<cutlass::float_e4m3_t, 8>> \
46
<<<1, num_experts, 0, stream>>>( \
@@ -66,23 +65,34 @@ __global__ void get_group_gemm_starts(
65
namespace {
66
67
void run_get_group_gemm_starts(
69
- torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
70
- torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
71
- torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
72
- torch::Tensor& b_group_scales_ptrs, torch::Tensor const& a_tensors,
73
- torch::Tensor const& b_tensors, torch::Tensor& out_tensors,
74
- torch::Tensor const& a_scales, torch::Tensor const& b_scales,
75
- torch::Tensor const& b_group_scales, const int64_t b_group_size) {
76
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
77
- TORCH_CHECK(b_tensors.dtype() == torch::kInt32); // int4 8x packed into int32
78
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
79
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
80
- TORCH_CHECK(b_group_scales.dtype() ==
81
- torch::kFloat8_e4m3fn); // the underlying torch type is e4m3
82
- TORCH_CHECK(out_tensors.dtype() ==
83
- torch::kBFloat16); // only support bf16 for now
68
+ torch::stable::Tensor const& expert_offsets, torch::stable::Tensor& a_ptrs,
69
+ torch::stable::Tensor& b_ptrs, torch::stable::Tensor& out_ptrs,
70
+ torch::stable::Tensor& a_scales_ptrs, torch::stable::Tensor& b_scales_ptrs,
71
+ torch::stable::Tensor& b_group_scales_ptrs,
72
+ torch::stable::Tensor const& a_tensors,
73
+ torch::stable::Tensor const& b_tensors, torch::stable::Tensor& out_tensors,
74
+ torch::stable::Tensor const& a_scales,
75
+ torch::stable::Tensor const& b_scales,
76
+ torch::stable::Tensor const& b_group_scales, const int64_t b_group_size) {
77
+ STD_TORCH_CHECK(a_tensors.scalar_type() ==
78
+ torch::headeronly::ScalarType::Float8_e4m3fn);
79
+ STD_TORCH_CHECK(
80
+ b_tensors.scalar_type() ==
81
+ torch::headeronly::ScalarType::Int); // int4 8x packed into int32
82
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
83
+ torch::headeronly::ScalarType::Float);
84
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
85
+ torch::headeronly::ScalarType::Float);
86
+ STD_TORCH_CHECK(
87
+ b_group_scales.scalar_type() ==
88
+ torch::headeronly::ScalarType::Float8_e4m3fn); // the underlying torch
89
+ // type is e4m3
90
+ STD_TORCH_CHECK(
91
+ out_tensors.scalar_type() ==
92
+ torch::headeronly::ScalarType::BFloat16); // only support bf16 for now
93
// expect int64_t to avoid overflow during offset calculations
85
- TORCH_CHECK(expert_offsets.dtype() == torch::kInt64);
94
+ STD_TORCH_CHECK(expert_offsets.scalar_type() ==
95
+ torch::headeronly::ScalarType::Long);
96
97
int num_experts = static_cast<int>(expert_offsets.size(0));
98
// logical k, n
@@ -90,15 +100,16 @@ void run_get_group_gemm_starts(
100
int64_t k = a_tensors.size(1);
101
int64_t scale_k = cutlass::ceil_div(k, b_group_size);
102
93
- auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
103
+ auto stream = get_current_cuda_stream(a_tensors.get_device_index());
104
105
if (false) {
106
}
97
- __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
98
- __CALL_GET_STARTS_KERNEL(torch::kFloat16, half)
107
+ __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::BFloat16,
108
+ cutlass::bfloat16_t)
109
+ __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::Half, half)
110
else {
100
- TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
111
+ STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
112
}
113
}
114
104
-} // namespace
\ No newline at end of file
115
+} // namespace
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu
renamed
+83
-63
@@ -14,13 +14,12 @@
14
#include "cutlass/util/mixed_dtype_utils.hpp"
15
16
// vllm includes
17
-#include <ATen/cuda/CUDAContext.h>
18
-#include <c10/cuda/CUDAGuard.h>
19
-#include <torch/all.h>
17
+#include <torch/csrc/stable/library.h>
18
+#include <torch/csrc/stable/tensor.h>
19
+#include "libtorch_stable/torch_utils.h"
20
#include "cutlass_extensions/torch_utils.hpp"
21
#include "cutlass_extensions/common.hpp"
22
23
-#include "core/registration.h"
23
#include "get_group_starts.cuh"
24
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
25
#include "w4a8_utils.cuh"
@@ -168,31 +167,40 @@ struct W4A8GroupedGemmKernel {
167
static_assert(sizeof(LayoutB_Reordered) % sizeof(int32_t) == 0,
168
"LayoutB_Reordered size must be divisible by 4 bytes");
169
171
- static void grouped_mm(
172
- torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
173
- const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
174
- const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
175
- const int64_t b_group_size, const torch::Tensor& expert_offsets,
176
- const torch::Tensor& problem_sizes_torch, const torch::Tensor& a_strides,
177
- const torch::Tensor& b_strides, const torch::Tensor& c_strides,
178
- const torch::Tensor& group_scale_strides) {
170
+ static void grouped_mm(torch::stable::Tensor& out_tensors,
171
+ const torch::stable::Tensor& a_tensors,
172
+ const torch::stable::Tensor& b_tensors,
173
+ const torch::stable::Tensor& a_scales,
174
+ const torch::stable::Tensor& b_scales,
175
+ const torch::stable::Tensor& b_group_scales,
176
+ const int64_t b_group_size,
177
+ const torch::stable::Tensor& expert_offsets,
178
+ const torch::stable::Tensor& problem_sizes_torch,
179
+ const torch::stable::Tensor& a_strides,
180
+ const torch::stable::Tensor& b_strides,
181
+ const torch::stable::Tensor& c_strides,
182
+ const torch::stable::Tensor& group_scale_strides) {
183
auto device = a_tensors.device();
184
auto device_id = device.index();
181
- const at::cuda::OptionalCUDAGuard device_guard(device);
182
- auto stream = at::cuda::getCurrentCUDAStream(device_id);
185
+ const torch::stable::accelerator::DeviceGuard device_guard(device_id);
186
+ auto stream = get_current_cuda_stream(device_id);
187
188
int num_experts = static_cast<int>(expert_offsets.size(0));
189
int n = static_cast<int>(b_tensors.size(1));
190
int k = static_cast<int>(b_tensors.size(2)) * PackFactor;
191
188
- auto options_int =
189
- torch::TensorOptions().dtype(torch::kInt64).device(device);
190
- torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
191
- torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
192
- torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
193
- torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
194
- torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
195
- torch::Tensor b_group_scales_ptrs = torch::empty(num_experts, options_int);
192
+ torch::stable::Tensor a_ptrs = torch::stable::empty(
193
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
194
+ torch::stable::Tensor b_ptrs = torch::stable::empty(
195
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
196
+ torch::stable::Tensor out_ptrs = torch::stable::empty(
197
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
198
+ torch::stable::Tensor a_scales_ptrs = torch::stable::empty(
199
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
200
+ torch::stable::Tensor b_scales_ptrs = torch::stable::empty(
201
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
202
+ torch::stable::Tensor b_group_scales_ptrs = torch::stable::empty(
203
+ num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
204
205
// get the correct offsets to pass to gemm
206
run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
@@ -247,9 +255,9 @@ struct W4A8GroupedGemmKernel {
255
256
// Allocate workspace
257
size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
250
- torch::Tensor workspace =
251
- torch::empty(workspace_size,
252
- torch::TensorOptions().dtype(torch::kU8).device(device));
258
+ torch::stable::Tensor workspace = torch::stable::empty(
259
+ workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
260
+ device);
261
262
// Run GEMM
263
GemmShuffled gemm;
@@ -294,14 +302,20 @@ using Kernel_256x128_2x1x1_Coop =
302
using Kernel_128x256_2x1x1_Coop =
303
W4A8GroupedGemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>, Coop, CoopEpi>;
304
297
-void mm_dispatch(
298
- torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
299
- const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
300
- const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
301
- const int64_t b_group_size, const torch::Tensor& expert_offsets,
302
- const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
303
- const torch::Tensor& b_strides, const torch::Tensor& c_strides,
304
- const torch::Tensor& group_scale_strides, const std::string& schedule) {
305
+void mm_dispatch(torch::stable::Tensor& out_tensors,
306
+ const torch::stable::Tensor& a_tensors,
307
+ const torch::stable::Tensor& b_tensors,
308
+ const torch::stable::Tensor& a_scales,
309
+ const torch::stable::Tensor& b_scales,
310
+ const torch::stable::Tensor& b_group_scales,
311
+ const int64_t b_group_size,
312
+ const torch::stable::Tensor& expert_offsets,
313
+ const torch::stable::Tensor& problem_sizes,
314
+ const torch::stable::Tensor& a_strides,
315
+ const torch::stable::Tensor& b_strides,
316
+ const torch::stable::Tensor& c_strides,
317
+ const torch::stable::Tensor& group_scale_strides,
318
+ const std::string& schedule) {
319
if (schedule == "Kernel_128x16_1x1x1_Coop") {
320
Kernel_128x16_1x1x1_Coop::grouped_mm(
321
out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
@@ -358,18 +372,23 @@ void mm_dispatch(
372
b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
373
c_strides, group_scale_strides);
374
} else {
361
- TORCH_CHECK(false,
362
- "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
375
+ STD_TORCH_CHECK(false,
376
+ "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
377
}
378
}
379
366
-void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
367
- const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
368
- const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
369
- const int64_t b_group_size, const torch::Tensor& expert_offsets,
370
- const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
371
- const torch::Tensor& b_strides, const torch::Tensor& c_strides,
372
- const torch::Tensor& group_scale_strides,
380
+void mm(torch::stable::Tensor& out_tensors,
381
+ const torch::stable::Tensor& a_tensors,
382
+ const torch::stable::Tensor& b_tensors,
383
+ const torch::stable::Tensor& a_scales,
384
+ const torch::stable::Tensor& b_scales,
385
+ const torch::stable::Tensor& b_group_scales, const int64_t b_group_size,
386
+ const torch::stable::Tensor& expert_offsets,
387
+ const torch::stable::Tensor& problem_sizes,
388
+ const torch::stable::Tensor& a_strides,
389
+ const torch::stable::Tensor& b_strides,
390
+ const torch::stable::Tensor& c_strides,
391
+ const torch::stable::Tensor& group_scale_strides,
392
std::optional<std::string> maybe_schedule) {
393
// user has specified a schedule
394
if (maybe_schedule) {
@@ -406,26 +425,27 @@ void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
425
a_strides, b_strides, c_strides, group_scale_strides, schedule);
426
}
427
409
-std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
410
- torch::Tensor const& b_tensors) {
411
- TORCH_CHECK(b_tensors.dtype() == torch::kInt32);
412
- TORCH_CHECK(b_tensors.dim() == 3); // (experts, n, k)
413
- TORCH_CHECK(b_tensors.is_contiguous());
414
- TORCH_CHECK(b_tensors.is_cuda());
428
+std::tuple<torch::stable::Tensor, torch::stable::Tensor>
429
+encode_and_reorder_int4b(torch::stable::Tensor const& b_tensors) {
430
+ STD_TORCH_CHECK(b_tensors.scalar_type() ==
431
+ torch::headeronly::ScalarType::Int);
432
+ STD_TORCH_CHECK(b_tensors.dim() == 3); // (experts, n, k)
433
+ STD_TORCH_CHECK(b_tensors.is_contiguous());
434
+ STD_TORCH_CHECK(b_tensors.is_cuda());
435
436
int n = static_cast<int>(b_tensors.size(1));
437
int k = static_cast<int>(b_tensors.size(2)) * PackFactor; // logical k
438
439
// CUTLASS reorder_tensor requires k % 256 == 0 and n % 16 == 0.
440
// These misalignments cause silent OOB unless run under Compute Sanitizer.
421
- TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
422
- TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
441
+ STD_TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
442
+ STD_TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
443
444
// we will store the layout to an int32 tensor;
445
// this is the number of elements we need per layout
446
constexpr size_t layout_width = sizeof(LayoutB_Reordered) / sizeof(int32_t);
447
428
- torch::Tensor b_tensors_packed = torch::empty_like(b_tensors);
448
+ torch::stable::Tensor b_tensors_packed = torch::stable::empty_like(b_tensors);
449
int num_experts = static_cast<int>(b_tensors.size(0));
450
451
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(
455
size_t num_int4_elems = 1ull * num_experts * n * k;
456
bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(b_ptr, b_packed_ptr,
457
num_int4_elems);
438
- TORCH_CHECK(ok, "unified_encode_int4b failed");
458
+ STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
459
460
// construct the layout once; assumes each expert has the same layout
461
using LayoutType = LayoutB_Reordered;
@@ -456,28 +476,28 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
476
}
477
478
// save the packed layout to torch tensor so we can re-use it
459
- auto cpu_opts =
460
- torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
461
- torch::Tensor layout_cpu =
462
- torch::empty({num_experts, layout_width}, cpu_opts);
479
+ torch::stable::Tensor layout_cpu = torch::stable::empty(
480
+ {num_experts, layout_width}, torch::headeronly::ScalarType::Int,
481
+ std::nullopt, torch::stable::Device(torch::stable::DeviceType::CPU));
482
464
- int32_t* layout_data = layout_cpu.data_ptr<int32_t>();
483
+ int32_t* layout_data = layout_cpu.mutable_data_ptr<int32_t>();
484
for (int i = 0; i < num_experts; ++i) {
485
std::memcpy(layout_data + i * layout_width, // dst (int32*)
486
&layout_B_reordered, // src (LayoutType*)
487
sizeof(LayoutType)); // number of bytes
488
}
489
471
- torch::Tensor packed_layout =
472
- layout_cpu.to(b_tensors.device(), /*non_blocking=*/false);
490
+ torch::stable::Tensor packed_layout =
491
+ torch::stable::to(layout_cpu, b_tensors.device(),
492
+ /*non_blocking=*/false);
493
494
return {b_tensors_packed, packed_layout};
495
}
496
477
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
478
- m.impl("cutlass_w4a8_moe_mm", &mm);
479
- m.impl("cutlass_encode_and_reorder_int4b_grouped", &encode_and_reorder_int4b);
497
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
498
+ m.impl("cutlass_w4a8_moe_mm", TORCH_BOX(&mm));
499
+ m.impl("cutlass_encode_and_reorder_int4b_grouped",
500
+ TORCH_BOX(&encode_and_reorder_int4b));
501
}
502
503
} // namespace vllm::cutlass_w4a8_moe
483
-/////////////////////////////////////////////////////////////////////////////////////////////////
\ No newline at end of file
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu
renamed
+60
-58
@@ -3,14 +3,12 @@
3
// https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu
4
//
5
6
-#include <ATen/cuda/CUDAContext.h>
7
-#include <c10/cuda/CUDAGuard.h>
8
-#include <torch/all.h>
6
+#include <torch/csrc/stable/library.h>
7
+#include <torch/csrc/stable/tensor.h>
8
+#include "libtorch_stable/torch_utils.h"
9
#include "cutlass_extensions/torch_utils.hpp"
10
#include "w4a8_utils.cuh"
11
12
-#include "core/registration.h"
13
-
12
#include "cutlass/cutlass.h"
13
#include <limits>
14
@@ -161,31 +159,31 @@ struct W4A8GemmKernel {
159
using StrideD = typename GemmKernelShuffled::StrideD;
160
using StrideS = typename CollectiveMainloopShuffled::StrideScale;
161
164
- static torch::Tensor mm(torch::Tensor const& A,
165
- torch::Tensor const& B, // already packed
166
- torch::Tensor const& group_scales, // already packed
167
- int64_t group_size,
168
- torch::Tensor const& channel_scales,
169
- torch::Tensor const& token_scales,
170
- std::optional<at::ScalarType> const& maybe_out_type) {
162
+ static torch::stable::Tensor mm(
163
+ torch::stable::Tensor const& A,
164
+ torch::stable::Tensor const& B, // already packed
165
+ torch::stable::Tensor const& group_scales, // already packed
166
+ int64_t group_size, torch::stable::Tensor const& channel_scales,
167
+ torch::stable::Tensor const& token_scales,
168
+ std::optional<torch::headeronly::ScalarType> const& maybe_out_type) {
169
// TODO: param validation
170
int m = A.size(0);
171
int k = A.size(1);
172
int n = B.size(1);
173
174
// safely cast group_size to int
177
- TORCH_CHECK(group_size > 0 && group_size <= std::numeric_limits<int>::max(),
178
- "group_size out of supported range for int: ", group_size);
175
+ STD_TORCH_CHECK(
176
+ group_size > 0 && group_size <= std::numeric_limits<int>::max(),
177
+ "group_size out of supported range for int: ", group_size);
178
int const group_size_int = static_cast<int>(group_size);
179
180
// Allocate output
182
- const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
181
+ const torch::stable::accelerator::DeviceGuard device_guard(
182
+ A.get_device_index());
183
auto device = A.device();
184
- auto stream = at::cuda::getCurrentCUDAStream(device.index());
185
- torch::Tensor D =
186
- torch::empty({m, n}, torch::TensorOptions()
187
- .dtype(equivalent_scalar_type_v<ElementD>)
188
- .device(device));
184
+ auto stream = get_current_cuda_stream(device.index());
185
+ torch::stable::Tensor D = torch::stable::empty(
186
+ {m, n}, equivalent_scalar_type_v<ElementD>, std::nullopt, device);
187
// prepare arg pointers
188
auto A_ptr = static_cast<MmaType const*>(A.const_data_ptr());
189
auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
@@ -237,9 +235,9 @@ struct W4A8GemmKernel {
235
236
// Workspace
237
size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
240
- torch::Tensor workspace =
241
- torch::empty(workspace_size,
242
- torch::TensorOptions().dtype(torch::kU8).device(device));
238
+ torch::stable::Tensor workspace = torch::stable::empty(
239
+ workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
240
+ device);
241
242
// Run GEMM
243
GemmShuffled gemm;
@@ -269,14 +267,14 @@ using Kernel_128x64_1x1x1 = W4A8GemmKernel<Shape<_128, _64>, Shape<_1, _1, _1>>;
267
using Kernel_128x32_1x1x1 = W4A8GemmKernel<Shape<_128, _32>, Shape<_1, _1, _1>>;
268
using Kernel_128x16_1x1x1 = W4A8GemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>>;
269
272
-torch::Tensor mm_dispatch(torch::Tensor const& A,
273
- torch::Tensor const& B, // already packed
274
- torch::Tensor const& group_scales, // already packed
275
- int64_t group_size,
276
- torch::Tensor const& channel_scales,
277
- torch::Tensor const& token_scales,
278
- std::optional<at::ScalarType> const& maybe_out_type,
279
- const std::string& schedule) {
270
+torch::stable::Tensor mm_dispatch(
271
+ torch::stable::Tensor const& A,
272
+ torch::stable::Tensor const& B, // already packed
273
+ torch::stable::Tensor const& group_scales, // already packed
274
+ int64_t group_size, torch::stable::Tensor const& channel_scales,
275
+ torch::stable::Tensor const& token_scales,
276
+ std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
277
+ const std::string& schedule) {
278
if (schedule == "256x128_1x1x1") {
279
return Kernel_256x128_1x1x1::mm(A, B, group_scales, group_size,
280
channel_scales, token_scales,
@@ -318,17 +316,18 @@ torch::Tensor mm_dispatch(torch::Tensor const& A,
316
channel_scales, token_scales,
317
maybe_out_type);
318
}
321
- TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
319
+ STD_TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
320
return {};
321
}
322
325
-torch::Tensor mm(torch::Tensor const& A,
326
- torch::Tensor const& B, // already packed
327
- torch::Tensor const& group_scales, // already packed
328
- int64_t group_size, torch::Tensor const& channel_scales,
329
- torch::Tensor const& token_scales,
330
- std::optional<at::ScalarType> const& maybe_out_type,
331
- std::optional<std::string> maybe_schedule) {
323
+torch::stable::Tensor mm(
324
+ torch::stable::Tensor const& A,
325
+ torch::stable::Tensor const& B, // already packed
326
+ torch::stable::Tensor const& group_scales, // already packed
327
+ int64_t group_size, torch::stable::Tensor const& channel_scales,
328
+ torch::stable::Tensor const& token_scales,
329
+ std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
330
+ std::optional<std::string> maybe_schedule) {
331
// requested a specific schedule
332
if (maybe_schedule) {
333
return mm_dispatch(A, B, group_scales, group_size, channel_scales,
@@ -378,14 +377,15 @@ torch::Tensor mm(torch::Tensor const& A,
377
// ----------------------------------------------------------------------------
378
// Pre-processing utils
379
// ----------------------------------------------------------------------------
381
-torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
382
- TORCH_CHECK(scales.dtype() == torch::kFloat8_e4m3fn);
383
- TORCH_CHECK(scales.is_contiguous());
384
- TORCH_CHECK(scales.is_cuda());
385
-
386
- auto packed_scales = torch::empty(
387
- {scales.numel() * ScalePackSize},
388
- torch::TensorOptions().dtype(scales.dtype()).device(scales.device()));
380
+torch::stable::Tensor pack_scale_fp8(torch::stable::Tensor const& scales) {
381
+ STD_TORCH_CHECK(scales.scalar_type() ==
382
+ torch::headeronly::ScalarType::Float8_e4m3fn);
383
+ STD_TORCH_CHECK(scales.is_contiguous());
384
+ STD_TORCH_CHECK(scales.is_cuda());
385
+
386
+ auto packed_scales =
387
+ torch::stable::empty({scales.numel() * ScalePackSize},
388
+ scales.scalar_type(), std::nullopt, scales.device());
389
auto scales_ptr = static_cast<MmaType const*>(scales.const_data_ptr());
390
auto packed_scales_ptr =
391
static_cast<cutlass::Array<ElementScale, ScalePackSize>*>(
@@ -396,15 +396,16 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
396
return packed_scales;
397
}
398
399
-torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
400
- TORCH_CHECK(B.dtype() == torch::kInt32);
401
- TORCH_CHECK(B.dim() == 2);
399
+torch::stable::Tensor encode_and_reorder_int4b(torch::stable::Tensor const& B) {
400
+ STD_TORCH_CHECK(B.scalar_type() == torch::headeronly::ScalarType::Int);
401
+ STD_TORCH_CHECK(B.dim() == 2);
402
403
- torch::Tensor B_packed = torch::empty_like(B);
403
+ torch::stable::Tensor B_packed = torch::stable::empty_like(B);
404
405
int k = B.size(0) * PackFactor; // logical k
406
int n = B.size(1);
407
- TORCH_CHECK((n * k) % 32 == 0, "need multiples of 32 int4s for 16B chunks");
407
+ STD_TORCH_CHECK((n * k) % 32 == 0,
408
+ "need multiples of 32 int4s for 16B chunks");
409
410
auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
411
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) {
416
417
bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(B_ptr, B_packed_ptr,
418
n * k);
418
- TORCH_CHECK(ok, "unified_encode_int4b failed");
419
+ STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
420
cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
421
422
return B_packed;
423
}
424
424
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
425
- m.impl("cutlass_w4a8_mm", &mm);
426
- m.impl("cutlass_pack_scale_fp8", &pack_scale_fp8);
427
- m.impl("cutlass_encode_and_reorder_int4b", &encode_and_reorder_int4b);
425
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
426
+ m.impl("cutlass_w4a8_mm", TORCH_BOX(&mm));
427
+ m.impl("cutlass_pack_scale_fp8", TORCH_BOX(&pack_scale_fp8));
428
+ m.impl("cutlass_encode_and_reorder_int4b",
429
+ TORCH_BOX(&encode_and_reorder_int4b));
430
}
431
430
-} // namespace vllm::cutlass_w4a8
\ No newline at end of file
432
+} // namespace vllm::cutlass_w4a8
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu
renamed
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh
renamed
csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
renamed
+18
-16
@@ -14,16 +14,15 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
17
+#include <torch/csrc/stable/tensor.h>
18
+#include "libtorch_stable/torch_utils.h"
19
+#include "libtorch_stable/dispatch_utils.h"
20
+#include "cuda_vec_utils.cuh"
21
22
#include <cuda_runtime_api.h>
23
#include <cuda_runtime.h>
24
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
24
-
25
#include <cuda_fp8.h>
26
-#include "dispatch_utils.h"
26
27
#include "cuda_utils.h"
28
#include "launch_bounds_utils.h"
@@ -118,17 +117,19 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
117
118
} // namespace vllm
119
121
-void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
122
- torch::Tensor& output_sf,
123
- torch::Tensor& input, // [..., 2 * d]
124
- torch::Tensor& input_sf) {
120
+void silu_and_mul_nvfp4_quant_sm1xxa(
121
+ torch::stable::Tensor& output, // [..., d]
122
+ torch::stable::Tensor& output_sf,
123
+ torch::stable::Tensor& input, // [..., 2 * d]
124
+ torch::stable::Tensor& input_sf) {
125
int32_t m = input.size(0);
126
int32_t n = input.size(1) / 2;
127
128
- TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
129
- TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
130
- input.scalar_type() == at::ScalarType::BFloat16,
131
- "Unsupported input data type for quantize_to_fp4.");
128
+ STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
129
+ STD_TORCH_CHECK(
130
+ input.scalar_type() == torch::headeronly::ScalarType::Half ||
131
+ input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
132
+ "Unsupported input data type for quantize_to_fp4.");
133
134
int multiProcessorCount =
135
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -136,8 +137,9 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
137
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
138
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
139
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
139
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
140
- auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
140
+ const torch::stable::accelerator::DeviceGuard device_guard(
141
+ input.get_device_index());
142
+ auto stream = get_current_cuda_stream(input.get_device_index());
143
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
144
int const numBlocksPerSM =
145
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]
151
int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
152
dim3 grid(grid_x, grid_y);
153
152
- VLLM_DISPATCH_HALF_TYPES(
154
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
155
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
156
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
157
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
renamed
+199
-132
@@ -14,14 +14,12 @@
14
* limitations under the License.
15
*/
16
17
-#include "core/registration.h"
17
+#include <torch/csrc/stable/library.h>
18
+#include <torch/csrc/stable/tensor.h>
19
+#include "libtorch_stable/torch_utils.h"
20
19
-#include <torch/all.h>
21
#include <cutlass/arch/arch.h>
22
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
24
-#include <c10/cuda/CUDAStream.h>
23
#include "cutlass_extensions/common.hpp"
24
25
#include "cute/tensor.hpp"
@@ -122,7 +120,7 @@ __global__ void __get_group_gemm_starts(
120
#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(ELEMENT_AB_TYPE, SF_TYPE, \
121
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
122
LayoutSFB, ScaleConfig) \
125
- else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
123
+ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
124
__get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, \
125
LayoutSFA, LayoutSFB, ScaleConfig> \
126
<<<1, num_experts, 0, stream>>>( \
@@ -150,50 +148,64 @@ __global__ void __get_group_gemm_starts(
148
}
149
150
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
153
-void run_get_group_gemm_starts(
154
- const torch::Tensor& a_starts, const torch::Tensor& b_starts,
155
- const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts,
156
- const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts,
157
- const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
158
- const torch::Tensor& a_strides, const torch::Tensor& b_strides,
159
- const torch::Tensor& c_strides, int64_t a_stride_val, int64_t b_stride_val,
160
- int64_t c_stride_val,
161
- /*these are used for their base addresses*/
162
- torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
163
- torch::Tensor const& out_tensors, torch::Tensor const& a_scales,
164
- torch::Tensor const& b_scales, torch::Tensor const& alphas,
165
- torch::Tensor const& expert_offsets, torch::Tensor const& sf_offsets,
166
- torch::Tensor const& problem_sizes, int M, int N, int K) {
151
+void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts,
152
+ const torch::stable::Tensor& b_starts,
153
+ const torch::stable::Tensor& out_starts,
154
+ const torch::stable::Tensor& a_scales_starts,
155
+ const torch::stable::Tensor& b_scales_starts,
156
+ const torch::stable::Tensor& alpha_starts,
157
+ const torch::stable::Tensor& layout_sfa,
158
+ const torch::stable::Tensor& layout_sfb,
159
+ const torch::stable::Tensor& a_strides,
160
+ const torch::stable::Tensor& b_strides,
161
+ const torch::stable::Tensor& c_strides,
162
+ int64_t a_stride_val, int64_t b_stride_val,
163
+ int64_t c_stride_val,
164
+ /*these are used for their base addresses*/
165
+ torch::stable::Tensor const& a_tensors,
166
+ torch::stable::Tensor const& b_tensors,
167
+ torch::stable::Tensor const& out_tensors,
168
+ torch::stable::Tensor const& a_scales,
169
+ torch::stable::Tensor const& b_scales,
170
+ torch::stable::Tensor const& alphas,
171
+ torch::stable::Tensor const& expert_offsets,
172
+ torch::stable::Tensor const& sf_offsets,
173
+ torch::stable::Tensor const& problem_sizes,
174
+ int M, int N, int K) {
175
int num_experts = (int)expert_offsets.size(0);
168
- auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
176
+ auto stream = get_current_cuda_stream(a_tensors.get_device_index());
177
170
- TORCH_CHECK(out_tensors.size(1) == N,
171
- "Output tensor shape doesn't match expected shape");
172
- TORCH_CHECK(K / 2 == b_tensors.size(2),
173
- "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
174
- " dimension must match");
178
+ STD_TORCH_CHECK(out_tensors.size(1) == N,
179
+ "Output tensor shape doesn't match expected shape");
180
+ STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
181
+ "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
182
+ " dimension must match");
183
if (false) {
184
}
185
//(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB,
186
// ScaleConfig)
187
__CALL_GET_STARTS_KERNEL_BLOCKSCALE(
180
- cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kBFloat16,
181
- cutlass::bfloat16_t, LayoutSFA, LayoutSFB, ScaleConfig)
188
+ cutlass::float_e2m1_t, cutlass::float_ue4m3_t,
189
+ torch::headeronly::ScalarType::BFloat16, cutlass::bfloat16_t, LayoutSFA,
190
+ LayoutSFB, ScaleConfig)
191
__CALL_GET_STARTS_KERNEL_BLOCKSCALE(cutlass::float_e2m1_t,
183
- cutlass::float_ue4m3_t, torch::kFloat16,
184
- half, LayoutSFA, LayoutSFB, ScaleConfig)
192
+ cutlass::float_ue4m3_t,
193
+ torch::headeronly::ScalarType::Half, half,
194
+ LayoutSFA, LayoutSFB, ScaleConfig)
195
else {
186
- TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
196
+ STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
197
}
198
}
199
200
template <typename OutType>
201
void run_fp4_blockwise_scaled_group_mm_sm100(
192
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
193
- const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
194
- const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
195
- const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
196
- int N, int K) {
202
+ torch::stable::Tensor& output, const torch::stable::Tensor& a,
203
+ const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
204
+ const torch::stable::Tensor& b_blockscales,
205
+ const torch::stable::Tensor& alphas,
206
+ const torch::stable::Tensor& problem_sizes,
207
+ const torch::stable::Tensor& expert_offsets,
208
+ const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
209
using ProblemShape =
210
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
211
using ElementType = cutlass::float_e2m1_t;
@@ -272,20 +284,40 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
284
285
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
286
int num_experts = static_cast<int>(expert_offsets.size(0));
275
- auto options_int =
276
- torch::TensorOptions().dtype(torch::kInt64).device(a.device());
277
-
278
- torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
279
- torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
280
- torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
281
- torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
282
- torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
283
- torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
284
- torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
285
- torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
286
- torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
287
- torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
288
- torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
287
+
288
+ torch::stable::Tensor a_ptrs =
289
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
290
+ std::nullopt, a.device());
291
+ torch::stable::Tensor b_ptrs =
292
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
293
+ std::nullopt, a.device());
294
+ torch::stable::Tensor out_ptrs =
295
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
296
+ std::nullopt, a.device());
297
+ torch::stable::Tensor a_scales_ptrs =
298
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
299
+ std::nullopt, a.device());
300
+ torch::stable::Tensor b_scales_ptrs =
301
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
302
+ std::nullopt, a.device());
303
+ torch::stable::Tensor alpha_ptrs =
304
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
305
+ std::nullopt, a.device());
306
+ torch::stable::Tensor layout_sfa = torch::stable::empty(
307
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
308
+ a.device());
309
+ torch::stable::Tensor layout_sfb = torch::stable::empty(
310
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
311
+ a.device());
312
+ torch::stable::Tensor a_strides1 =
313
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
314
+ std::nullopt, a.device());
315
+ torch::stable::Tensor b_strides1 =
316
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
317
+ std::nullopt, a.device());
318
+ torch::stable::Tensor c_strides1 =
319
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
320
+ std::nullopt, a.device());
321
322
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
323
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(
340
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
341
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
342
scheduler.raster_order = RasterOrderOptions::AlongM;
311
- hw_info.device_id = a.get_device();
343
+ hw_info.device_id = a.get_device_index();
344
static std::unordered_map<int, int> cached_sm_counts;
345
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
346
cached_sm_counts[hw_info.device_id] =
@@ -350,32 +382,35 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
382
scheduler};
383
384
size_t workspace_size = Gemm::get_workspace_size(args);
353
- auto const workspace_options =
354
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
355
- auto workspace = torch::empty(workspace_size, workspace_options);
356
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
385
+ auto workspace =
386
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
387
+ std::nullopt, a.device());
388
+ const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
389
390
auto can_implement_status = gemm_op.can_implement(args);
359
- TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
360
- "Failed to implement GEMM: status=", (int)can_implement_status);
391
+ STD_TORCH_CHECK(
392
+ can_implement_status == cutlass::Status::kSuccess,
393
+ "Failed to implement GEMM: status=", (int)can_implement_status);
394
395
// Run the GEMM
396
auto status = gemm_op.initialize(args, workspace.data_ptr());
364
- TORCH_CHECK(status == cutlass::Status::kSuccess,
365
- "Failed to initialize GEMM: status=", (int)status,
366
- " workspace_size=", workspace_size, " num_experts=", num_experts,
367
- " M=", M, " N=", N, " K=", K);
397
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
398
+ "Failed to initialize GEMM: status=", (int)status,
399
+ " workspace_size=", workspace_size,
400
+ " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
401
402
status = gemm_op.run(args, workspace.data_ptr(), stream);
370
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
403
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
404
}
405
406
void run_fp4_blockwise_scaled_group_mm_sm120(
374
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
375
- const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
376
- const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
377
- const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
378
- int N, int K) {
407
+ torch::stable::Tensor& output, const torch::stable::Tensor& a,
408
+ const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
409
+ const torch::stable::Tensor& b_blockscales,
410
+ const torch::stable::Tensor& alphas,
411
+ const torch::stable::Tensor& problem_sizes,
412
+ const torch::stable::Tensor& expert_offsets,
413
+ const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
414
using ProblemShape =
415
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
416
using ElementType = cutlass::float_e2m1_t;
@@ -446,20 +481,40 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
481
482
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
483
int num_experts = static_cast<int>(expert_offsets.size(0));
449
- auto options_int =
450
- torch::TensorOptions().dtype(torch::kInt64).device(a.device());
451
-
452
- torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
453
- torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
454
- torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
455
- torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
456
- torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
457
- torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
458
- torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
459
- torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
460
- torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
461
- torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
462
- torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
484
+
485
+ torch::stable::Tensor a_ptrs =
486
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
487
+ std::nullopt, a.device());
488
+ torch::stable::Tensor b_ptrs =
489
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
490
+ std::nullopt, a.device());
491
+ torch::stable::Tensor out_ptrs =
492
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
493
+ std::nullopt, a.device());
494
+ torch::stable::Tensor a_scales_ptrs =
495
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
496
+ std::nullopt, a.device());
497
+ torch::stable::Tensor b_scales_ptrs =
498
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
499
+ std::nullopt, a.device());
500
+ torch::stable::Tensor alpha_ptrs =
501
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
502
+ std::nullopt, a.device());
503
+ torch::stable::Tensor layout_sfa = torch::stable::empty(
504
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
505
+ a.device());
506
+ torch::stable::Tensor layout_sfb = torch::stable::empty(
507
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
508
+ a.device());
509
+ torch::stable::Tensor a_strides1 =
510
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
511
+ std::nullopt, a.device());
512
+ torch::stable::Tensor b_strides1 =
513
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
514
+ std::nullopt, a.device());
515
+ torch::stable::Tensor c_strides1 =
516
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
517
+ std::nullopt, a.device());
518
519
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
520
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(
535
using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
536
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
537
scheduler.raster_order = RasterOrderOptions::AlongM;
483
- hw_info.device_id = a.get_device();
538
+ hw_info.device_id = a.get_device_index();
539
static std::unordered_map<int, int> cached_sm_counts;
540
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
541
cached_sm_counts[hw_info.device_id] =
@@ -523,33 +578,36 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
578
scheduler};
579
580
size_t workspace_size = Gemm::get_workspace_size(args);
526
- auto const workspace_options =
527
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
528
- auto workspace = torch::empty(workspace_size, workspace_options);
529
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
581
+ auto workspace =
582
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
583
+ std::nullopt, a.device());
584
+ const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
585
586
auto can_implement_status = gemm_op.can_implement(args);
532
- TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
533
- "Failed to implement GEMM: status=", (int)can_implement_status);
587
+ STD_TORCH_CHECK(
588
+ can_implement_status == cutlass::Status::kSuccess,
589
+ "Failed to implement GEMM: status=", (int)can_implement_status);
590
591
// Run the GEMM
592
auto status = gemm_op.initialize(args, workspace.data_ptr());
537
- TORCH_CHECK(status == cutlass::Status::kSuccess,
538
- "Failed to initialize GEMM: status=", (int)status,
539
- " workspace_size=", workspace_size, " num_experts=", num_experts,
540
- " M=", M, " N=", N, " K=", K);
593
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
594
+ "Failed to initialize GEMM: status=", (int)status,
595
+ " workspace_size=", workspace_size,
596
+ " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
597
598
status = gemm_op.run(args, workspace.data_ptr(), stream);
543
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
599
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
600
}
601
602
template <typename OutType>
603
void run_fp4_blockwise_scaled_group_mm(
548
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
549
- const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
550
- const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
551
- const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
552
- int N, int K) {
604
+ torch::stable::Tensor& output, const torch::stable::Tensor& a,
605
+ const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
606
+ const torch::stable::Tensor& b_blockscales,
607
+ const torch::stable::Tensor& alphas,
608
+ const torch::stable::Tensor& problem_sizes,
609
+ const torch::stable::Tensor& expert_offsets,
610
+ const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
611
int32_t version_num = get_sm_version_num();
612
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
613
if (version_num >= 120 && version_num < 130) {
@@ -567,7 +625,7 @@ void run_fp4_blockwise_scaled_group_mm(
625
return;
626
}
627
#endif
570
- TORCH_CHECK_NOT_IMPLEMENTED(
628
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
629
false,
630
"No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ",
631
version_num, ". Required capability: 100 or 120");
@@ -575,26 +633,31 @@ void run_fp4_blockwise_scaled_group_mm(
633
634
#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
635
(defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
578
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
579
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
636
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
637
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
638
#endif
639
582
-#define CHECK_TYPE(x, st, m) \
583
- TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
640
+#define CHECK_TYPE(x, st, m) \
641
+ STD_TORCH_CHECK(x.scalar_type() == st, \
642
+ ": Inconsistency of torch::stable::Tensor type:", m)
643
#define CHECK_TH_CUDA(x, m) \
585
- TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
644
+ STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
645
#define CHECK_CONTIGUOUS(x, m) \
587
- TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
646
+ STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
647
#define CHECK_INPUT(x, st, m) \
648
CHECK_TH_CUDA(x, m); \
649
CHECK_CONTIGUOUS(x, m); \
650
CHECK_TYPE(x, st, m)
651
593
-void cutlass_fp4_group_mm(
594
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
595
- const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
596
- const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
597
- const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
652
+void cutlass_fp4_group_mm(torch::stable::Tensor& output,
653
+ const torch::stable::Tensor& a,
654
+ const torch::stable::Tensor& b,
655
+ const torch::stable::Tensor& a_blockscale,
656
+ const torch::stable::Tensor& b_blockscales,
657
+ const torch::stable::Tensor& alphas,
658
+ const torch::stable::Tensor& problem_sizes,
659
+ const torch::stable::Tensor& expert_offsets,
660
+ const torch::stable::Tensor& sf_offsets) {
661
#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
662
(defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
663
// Input validation
@@ -602,30 +665,34 @@ void cutlass_fp4_group_mm(
665
CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
666
CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale");
667
CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales");
605
- CHECK_INPUT(alphas, at::ScalarType::Float, "alphas");
606
-
607
- TORCH_CHECK(a_blockscale.dim() == 2,
608
- "expected a_blockscale to be of shape [num_experts, rounded_m,"
609
- " k // group_size], observed rank: ",
610
- a_blockscale.dim())
611
- TORCH_CHECK(b_blockscales.dim() == 3,
612
- "expected b_blockscale to be of shape: "
613
- " [num_experts, n, k // group_size], observed rank: ",
614
- b_blockscales.dim())
615
- TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be a 2D tensor");
616
- TORCH_CHECK(problem_sizes.size(1) == 3,
617
- "problem_sizes must have the shape (num_experts, 3)");
618
- TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
619
- "Number of experts in problem_sizes must match expert_offsets");
620
- TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
621
- "problem_sizes must be int32.");
668
+ CHECK_INPUT(alphas, torch::headeronly::ScalarType::Float, "alphas");
669
+
670
+ STD_TORCH_CHECK(
671
+ a_blockscale.dim() == 2,
672
+ "expected a_blockscale to be of shape [num_experts, rounded_m,"
673
+ " k // group_size], observed rank: ",
674
+ a_blockscale.dim())
675
+ STD_TORCH_CHECK(b_blockscales.dim() == 3,
676
+ "expected b_blockscale to be of shape: "
677
+ " [num_experts, n, k // group_size], observed rank: ",
678
+ b_blockscales.dim())
679
+ STD_TORCH_CHECK(problem_sizes.dim() == 2,
680
+ "problem_sizes must be a 2D tensor");
681
+ STD_TORCH_CHECK(problem_sizes.size(1) == 3,
682
+ "problem_sizes must have the shape (num_experts, 3)");
683
+ STD_TORCH_CHECK(
684
+ problem_sizes.size(0) == expert_offsets.size(0),
685
+ "Number of experts in problem_sizes must match expert_offsets");
686
+ STD_TORCH_CHECK(
687
+ problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
688
+ "problem_sizes must be int32.");
689
690
int M = static_cast<int>(a.size(0));
691
int N = static_cast<int>(b.size(1));
692
int E = static_cast<int>(b.size(0));
693
int K = static_cast<int>(2 * b.size(2));
694
628
- if (output.scalar_type() == torch::kBFloat16) {
695
+ if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
696
run_fp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
697
output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
698
expert_offsets, sf_offsets, M, N, K);
@@ -633,7 +700,7 @@ void cutlass_fp4_group_mm(
700
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
701
int32_t version_num = get_sm_version_num();
702
if (version_num >= 120 && version_num < 130) {
636
- TORCH_CHECK_NOT_IMPLEMENTED(
703
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
704
false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ",
705
output.scalar_type());
706
}
@@ -643,7 +710,7 @@ void cutlass_fp4_group_mm(
710
expert_offsets, sf_offsets, M, N, K);
711
}
712
#else
646
- TORCH_CHECK_NOT_IMPLEMENTED(
713
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
714
false,
715
"No compiled cutlass_fp4_group_mm kernel, vLLM must "
716
"be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 "
@@ -651,6 +718,6 @@ void cutlass_fp4_group_mm(
718
#endif
719
}
720
654
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
655
- m.impl("cutlass_fp4_group_mm", &cutlass_fp4_group_mm);
721
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
722
+ m.impl("cutlass_fp4_group_mm", TORCH_BOX(&cutlass_fp4_group_mm));
723
}
csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
renamed
+56
-52
@@ -14,16 +14,15 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
17
+#include <torch/csrc/stable/tensor.h>
18
+#include "libtorch_stable/torch_utils.h"
19
+#include "libtorch_stable/dispatch_utils.h"
20
+#include "cuda_vec_utils.cuh"
21
22
#include <cuda_runtime_api.h>
23
#include <cuda_runtime.h>
24
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
24
-
25
#include <cuda_fp8.h>
26
-#include "dispatch_utils.h"
26
27
#include "cuda_utils.h"
28
#include "nvfp4_utils.cuh"
@@ -327,25 +326,28 @@ void quant_impl(void* output, void* output_scale, void* input,
326
} // namespace vllm
327
328
/*Quantization entry for fp4 experts quantization*/
330
-#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
329
+#define CHECK_TH_CUDA(x, m) \
330
+ STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
331
#define CHECK_CONTIGUOUS(x, m) \
332
- TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
332
+ STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
333
#define CHECK_INPUT(x, m) \
334
CHECK_TH_CUDA(x, m); \
335
CHECK_CONTIGUOUS(x, m);
336
337
-constexpr auto HALF = at::ScalarType::Half;
338
-constexpr auto BF16 = at::ScalarType::BFloat16;
339
-constexpr auto FLOAT = at::ScalarType::Float;
340
-constexpr auto INT = at::ScalarType::Int;
341
-constexpr auto UINT8 = at::ScalarType::Byte;
337
+constexpr auto HALF = torch::headeronly::ScalarType::Half;
338
+constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
339
+constexpr auto FLOAT = torch::headeronly::ScalarType::Float;
340
+constexpr auto INT = torch::headeronly::ScalarType::Int;
341
+constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
342
343
// Common validation for fp4 experts quantization entry points.
344
static void validate_fp4_experts_quant_inputs(
345
- torch::Tensor const& output, torch::Tensor const& output_scale,
346
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
347
- torch::Tensor const& input_offset_by_experts,
348
- torch::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
345
+ torch::stable::Tensor const& output,
346
+ torch::stable::Tensor const& output_scale,
347
+ torch::stable::Tensor const& input,
348
+ torch::stable::Tensor const& input_global_scale,
349
+ torch::stable::Tensor const& input_offset_by_experts,
350
+ torch::stable::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
351
int64_t k) {
352
CHECK_INPUT(output, "output");
353
CHECK_INPUT(output_scale, "output_scale");
@@ -354,41 +356,42 @@ static void validate_fp4_experts_quant_inputs(
356
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
357
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
358
357
- TORCH_CHECK(output.dim() == 2);
358
- TORCH_CHECK(output_scale.dim() == 2);
359
- TORCH_CHECK(input.dim() == 2);
360
- TORCH_CHECK(input_global_scale.dim() == 1);
361
- TORCH_CHECK(input_offset_by_experts.dim() == 1);
362
- TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
363
-
364
- TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
365
- TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
366
- TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
367
- TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
359
+ STD_TORCH_CHECK(output.dim() == 2);
360
+ STD_TORCH_CHECK(output_scale.dim() == 2);
361
+ STD_TORCH_CHECK(input.dim() == 2);
362
+ STD_TORCH_CHECK(input_global_scale.dim() == 1);
363
+ STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
364
+ STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
365
+
366
+ STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
367
+ STD_TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
368
+ STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
369
+ STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
370
// output is uint8 (two nvfp4 values are packed into one uint8)
371
// output_scale is int32 (four fp8 values are packed into one int32)
370
- TORCH_CHECK(output.scalar_type() == UINT8);
371
- TORCH_CHECK(output_scale.scalar_type() == INT);
372
+ STD_TORCH_CHECK(output.scalar_type() == UINT8);
373
+ STD_TORCH_CHECK(output_scale.scalar_type() == INT);
374
375
const int BLOCK_SIZE = 16;
374
- TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
376
+ STD_TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
377
auto n_experts = input_global_scale.size(0);
376
- TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
377
- TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
378
- TORCH_CHECK(output.size(0) == m_topk);
379
- TORCH_CHECK(output.size(1) == k / 2);
378
+ STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
379
+ STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
380
+ STD_TORCH_CHECK(output.size(0) == m_topk);
381
+ STD_TORCH_CHECK(output.size(1) == k / 2);
382
int scales_k = k / BLOCK_SIZE;
383
// 4 means the swizzle requirement by nvidia nvfp4.
384
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
385
// 4 means 4 fp8 values are packed into one int32
384
- TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
386
+ STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
387
}
388
389
void scaled_fp4_experts_quant_sm1xxa(
388
- torch::Tensor& output, torch::Tensor& output_scale,
389
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
390
- torch::Tensor const& input_offset_by_experts,
391
- torch::Tensor const& output_scale_offset_by_experts) {
390
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
391
+ torch::stable::Tensor const& input,
392
+ torch::stable::Tensor const& input_global_scale,
393
+ torch::stable::Tensor const& input_offset_by_experts,
394
+ torch::stable::Tensor const& output_scale_offset_by_experts) {
395
auto m_topk = input.size(0);
396
auto k = input.size(1);
397
@@ -397,11 +400,11 @@ void scaled_fp4_experts_quant_sm1xxa(
400
output_scale_offset_by_experts, m_topk, k);
401
402
auto n_experts = input_global_scale.size(0);
400
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
401
- const cudaStream_t stream =
402
- at::cuda::getCurrentCUDAStream(input.get_device());
403
+ const torch::stable::accelerator::DeviceGuard device_guard(
404
+ input.get_device_index());
405
+ const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
406
404
- VLLM_DISPATCH_HALF_TYPES(
407
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
408
input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
409
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
410
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
@@ -413,14 +416,15 @@ void scaled_fp4_experts_quant_sm1xxa(
416
}
417
418
void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
416
- torch::Tensor& output, torch::Tensor& output_scale,
417
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
418
- torch::Tensor const& input_offset_by_experts,
419
- torch::Tensor const& output_scale_offset_by_experts) {
419
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
420
+ torch::stable::Tensor const& input,
421
+ torch::stable::Tensor const& input_global_scale,
422
+ torch::stable::Tensor const& input_offset_by_experts,
423
+ torch::stable::Tensor const& output_scale_offset_by_experts) {
424
auto m_topk = input.size(0);
425
// Input has gate || up layout, so k = input.size(1) / 2
426
auto k_times_2 = input.size(1);
423
- TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
427
+ STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
428
auto k = k_times_2 / 2;
429
430
validate_fp4_experts_quant_inputs(output, output_scale, input,
@@ -428,11 +432,11 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
432
output_scale_offset_by_experts, m_topk, k);
433
434
auto n_experts = input_global_scale.size(0);
431
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
432
- const cudaStream_t stream =
433
- at::cuda::getCurrentCUDAStream(input.get_device());
435
+ const torch::stable::accelerator::DeviceGuard device_guard(
436
+ input.get_device_index());
437
+ const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
438
435
- VLLM_DISPATCH_HALF_TYPES(
439
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
440
input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] {
441
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
442
vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu
new
+172
@@ -0,0 +1,172 @@
1
+/*
2
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
+ *
4
+ * Licensed under the Apache License, Version 2.0 (the "License");
5
+ * you may not use this file except in compliance with the License.
6
+ * You may obtain a copy of the License at
7
+ *
8
+ * http://www.apache.org/licenses/LICENSE-2.0
9
+ *
10
+ * Unless required by applicable law or agreed to in writing, software
11
+ * distributed under the License is distributed on an "AS IS" BASIS,
12
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ * See the License for the specific language governing permissions and
14
+ * limitations under the License.
15
+ */
16
+
17
+#include <torch/csrc/stable/tensor.h>
18
+
19
+#include "libtorch_stable/torch_utils.h"
20
+
21
+#include "cutlass_extensions/common.hpp"
22
+#include "nvfp4_utils.cuh"
23
+
24
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
25
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
26
+void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
27
+ torch::stable::Tensor const& input,
28
+ torch::stable::Tensor const& output_sf,
29
+ torch::stable::Tensor const& input_sf,
30
+ bool is_sf_swizzled_layout);
31
+#endif
32
+
33
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
34
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
35
+void scaled_fp4_experts_quant_sm1xxa(
36
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
37
+ torch::stable::Tensor const& input,
38
+ torch::stable::Tensor const& input_global_scale,
39
+ torch::stable::Tensor const& input_offset_by_experts,
40
+ torch::stable::Tensor const& output_scale_offset_by_experts);
41
+#endif
42
+
43
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
44
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
45
+void silu_and_mul_nvfp4_quant_sm1xxa(torch::stable::Tensor& output,
46
+ torch::stable::Tensor& output_sf,
47
+ torch::stable::Tensor& input,
48
+ torch::stable::Tensor& input_sf);
49
+#endif
50
+
51
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
52
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
53
+void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
54
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
55
+ torch::stable::Tensor const& input,
56
+ torch::stable::Tensor const& input_global_scale,
57
+ torch::stable::Tensor const& input_offset_by_experts,
58
+ torch::stable::Tensor const& output_scale_offset_by_experts);
59
+#endif
60
+
61
+static bool nvfp4_quant_sm_supported() {
62
+ const int32_t sm = get_sm_version_num();
63
+#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
64
+ if (sm >= 100 && sm < 120) return true;
65
+#endif
66
+#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120
67
+ if (sm >= 120 && sm < 130) return true;
68
+#endif
69
+ return false;
70
+}
71
+
72
+void scaled_fp4_quant_out(torch::stable::Tensor const& input,
73
+ torch::stable::Tensor const& input_sf,
74
+ bool is_sf_swizzled_layout,
75
+ torch::stable::Tensor& output,
76
+ torch::stable::Tensor& output_sf) {
77
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
78
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
79
+ STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
80
+ "No compiled nvfp4 quantization kernel for SM ",
81
+ get_sm_version_num(),
82
+ ". Recompile with the appropriate CUDA arch.");
83
+ return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
84
+ is_sf_swizzled_layout);
85
+#endif
86
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
87
+ "No compiled nvfp4 quantization kernel");
88
+}
89
+
90
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
91
+ torch::stable::Tensor const& input, torch::stable::Tensor const& input_sf,
92
+ bool is_sf_swizzled_layout) {
93
+ int64_t n = input.size(-1);
94
+ int64_t m = input.numel() / n;
95
+ auto device = input.device();
96
+
97
+ // Two fp4 values packed into a uint8
98
+ auto output = torch::stable::empty(
99
+ {m, n / 2}, torch::headeronly::ScalarType::Byte, std::nullopt, device);
100
+
101
+ torch::stable::Tensor output_sf;
102
+ if (is_sf_swizzled_layout) {
103
+ auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
104
+ output_sf = torch::stable::empty(
105
+ {sf_m, sf_n}, torch::headeronly::ScalarType::Int, std::nullopt, device);
106
+ } else {
107
+ output_sf = torch::stable::empty({m, n / CVT_FP4_SF_VEC_SIZE},
108
+ torch::headeronly::ScalarType::Byte,
109
+ std::nullopt, device);
110
+ }
111
+
112
+ scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
113
+ output_sf);
114
+ return {output, output_sf};
115
+}
116
+
117
+void scaled_fp4_experts_quant(
118
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
119
+ torch::stable::Tensor const& input,
120
+ torch::stable::Tensor const& input_global_scale,
121
+ torch::stable::Tensor const& input_offset_by_experts,
122
+ torch::stable::Tensor const& output_scale_offset_by_experts) {
123
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
124
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
125
+ STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
126
+ "No compiled nvfp4 experts quantization kernel for SM ",
127
+ get_sm_version_num(),
128
+ ". Recompile with the appropriate CUDA arch.");
129
+ return scaled_fp4_experts_quant_sm1xxa(
130
+ output, output_scale, input, input_global_scale, input_offset_by_experts,
131
+ output_scale_offset_by_experts);
132
+#endif
133
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
134
+ false, "No compiled nvfp4 experts quantization kernel");
135
+}
136
+
137
+void silu_and_mul_nvfp4_quant(torch::stable::Tensor& output,
138
+ torch::stable::Tensor& output_sf,
139
+ torch::stable::Tensor& input,
140
+ torch::stable::Tensor& input_sf) {
141
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
142
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
143
+ STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
144
+ "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
145
+ get_sm_version_num(),
146
+ ". Recompile with the appropriate CUDA arch.");
147
+ return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
148
+#endif
149
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
150
+ false, "No compiled silu_and_mul nvfp4 quantization kernel");
151
+}
152
+
153
+void silu_and_mul_scaled_fp4_experts_quant(
154
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
155
+ torch::stable::Tensor const& input,
156
+ torch::stable::Tensor const& input_global_scale,
157
+ torch::stable::Tensor const& input_offset_by_experts,
158
+ torch::stable::Tensor const& output_scale_offset_by_experts) {
159
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
160
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
161
+ STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
162
+ "No compiled silu_and_mul nvfp4 experts quantization kernel "
163
+ "for SM ",
164
+ get_sm_version_num(),
165
+ ". Recompile with the appropriate CUDA arch.");
166
+ return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
167
+ output, output_scale, input, input_global_scale, input_offset_by_experts,
168
+ output_scale_offset_by_experts);
169
+#endif
170
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
171
+ false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
172
+}
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
renamed
+36
-34
@@ -14,16 +14,16 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
17
+#include <torch/csrc/stable/tensor.h>
18
19
#include <cuda_runtime_api.h>
20
#include <cuda_runtime.h>
21
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
24
-
22
#include <cuda_fp8.h>
26
-#include "dispatch_utils.h"
23
+
24
+#include "libtorch_stable/torch_utils.h"
25
+#include "libtorch_stable/dispatch_utils.h"
26
+#include "cuda_vec_utils.cuh"
27
28
#include "cuda_utils.h"
29
#include "launch_bounds_utils.h"
@@ -173,18 +173,19 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
173
174
} // namespace vllm
175
176
-void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
177
- torch::Tensor const& input,
178
- torch::Tensor const& output_sf,
179
- torch::Tensor const& input_sf,
176
+void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
177
+ torch::stable::Tensor const& input,
178
+ torch::stable::Tensor const& output_sf,
179
+ torch::stable::Tensor const& input_sf,
180
bool is_sf_swizzled_layout) {
181
int32_t m = input.size(0);
182
int32_t n = input.size(1);
183
184
- TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
185
- TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
186
- input.scalar_type() == at::ScalarType::BFloat16,
187
- "Unsupported input data type for quantize_to_fp4.");
184
+ STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
185
+ STD_TORCH_CHECK(
186
+ input.scalar_type() == torch::headeronly::ScalarType::Half ||
187
+ input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
188
+ "Unsupported input data type for quantize_to_fp4.");
189
190
int multiProcessorCount =
191
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -192,8 +193,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
193
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
194
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
195
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
195
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
196
- auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
196
+ const torch::stable::accelerator::DeviceGuard device_guard(
197
+ input.get_device_index());
198
+ auto stream = get_current_cuda_stream(input.get_device_index());
199
200
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
201
@@ -213,15 +215,15 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
215
std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
216
dim3 grid(grid_x, grid_y);
217
216
- VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
217
- using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
218
- auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
219
- // NOTE: We don't support e8m0 scales at this moment.
220
- vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
221
- m, n, num_padded_cols, input_ptr, input_sf_ptr,
222
- reinterpret_cast<uint32_t*>(output_ptr),
223
- reinterpret_cast<uint32_t*>(sf_out));
224
- });
218
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
219
+ input.scalar_type(), "nvfp4_quant_kernel", [&] {
220
+ using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
221
+ auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
222
+ vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
223
+ m, n, num_padded_cols, input_ptr, input_sf_ptr,
224
+ reinterpret_cast<uint32_t*>(output_ptr),
225
+ reinterpret_cast<uint32_t*>(sf_out));
226
+ });
227
} else {
228
int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD;
229
int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
@@ -229,15 +231,15 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
231
m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
232
dim3 grid(grid_x, grid_y);
233
232
- VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
233
- using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
234
- auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
235
- // NOTE: We don't support e8m0 scales at this moment.
236
- vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
237
- <<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, num_packed_cols,
238
- input_ptr, input_sf_ptr,
239
- reinterpret_cast<uint32_t*>(output_ptr),
240
- reinterpret_cast<uint32_t*>(sf_out));
241
- });
234
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
235
+ input.scalar_type(), "nvfp4_quant_kernel", [&] {
236
+ using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
237
+ auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
238
+ vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
239
+ <<<grid, block, 0, stream>>>(
240
+ m, n, sf_n_unpadded, num_packed_cols, input_ptr, input_sf_ptr,
241
+ reinterpret_cast<uint32_t*>(output_ptr),
242
+ reinterpret_cast<uint32_t*>(sf_out));
243
+ });
244
}
245
}
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu
renamed
+28
-20
@@ -14,32 +14,39 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
18
-#include <c10/cuda/CUDAGuard.h>
17
+#include <torch/csrc/stable/tensor.h>
18
+
19
+#include "libtorch_stable/torch_utils.h"
20
+
21
#include "cutlass_extensions/common.hpp"
22
23
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
22
-void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
23
- torch::Tensor const& B,
24
- torch::Tensor const& A_sf,
25
- torch::Tensor const& B_sf,
26
- torch::Tensor const& alpha);
24
+void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
25
+ torch::stable::Tensor const& A,
26
+ torch::stable::Tensor const& B,
27
+ torch::stable::Tensor const& A_sf,
28
+ torch::stable::Tensor const& B_sf,
29
+ torch::stable::Tensor const& alpha);
30
#endif
31
32
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
30
-void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
31
- torch::Tensor const& B,
32
- torch::Tensor const& A_sf,
33
- torch::Tensor const& B_sf,
34
- torch::Tensor const& alpha);
33
+void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
34
+ torch::stable::Tensor const& A,
35
+ torch::stable::Tensor const& B,
36
+ torch::stable::Tensor const& A_sf,
37
+ torch::stable::Tensor const& B_sf,
38
+ torch::stable::Tensor const& alpha);
39
#endif
40
37
-void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
38
- const torch::Tensor& B, const torch::Tensor& A_sf,
39
- const torch::Tensor& B_sf,
40
- const torch::Tensor& alpha) {
41
- // Make sure we’re on A’s device.
42
- const c10::cuda::OptionalCUDAGuard device_guard(device_of(A));
41
+void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
42
+ const torch::stable::Tensor& A,
43
+ const torch::stable::Tensor& B,
44
+ const torch::stable::Tensor& A_sf,
45
+ const torch::stable::Tensor& B_sf,
46
+ const torch::stable::Tensor& alpha) {
47
+ // Make sure we're on A's device.
48
+ const torch::stable::accelerator::DeviceGuard device_guard(
49
+ A.get_device_index());
50
const int32_t sm = get_sm_version_num();
51
52
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
@@ -56,8 +63,9 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
63
}
64
#endif
65
59
- TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 mm kernel for SM ", sm,
60
- ". Recompile with CUDA >= 12.8 and CC >= 100.");
66
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
67
+ false, "No compiled nvfp4 mm kernel for SM ", sm,
68
+ ". Recompile with CUDA >= 12.8 and CC >= 100.");
69
}
70
71
bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability) {
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu
renamed
+82
-75
@@ -14,10 +14,9 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
17
+#include <torch/csrc/stable/tensor.h>
18
19
-#include <ATen/cuda/CUDAContext.h>
20
-#include <c10/cuda/CUDAGuard.h>
19
+#include "libtorch_stable/torch_utils.h"
20
21
#include "cutlass_extensions/common.hpp"
22
@@ -127,8 +126,9 @@ struct Fp4GemmSm100 {
126
127
template <typename Config>
128
typename Config::Gemm::Arguments args_from_options(
130
- at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
131
- at::Tensor const& A_sf, at::Tensor const& B_sf, at::Tensor const& alpha,
129
+ torch::stable::Tensor& D, torch::stable::Tensor const& A,
130
+ torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf,
131
+ torch::stable::Tensor const& B_sf, torch::stable::Tensor const& alpha,
132
int64_t M, int64_t N, int64_t K) {
133
using ElementA = typename Config::Gemm::ElementA;
134
using ElementB = typename Config::Gemm::ElementB;
@@ -174,19 +174,20 @@ typename Config::Gemm::Arguments args_from_options(
174
}
175
176
template <typename Config>
177
-void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
178
- at::Tensor const& A_sf, at::Tensor const& B_sf,
179
- at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
180
- cudaStream_t stream) {
177
+void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A,
178
+ torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf,
179
+ torch::stable::Tensor const& B_sf,
180
+ torch::stable::Tensor const& alpha, int64_t m, int64_t n,
181
+ int64_t k, cudaStream_t stream) {
182
typename Config::Gemm gemm;
183
184
auto arguments =
185
args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k);
186
187
size_t workspace_size = Config::Gemm::get_workspace_size(arguments);
187
- auto const workspace_options =
188
- torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
189
- auto workspace = torch::empty(workspace_size, workspace_options);
188
+ auto workspace =
189
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
190
+ std::nullopt, A.device());
191
192
CUTLASS_CHECK(gemm.can_implement(arguments));
193
@@ -197,12 +198,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
198
199
// Dispatch function to select appropriate config based on M
200
template <typename OutType>
200
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
201
- torch::Tensor const& B,
202
- torch::Tensor const& A_sf,
203
- torch::Tensor const& B_sf,
204
- torch::Tensor const& alpha, int64_t m, int64_t n,
205
- int64_t k, cudaStream_t stream) {
201
+void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
202
+ torch::stable::Tensor const& A,
203
+ torch::stable::Tensor const& B,
204
+ torch::stable::Tensor const& A_sf,
205
+ torch::stable::Tensor const& B_sf,
206
+ torch::stable::Tensor const& alpha, int64_t m,
207
+ int64_t n, int64_t k, cudaStream_t stream) {
208
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
209
210
if (mp2 <= 16) {
@@ -222,61 +224,65 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
224
225
#else
226
template <typename OutType>
225
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
226
- torch::Tensor const& B,
227
- torch::Tensor const& A_sf,
228
- torch::Tensor const& B_sf,
229
- torch::Tensor const& alpha, int64_t m, int64_t n,
230
- int64_t k, cudaStream_t stream) {
231
- TORCH_CHECK(false,
232
- "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
233
- "a CUTLASS 3.8 source directory to enable support.");
227
+void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
228
+ torch::stable::Tensor const& A,
229
+ torch::stable::Tensor const& B,
230
+ torch::stable::Tensor const& A_sf,
231
+ torch::stable::Tensor const& B_sf,
232
+ torch::stable::Tensor const& alpha, int64_t m,
233
+ int64_t n, int64_t k, cudaStream_t stream) {
234
+ STD_TORCH_CHECK(false,
235
+ "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
236
+ "a CUTLASS 3.8 source directory to enable support.");
237
}
238
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
239
237
-#define CHECK_TYPE(x, st, m) \
238
- TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
240
+#define CHECK_TYPE(x, st, m) \
241
+ STD_TORCH_CHECK(x.scalar_type() == st, \
242
+ ": Inconsistency of torch::stable::Tensor type:", m)
243
#define CHECK_TH_CUDA(x, m) \
240
- TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
244
+ STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
245
#define CHECK_CONTIGUOUS(x, m) \
242
- TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
246
+ STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
247
#define CHECK_INPUT(x, st, m) \
248
CHECK_TH_CUDA(x, m); \
249
CHECK_CONTIGUOUS(x, m); \
250
CHECK_TYPE(x, st, m)
251
248
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
249
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
252
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
253
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
254
251
-void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
252
- torch::Tensor const& B,
253
- torch::Tensor const& A_sf,
254
- torch::Tensor const& B_sf,
255
- torch::Tensor const& alpha) {
255
+void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
256
+ torch::stable::Tensor const& A,
257
+ torch::stable::Tensor const& B,
258
+ torch::stable::Tensor const& A_sf,
259
+ torch::stable::Tensor const& B_sf,
260
+ torch::stable::Tensor const& alpha) {
261
CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
262
CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
263
264
CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
265
CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
266
262
- CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
267
+ CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha");
268
264
- TORCH_CHECK(A.dim() == 2, "a must be a matrix");
265
- TORCH_CHECK(B.dim() == 2, "b must be a matrix");
266
- TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
267
- "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
268
- A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
269
+ STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix");
270
+ STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix");
271
+ STD_TORCH_CHECK(A.size(1) == B.size(1),
272
+ "a and b shapes cannot be multiplied (", A.size(0), "x",
273
+ A.size(1), " and ", B.size(0), "x", B.size(1), ")");
274
270
- auto const m = A.sizes()[0];
271
- auto const n = B.sizes()[0];
272
- auto const k = A.sizes()[1] * 2;
275
+ auto const m = A.size(0);
276
+ auto const n = B.size(0);
277
+ auto const k = A.size(1) * 2;
278
279
constexpr int alignment = 32;
275
- TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
276
- ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
277
- "), k: ", k, ".");
278
- TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
279
- ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
280
+ STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ",
281
+ alignment, ", but got a shape: (", A.size(0), "x", A.size(1),
282
+ "), k: ", k, ".");
283
+ STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ",
284
+ alignment, ", but got b shape: (", B.size(0), "x", B.size(1),
285
+ ").");
286
287
auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
288
int rounded_m = round_up(m, 128);
@@ -285,33 +291,34 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
291
// integer.
292
int rounded_k = round_up(k / 16, 4);
293
288
- TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
289
- TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
290
- TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
291
- "scale_a and scale_b shapes cannot be multiplied (",
292
- A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
293
- "x", B_sf.sizes()[1], ")");
294
- TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
295
- "scale_a must be padded and swizzled to a shape (", rounded_m,
296
- "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
297
- A_sf.sizes()[1], ")");
298
- TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
299
- "scale_b must be padded and swizzled to a shape (", rounded_n,
300
- "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
301
- B_sf.sizes()[1], ")");
302
-
303
- auto out_dtype = D.dtype();
304
- const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
305
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
306
-
307
- if (out_dtype == at::ScalarType::Half) {
294
+ STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
295
+ STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
296
+ STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1),
297
+ "scale_a and scale_b shapes cannot be multiplied (",
298
+ A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x",
299
+ B_sf.size(1), ")");
300
+ STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
301
+ "scale_a must be padded and swizzled to a shape (", rounded_m,
302
+ "x", rounded_k, "), but got a shape (", A_sf.size(0), "x",
303
+ A_sf.size(1), ")");
304
+ STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
305
+ "scale_b must be padded and swizzled to a shape (", rounded_n,
306
+ "x", rounded_k, "), but got a shape (", B_sf.size(0), "x",
307
+ B_sf.size(1), ")");
308
+
309
+ auto out_dtype = D.scalar_type();
310
+ const torch::stable::accelerator::DeviceGuard device_guard(
311
+ A.get_device_index());
312
+ const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
313
+
314
+ if (out_dtype == torch::headeronly::ScalarType::Half) {
315
cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n,
316
k, stream);
310
- } else if (out_dtype == at::ScalarType::BFloat16) {
317
+ } else if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
318
cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha,
319
m, n, k, stream);
320
} else {
314
- TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype,
315
- ")");
321
+ STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (",
322
+ out_dtype, ")");
323
}
324
}
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
renamed
+86
-79
@@ -14,10 +14,9 @@
14
* limitations under the License.
15
*/
16
17
-#include <torch/all.h>
17
+#include <torch/csrc/stable/tensor.h>
18
19
-#include <ATen/cuda/CUDAContext.h>
20
-#include <c10/cuda/CUDAGuard.h>
19
+#include "libtorch_stable/torch_utils.h"
20
21
#include "cutlass_extensions/common.hpp"
22
@@ -34,19 +33,20 @@
33
34
using namespace cute;
35
37
-#define CHECK_TYPE(x, st, m) \
38
- TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
36
+#define CHECK_TYPE(x, st, m) \
37
+ STD_TORCH_CHECK(x.scalar_type() == st, \
38
+ ": Inconsistency of torch::stable::Tensor type:", m)
39
#define CHECK_TH_CUDA(x, m) \
40
- TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
40
+ STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
41
#define CHECK_CONTIGUOUS(x, m) \
42
- TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
42
+ STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
43
#define CHECK_INPUT(x, st, m) \
44
CHECK_TH_CUDA(x, m); \
45
CHECK_CONTIGUOUS(x, m); \
46
CHECK_TYPE(x, st, m)
47
48
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
49
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
48
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
49
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
50
51
struct sm120_fp4_config_M256 {
52
using ClusterShape = Shape<_1, _1, _1>;
@@ -109,12 +109,13 @@ struct Fp4GemmSm120 {
109
};
110
111
template <typename Gemm>
112
-typename Gemm::Arguments args_from_options(at::Tensor& D, at::Tensor const& A,
113
- at::Tensor const& B,
114
- at::Tensor const& A_sf,
115
- at::Tensor const& B_sf,
116
- torch::Tensor const& alpha, int M,
117
- int N, int K) {
112
+typename Gemm::Arguments args_from_options(torch::stable::Tensor& D,
113
+ torch::stable::Tensor const& A,
114
+ torch::stable::Tensor const& B,
115
+ torch::stable::Tensor const& A_sf,
116
+ torch::stable::Tensor const& B_sf,
117
+ torch::stable::Tensor const& alpha,
118
+ int M, int N, int K) {
119
using ElementA = typename Gemm::ElementA;
120
using ElementB = typename Gemm::ElementB;
121
using ElementD = typename Gemm::ElementD;
@@ -158,18 +159,19 @@ typename Gemm::Arguments args_from_options(at::Tensor& D, at::Tensor const& A,
159
}
160
161
template <typename Gemm>
161
-void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
162
- at::Tensor const& A_sf, at::Tensor const& B_sf,
163
- torch::Tensor const& alpha, int M, int N, int K,
162
+void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A,
163
+ torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf,
164
+ torch::stable::Tensor const& B_sf,
165
+ torch::stable::Tensor const& alpha, int M, int N, int K,
166
cudaStream_t stream) {
167
Gemm gemm;
168
169
auto arguments = args_from_options<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
170
171
size_t workspace_size = Gemm::get_workspace_size(arguments);
170
- auto const workspace_options =
171
- torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
172
- auto workspace = torch::empty(workspace_size, workspace_options);
172
+ auto workspace =
173
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
174
+ std::nullopt, A.device());
175
176
CUTLASS_CHECK(gemm.can_implement(arguments));
177
@@ -178,12 +180,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
180
CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
181
}
182
181
-void cutlass_fp4_bf16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
182
- torch::Tensor const& B,
183
- torch::Tensor const& A_sf,
184
- torch::Tensor const& B_sf,
185
- torch::Tensor const& alpha, int m, int n,
186
- int k, cudaStream_t stream) {
183
+void cutlass_fp4_bf16_gemm_dispatch(torch::stable::Tensor& D,
184
+ torch::stable::Tensor const& A,
185
+ torch::stable::Tensor const& B,
186
+ torch::stable::Tensor const& A_sf,
187
+ torch::stable::Tensor const& B_sf,
188
+ torch::stable::Tensor const& alpha, int m,
189
+ int n, int k, cudaStream_t stream) {
190
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
191
if (mp2 <= 256) {
192
runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
@@ -194,12 +197,13 @@ void cutlass_fp4_bf16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
197
}
198
}
199
197
-void cutlass_fp4_f16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
198
- torch::Tensor const& B,
199
- torch::Tensor const& A_sf,
200
- torch::Tensor const& B_sf,
201
- torch::Tensor const& alpha, int m, int n,
202
- int k, cudaStream_t stream) {
200
+void cutlass_fp4_f16_gemm_dispatch(torch::stable::Tensor& D,
201
+ torch::stable::Tensor const& A,
202
+ torch::stable::Tensor const& B,
203
+ torch::stable::Tensor const& A_sf,
204
+ torch::stable::Tensor const& B_sf,
205
+ torch::stable::Tensor const& alpha, int m,
206
+ int n, int k, cudaStream_t stream) {
207
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
208
if (mp2 <= 256) {
209
runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
@@ -210,11 +214,12 @@ void cutlass_fp4_f16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
214
}
215
}
216
213
-void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
214
- torch::Tensor const& B,
215
- torch::Tensor const& A_sf,
216
- torch::Tensor const& B_sf,
217
- torch::Tensor const& alpha) {
217
+void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
218
+ torch::stable::Tensor const& A,
219
+ torch::stable::Tensor const& B,
220
+ torch::stable::Tensor const& A_sf,
221
+ torch::stable::Tensor const& B_sf,
222
+ torch::stable::Tensor const& alpha) {
223
#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
224
CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
225
CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
@@ -222,24 +227,25 @@ void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
227
CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
228
CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
229
225
- CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
230
+ CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha");
231
227
- TORCH_CHECK(A.dim() == 2, "a must be a matrix");
228
- TORCH_CHECK(B.dim() == 2, "b must be a matrix");
229
- TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
230
- "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
231
- A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
232
+ STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix");
233
+ STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix");
234
+ STD_TORCH_CHECK(A.size(1) == B.size(1),
235
+ "a and b shapes cannot be multiplied (", A.size(0), "x",
236
+ A.size(1), " and ", B.size(0), "x", B.size(1), ")");
237
233
- auto const m = A.sizes()[0];
234
- auto const n = B.sizes()[0];
235
- auto const k = A.sizes()[1] * 2;
238
+ auto const m = A.size(0);
239
+ auto const n = B.size(0);
240
+ auto const k = A.size(1) * 2;
241
242
constexpr int alignment = 32;
238
- TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
239
- ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
240
- "), k: ", k, ".");
241
- TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
242
- ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
243
+ STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ",
244
+ alignment, ", but got a shape: (", A.size(0), "x", A.size(1),
245
+ "), k: ", k, ".");
246
+ STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ",
247
+ alignment, ", but got b shape: (", B.size(0), "x", B.size(1),
248
+ ").");
249
250
auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
251
int rounded_m = round_up(m, 128);
@@ -248,38 +254,39 @@ void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
254
// integer.
255
int rounded_k = round_up(k / 16, 4);
256
251
- TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
252
- TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
253
- TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
254
- "scale_a and scale_b shapes cannot be multiplied (",
255
- A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
256
- "x", B_sf.sizes()[1], ")");
257
- TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
258
- "scale_a must be padded and swizzled to a shape (", rounded_m,
259
- "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
260
- A_sf.sizes()[1], ")");
261
- TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
262
- "scale_b must be padded and swizzled to a shape (", rounded_n,
263
- "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
264
- B_sf.sizes()[1], ")");
265
-
266
- auto out_dtype = D.dtype();
267
- const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
268
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
269
-
270
- if (out_dtype == at::ScalarType::BFloat16) {
257
+ STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
258
+ STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
259
+ STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1),
260
+ "scale_a and scale_b shapes cannot be multiplied (",
261
+ A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x",
262
+ B_sf.size(1), ")");
263
+ STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
264
+ "scale_a must be padded and swizzled to a shape (", rounded_m,
265
+ "x", rounded_k, "), but got a shape (", A_sf.size(0), "x",
266
+ A_sf.size(1), ")");
267
+ STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
268
+ "scale_b must be padded and swizzled to a shape (", rounded_n,
269
+ "x", rounded_k, "), but got a shape (", B_sf.size(0), "x",
270
+ B_sf.size(1), ")");
271
+
272
+ auto out_dtype = D.scalar_type();
273
+ const torch::stable::accelerator::DeviceGuard device_guard(
274
+ A.get_device_index());
275
+ const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
276
+
277
+ if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
278
return cutlass_fp4_bf16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
279
stream);
273
- } else if (out_dtype == at::ScalarType::Half) {
280
+ } else if (out_dtype == torch::headeronly::ScalarType::Half) {
281
return cutlass_fp4_f16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
282
stream);
283
} else {
277
- TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
278
- out_dtype, ")");
284
+ STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
285
+ out_dtype, ")");
286
}
287
#else
281
- TORCH_CHECK(false,
282
- "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
283
- "a CUTLASS 3.8 source directory to enable support.");
288
+ STD_TORCH_CHECK(false,
289
+ "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
290
+ "a CUTLASS 3.8 source directory to enable support.");
291
#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
285
-}
\ No newline at end of file
292
+}
csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh
renamed
+1
-1
@@ -20,7 +20,7 @@
20
#include <cuda_fp8.h>
21
#include <utility>
22
23
-#include "../../cuda_vec_utils.cuh"
23
+#include "cuda_vec_utils.cuh"
24
25
#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
26
CUDA_VERSION >= 12090
csrc/libtorch_stable/torch_bindings.cpp
+110
@@ -103,6 +103,102 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
103
ops.def(
104
"cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
105
"bool");
106
+
107
+ // CUTLASS nvfp4 block scaled GEMM
108
+ ops.def(
109
+ "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
110
+ " Tensor block_scale_a, Tensor block_scale_b,"
111
+ " Tensor alpha) -> ()");
112
+
113
+ // cutlass nvfp4 block scaled group GEMM
114
+ ops.def(
115
+ "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
116
+ " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
117
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
118
+
119
+ // Compute NVFP4 block quantized tensor.
120
+ ops.def(
121
+ "scaled_fp4_quant(Tensor input,"
122
+ " Tensor input_scale, bool "
123
+ "is_sf_swizzled_layout) -> (Tensor, Tensor)");
124
+
125
+ // Out variant
126
+ // TODO: Add out_variant tag once PyTorch supports it (added in 2.11)
127
+ // This registration is now migrated to stable ABI
128
+ // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not
129
+ // yet in torch/headeronly), the tag should be applied from Python
130
+ // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,))
131
+ // with the .impl remaining in C++.
132
+ // See pytorch/pytorch#176117.
133
+ ops.def(
134
+ "scaled_fp4_quant.out(Tensor input,"
135
+ " Tensor input_scale, bool "
136
+ "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
137
+ "-> ()");
138
+
139
+ // Compute NVFP4 experts quantization.
140
+ ops.def(
141
+ "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
142
+ "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
143
+ "Tensor output_scale_offset_by_experts) -> ()");
144
+
145
+ // Fused SiLU+Mul+NVFP4 experts quantization.
146
+ ops.def(
147
+ "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
148
+ "output_scale,"
149
+ "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
150
+ "Tensor output_scale_offset_by_experts) -> ()");
151
+
152
+ // Fused SiLU+Mul+NVFP4 quantization.
153
+ ops.def(
154
+ "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
155
+ "Tensor input, Tensor input_global_scale) -> ()");
156
+
157
+ // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
158
+ // of the given capability
159
+ ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
160
+
161
+ // CUTLASS w4a8 GEMM
162
+ ops.def(
163
+ "cutlass_w4a8_mm("
164
+ " Tensor A,"
165
+ " Tensor B,"
166
+ " Tensor group_scales,"
167
+ " int group_size,"
168
+ " Tensor channel_scales,"
169
+ " Tensor token_scales,"
170
+ " ScalarType? out_type,"
171
+ " str? maybe_schedule"
172
+ ") -> Tensor");
173
+
174
+ // pack scales
175
+ ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
176
+
177
+ // encode and reorder weight matrix
178
+ ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
179
+
180
+ // CUTLASS w4a8 grouped GEMM
181
+ ops.def(
182
+ "cutlass_w4a8_moe_mm("
183
+ " Tensor! out_tensors,"
184
+ " Tensor a_tensors,"
185
+ " Tensor b_tensors,"
186
+ " Tensor a_scales,"
187
+ " Tensor b_scales,"
188
+ " Tensor b_group_scales,"
189
+ " int b_group_size,"
190
+ " Tensor expert_offsets,"
191
+ " Tensor problem_sizes,"
192
+ " Tensor a_strides,"
193
+ " Tensor b_strides,"
194
+ " Tensor c_strides,"
195
+ " Tensor group_scale_strides,"
196
+ " str? maybe_schedule"
197
+ ") -> ()");
198
+
199
+ ops.def(
200
+ "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
201
+ "Tensor)");
202
#endif
203
}
204
@@ -128,6 +224,18 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
224
TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets));
225
ops.impl("get_cutlass_batched_moe_mm_data",
226
TORCH_BOX(&get_cutlass_batched_moe_mm_data));
227
+
228
+ // FP4/NVFP4 ops
229
+ ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm));
230
+ ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func));
231
+ ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out));
232
+ ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant));
233
+ ops.impl("silu_and_mul_scaled_fp4_experts_quant",
234
+ TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
235
+ ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
236
+
237
+ // W4A8 ops: impl registrations are in the source files
238
+ // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
239
#endif
240
}
241
@@ -143,6 +251,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
251
TORCH_BOX(&cutlass_group_gemm_supported));
252
ops.impl("cutlass_scaled_mm_supports_block_fp8",
253
TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8));
254
+ ops.impl("cutlass_scaled_mm_supports_fp4",
255
+ TORCH_BOX(&cutlass_scaled_mm_supports_fp4));
256
#endif
257
}
258
csrc/libtorch_stable/torch_utils.h
+1
@@ -2,6 +2,7 @@
2
3
#include <torch/csrc/inductor/aoti_torch/c/shim.h>
4
#include <torch/csrc/stable/accelerator.h>
5
+#include <torch/csrc/stable/ops.h>
6
#include <torch/csrc/stable/tensor.h>
7
#include <torch/headeronly/util/shim_utils.h>
8
csrc/ops.h
-44
@@ -152,12 +152,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
152
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
153
torch::Tensor& scale);
154
155
-#ifndef USE_ROCM
156
-void silu_and_mul_nvfp4_quant(torch::Tensor& out,
157
- torch::Tensor& output_block_scale,
158
- torch::Tensor& input,
159
- torch::Tensor& input_global_scale);
160
-#endif
155
void persistent_masked_m_silu_mul_quant(
156
const at::Tensor& input, // (E, T, 2*H)
157
const at::Tensor& counts, // (E)
@@ -225,44 +219,6 @@ torch::Tensor ggml_moe_a8_vec(torch::Tensor X, torch::Tensor W,
219
220
int64_t ggml_moe_get_block_size(int64_t type);
221
228
-#ifndef USE_ROCM
229
-
230
-bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
231
-
232
-void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
233
- torch::Tensor const& B, torch::Tensor const& A_sf,
234
- torch::Tensor const& B_sf,
235
- torch::Tensor const& alpha);
236
-
237
-void cutlass_fp4_group_mm(
238
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
239
- const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
240
- const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
241
- const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets);
242
-
243
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
244
- torch::Tensor const& input, torch::Tensor const& input_scale,
245
- bool is_sf_swizzled_layout);
246
-
247
-void scaled_fp4_quant_out(torch::Tensor const& input,
248
- torch::Tensor const& input_scale,
249
- bool is_sf_swizzled_layout, torch::Tensor& output,
250
- torch::Tensor& output_scale);
251
-
252
-void scaled_fp4_experts_quant(
253
- torch::Tensor& output, torch::Tensor& output_scale,
254
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
255
- torch::Tensor const& input_offset_by_experts,
256
- torch::Tensor const& output_scale_offset_by_experts);
257
-
258
-void silu_and_mul_scaled_fp4_experts_quant(
259
- torch::Tensor& output, torch::Tensor& output_scale,
260
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
261
- torch::Tensor const& input_offset_by_experts,
262
- torch::Tensor const& output_scale_offset_by_experts);
263
-
264
-#endif
265
-
222
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
223
torch::Tensor const& scale,
224
std::optional<torch::Tensor> const& azp);
csrc/quantization/fp4/nvfp4_quant_entry.cu
deleted
-163
@@ -1,163 +0,0 @@
1
-/*
2
- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
- *
4
- * Licensed under the Apache License, Version 2.0 (the "License");
5
- * you may not use this file except in compliance with the License.
6
- * You may obtain a copy of the License at
7
- *
8
- * http://www.apache.org/licenses/LICENSE-2.0
9
- *
10
- * Unless required by applicable law or agreed to in writing, software
11
- * distributed under the License is distributed on an "AS IS" BASIS,
12
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
- * See the License for the specific language governing permissions and
14
- * limitations under the License.
15
- */
16
-
17
-#include <torch/all.h>
18
-
19
-#include "cutlass_extensions/common.hpp"
20
-#include "nvfp4_utils.cuh"
21
-
22
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
23
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
24
-void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
25
- torch::Tensor const& input,
26
- torch::Tensor const& output_sf,
27
- torch::Tensor const& input_sf,
28
- bool is_sf_swizzled_layout);
29
-#endif
30
-
31
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
32
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
33
-void scaled_fp4_experts_quant_sm1xxa(
34
- torch::Tensor& output, torch::Tensor& output_scale,
35
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
36
- torch::Tensor const& input_offset_by_experts,
37
- torch::Tensor const& output_scale_offset_by_experts);
38
-#endif
39
-
40
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
41
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
42
-void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
43
- torch::Tensor& output_sf,
44
- torch::Tensor& input,
45
- torch::Tensor& input_sf);
46
-#endif
47
-
48
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
49
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
50
-void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
51
- torch::Tensor& output, torch::Tensor& output_scale,
52
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
53
- torch::Tensor const& input_offset_by_experts,
54
- torch::Tensor const& output_scale_offset_by_experts);
55
-#endif
56
-
57
-static bool nvfp4_quant_sm_supported() {
58
- const int32_t sm = get_sm_version_num();
59
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
60
- if (sm >= 100 && sm < 120) return true;
61
-#endif
62
-#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120
63
- if (sm >= 120 && sm < 130) return true;
64
-#endif
65
- return false;
66
-}
67
-
68
-void scaled_fp4_quant_out(torch::Tensor const& input,
69
- torch::Tensor const& input_sf,
70
- bool is_sf_swizzled_layout, torch::Tensor& output,
71
- torch::Tensor& output_sf) {
72
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
73
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
74
- TORCH_CHECK(nvfp4_quant_sm_supported(),
75
- "No compiled nvfp4 quantization kernel for SM ",
76
- get_sm_version_num(),
77
- ". Recompile with the appropriate CUDA arch.");
78
- return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
79
- is_sf_swizzled_layout);
80
-#endif
81
- TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
82
-}
83
-
84
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
85
- torch::Tensor const& input, torch::Tensor const& input_sf,
86
- bool is_sf_swizzled_layout) {
87
- int64_t n = input.size(-1);
88
- int64_t m = input.numel() / n;
89
- auto device = input.device();
90
-
91
- // Two fp4 values packed into a uint8
92
- auto output = torch::empty(
93
- {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
94
-
95
- torch::Tensor output_sf;
96
- if (is_sf_swizzled_layout) {
97
- auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
98
- output_sf = torch::empty(
99
- {sf_m, sf_n},
100
- torch::TensorOptions().device(device).dtype(torch::kInt32));
101
- } else {
102
- output_sf = torch::empty(
103
- {m, n / CVT_FP4_SF_VEC_SIZE},
104
- torch::TensorOptions().device(device).dtype(torch::kUInt8));
105
- }
106
-
107
- scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
108
- output_sf);
109
- return {output, output_sf};
110
-}
111
-
112
-void scaled_fp4_experts_quant(
113
- torch::Tensor& output, torch::Tensor& output_scale,
114
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
115
- torch::Tensor const& input_offset_by_experts,
116
- torch::Tensor const& output_scale_offset_by_experts) {
117
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
118
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
119
- TORCH_CHECK(nvfp4_quant_sm_supported(),
120
- "No compiled nvfp4 experts quantization kernel for SM ",
121
- get_sm_version_num(),
122
- ". Recompile with the appropriate CUDA arch.");
123
- return scaled_fp4_experts_quant_sm1xxa(
124
- output, output_scale, input, input_global_scale, input_offset_by_experts,
125
- output_scale_offset_by_experts);
126
-#endif
127
- TORCH_CHECK_NOT_IMPLEMENTED(false,
128
- "No compiled nvfp4 experts quantization kernel");
129
-}
130
-
131
-void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
132
- torch::Tensor& input, torch::Tensor& input_sf) {
133
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
134
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
135
- TORCH_CHECK(nvfp4_quant_sm_supported(),
136
- "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
137
- get_sm_version_num(),
138
- ". Recompile with the appropriate CUDA arch.");
139
- return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
140
-#endif
141
- TORCH_CHECK_NOT_IMPLEMENTED(
142
- false, "No compiled silu_and_mul nvfp4 quantization kernel");
143
-}
144
-
145
-void silu_and_mul_scaled_fp4_experts_quant(
146
- torch::Tensor& output, torch::Tensor& output_scale,
147
- torch::Tensor const& input, torch::Tensor const& input_global_scale,
148
- torch::Tensor const& input_offset_by_experts,
149
- torch::Tensor const& output_scale_offset_by_experts) {
150
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
151
- (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
152
- TORCH_CHECK(nvfp4_quant_sm_supported(),
153
- "No compiled silu_and_mul nvfp4 experts quantization kernel "
154
- "for SM ",
155
- get_sm_version_num(),
156
- ". Recompile with the appropriate CUDA arch.");
157
- return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
158
- output, output_scale, input, input_global_scale, input_offset_by_experts,
159
- output_scale_offset_by_experts);
160
-#endif
161
- TORCH_CHECK_NOT_IMPLEMENTED(
162
- false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
163
-}
csrc/torch_bindings.cpp
-100
@@ -109,13 +109,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
109
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
110
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
111
112
-#ifndef USE_ROCM
113
- ops.def(
114
- "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
115
- "Tensor input, Tensor input_global_scale) -> ()");
116
- ops.impl("silu_and_mul_nvfp4_quant", torch::kCUDA, &silu_and_mul_nvfp4_quant);
117
-#endif
118
-
112
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
113
ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
114
@@ -332,47 +325,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
325
"Tensor? qzeros_or_none, bool inplace) -> Tensor");
326
// conditionally compiled so impl registrations are in source file
327
335
- // CUTLASS w4a8 GEMM
336
- ops.def(
337
- "cutlass_w4a8_mm("
338
- " Tensor A,"
339
- " Tensor B,"
340
- " Tensor group_scales,"
341
- " int group_size,"
342
- " Tensor channel_scales,"
343
- " Tensor token_scales,"
344
- " ScalarType? out_type,"
345
- " str? maybe_schedule"
346
- ") -> Tensor");
347
- // pack scales
348
- ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
349
- // encode and reorder weight matrix
350
- ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
351
- // conditionally compiled so impl registration is in source file
352
-
353
- // CUTLASS w4a8 grouped GEMM
354
- ops.def(
355
- "cutlass_w4a8_moe_mm("
356
- " Tensor! out_tensors,"
357
- " Tensor a_tensors,"
358
- " Tensor b_tensors,"
359
- " Tensor a_scales,"
360
- " Tensor b_scales,"
361
- " Tensor b_group_scales,"
362
- " int b_group_size,"
363
- " Tensor expert_offsets,"
364
- " Tensor problem_sizes,"
365
- " Tensor a_strides,"
366
- " Tensor b_strides,"
367
- " Tensor c_strides,"
368
- " Tensor group_scale_strides,"
369
- " str? maybe_schedule"
370
- ") -> ()");
371
- ops.def(
372
- "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
373
- "Tensor)");
374
- // conditionally compiled so impl registration is in source file
375
-
328
#endif
329
330
// Dequantization for GGML.
@@ -409,20 +361,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
361
ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
362
363
#ifndef USE_ROCM
412
- // CUTLASS nvfp4 block scaled GEMM
413
- ops.def(
414
- "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
415
- " Tensor block_scale_a, Tensor block_scale_b,"
416
- " Tensor alpha) -> ()");
417
- ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
418
-
419
- // cutlass nvfp4 block scaled group GEMM
420
- ops.def(
421
- "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
422
- " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
423
- " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
424
- // conditionally compiled so impl registration is in source file
425
-
364
// Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
365
ops.def(
366
"mxfp8_experts_quant("
@@ -455,44 +393,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
393
"-> int");
394
// conditionally compiled so impl in source file
395
458
- // Compute NVFP4 block quantized tensor.
459
- ops.def(
460
- "scaled_fp4_quant(Tensor input,"
461
- " Tensor input_scale, bool "
462
- "is_sf_swizzled_layout) -> (Tensor, Tensor)");
463
- ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
464
-
465
- // Out variant
466
- // TODO: Add {at::Tag::out_variant} tag and update all call sites
467
- // to use the functional variant once vLLM upgrades PyTorch.
468
- // See pytorch/pytorch#176117.
469
- ops.def(
470
- "scaled_fp4_quant.out(Tensor input,"
471
- " Tensor input_scale, bool "
472
- "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
473
- "-> ()");
474
- ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
475
-
476
- // Compute NVFP4 experts quantization.
477
- ops.def(
478
- "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
479
- "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
480
- "Tensor output_scale_offset_by_experts) -> ()");
481
- ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
482
-
483
- // Fused SiLU+Mul+NVFP4 experts quantization.
484
- ops.def(
485
- "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
486
- "output_scale,"
487
- "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
488
- "Tensor output_scale_offset_by_experts) -> ()");
489
- ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
490
- &silu_and_mul_scaled_fp4_experts_quant);
491
-
492
- // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
493
- // of the given capability
494
- ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
495
- ops.impl("cutlass_scaled_mm_supports_fp4", &cutlass_scaled_mm_supports_fp4);
396
#endif
397
398
// Quantized GEMM for GPTQ.