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, 19:45 UTC
[3/n] Migrate cutlass/scaled_mm_entry.cu torch stable ABI (#37221)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
mikaylagawarecki committed
Mar 30, 2026 at 14:20 UTC
ab1a6a43fa9500697dd01e73aa372c8777cd7a5b
54 files changed
+1845
-1613
CMakeLists.txt
+211
-207
@@ -340,7 +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/w8a8/cutlass/scaled_mm_entry.cu"
343
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
344
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
345
"csrc/cutlass_extensions/common.cpp")
@@ -490,132 +489,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
489
" in CUDA target architectures")
490
endif()
491
493
-
494
- set(SCALED_MM_3X_ARCHS)
495
- # The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
496
- # CUDA 12.0 or later
497
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
498
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS)
499
- set(SRCS
500
- "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu"
501
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu"
502
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu"
503
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu"
504
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu")
505
- set_gencode_flags_for_srcs(
506
- SRCS "${SRCS}"
507
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
508
- list(APPEND VLLM_EXT_SRC "${SRCS}")
509
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
510
- # Let scaled_mm_c2x know it doesn't need to build these arches
511
- list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
512
- message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
513
- else()
514
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS)
515
- message(STATUS "Not building scaled_mm_c3x_sm90 as CUDA Compiler version is "
516
- "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
517
- "later if you intend on running FP8 quantized models on "
518
- "Hopper.")
519
- else()
520
- message(STATUS "Not building scaled_mm_c3x_sm90 as no compatible archs found "
521
- "in CUDA target architectures")
522
- endif()
523
- endif()
524
-
525
-
526
- # The cutlass_scaled_mm kernels for Blackwell SM12x (c3x, i.e. CUTLASS 3.x) require
527
- # CUDA 12.8 or later
528
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
529
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0f" "${CUDA_ARCHS}")
530
- else()
531
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
532
- endif()
533
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
534
- set(SRCS
535
- "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu"
536
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu"
537
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu"
538
- )
539
- set_gencode_flags_for_srcs(
540
- SRCS "${SRCS}"
541
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
542
- list(APPEND VLLM_EXT_SRC "${SRCS}")
543
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
544
- # Let scaled_mm_c2x know it doesn't need to build these arches
545
- list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
546
- message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
547
- else()
548
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
549
- message(STATUS "Not building scaled_mm_c3x_sm120 as CUDA Compiler version is "
550
- "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
551
- "later if you intend on running FP8 quantized models on "
552
- "Blackwell.")
553
- else()
554
- message(STATUS "Not building scaled_mm_c3x_120 as no compatible archs found "
555
- "in CUDA target architectures")
556
- endif()
557
- endif()
558
-
559
-
560
- # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
561
- # require CUDA 12.8 or later
562
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
563
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
564
- else()
565
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
566
- endif()
567
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
568
- set(SRCS
569
- "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu"
570
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu"
571
- "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu"
572
- )
573
- set_gencode_flags_for_srcs(
574
- SRCS "${SRCS}"
575
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
576
- list(APPEND VLLM_EXT_SRC "${SRCS}")
577
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
578
- # Let scaled_mm_c2x know it doesn't need to build these arches
579
- list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
580
- message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
581
- else()
582
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
583
- message(STATUS "Not building scaled_mm_c3x_sm100 as CUDA Compiler version is "
584
- "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
585
- "later if you intend on running FP8 quantized models on "
586
- "Blackwell.")
587
- else()
588
- message(STATUS "Not building scaled_mm_c3x_100 as no compatible archs found "
589
- "in CUDA target architectures")
590
- endif()
591
- endif()
592
-
593
- #
594
- # For the cutlass_scaled_mm kernels we want to build the c2x (CUTLASS 2.x)
595
- # kernels for the remaining archs that are not already built for 3x.
596
- # (Build 8.9 for FP8)
597
- cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS
598
- "7.5;8.0;8.7;8.9+PTX" "${CUDA_ARCHS}")
599
- # subtract out the archs that are already built for 3x
600
- list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS})
601
- if (SCALED_MM_2X_ARCHS)
602
- set(SRCS "csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu")
603
- set_gencode_flags_for_srcs(
604
- SRCS "${SRCS}"
605
- CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
606
- list(APPEND VLLM_EXT_SRC "${SRCS}")
607
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
608
- message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
609
- else()
610
- if (SCALED_MM_3X_ARCHS)
611
- message(STATUS "Not building scaled_mm_c2x as all archs are already built"
612
- " for and covered by scaled_mm_c3x")
613
- else()
614
- message(STATUS "Not building scaled_mm_c2x as no compatible archs found "
615
- "in CUDA target architectures")
616
- endif()
617
- endif()
618
-
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)
@@ -693,55 +566,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
566
set(MLA_ARCHS)
567
endif()
568
696
- # CUTLASS MoE kernels
697
-
698
- # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works
699
- # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled
700
- # if it's possible to compile MoE kernels that use its output.
701
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
702
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
703
- set(SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu")
704
- set_gencode_flags_for_srcs(
705
- SRCS "${SRCS}"
706
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
707
- list(APPEND VLLM_EXT_SRC "${SRCS}")
708
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
709
- message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
710
- else()
711
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
712
- message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
713
- "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
714
- "if you intend on running FP8 quantized MoE models on Hopper.")
715
- else()
716
- message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
717
- "in CUDA target architectures.")
718
- endif()
719
- endif()
720
-
721
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
722
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
723
- else()
724
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
725
- endif()
726
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
727
- set(SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
728
- set_gencode_flags_for_srcs(
729
- SRCS "${SRCS}"
730
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
731
- list(APPEND VLLM_EXT_SRC "${SRCS}")
732
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
733
- message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
734
- else()
735
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
736
- message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
737
- "not >= 12.8, we recommend upgrading to CUDA 12.8 or later "
738
- "if you intend on running FP8 quantized MoE models on Blackwell.")
739
- else()
740
- message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
741
- "in CUDA target architectures.")
742
- endif()
743
- endif()
744
-
569
# Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
570
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
571
cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -787,36 +611,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
611
"in CUDA target architectures.")
612
endif()
613
790
- # moe_data.cu is used by all CUTLASS MoE kernels.
791
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
792
- cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
793
- else()
794
- cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
795
- endif()
796
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
797
- set(SRCS "csrc/quantization/w8a8/cutlass/moe/moe_data.cu")
798
- set_gencode_flags_for_srcs(
799
- SRCS "${SRCS}"
800
- CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}")
801
- list(APPEND VLLM_EXT_SRC "${SRCS}")
802
- message(STATUS "Building moe_data for archs: ${CUTLASS_MOE_DATA_ARCHS}")
803
- else()
804
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
805
- message(STATUS "Not building moe_data as CUDA Compiler version is "
806
- "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
807
- "if you intend on running FP8 quantized MoE models on Hopper or Blackwell.")
808
- else()
809
- message(STATUS "Not building moe_data as no compatible archs found "
810
- "in CUDA target architectures.")
811
- endif()
812
- endif()
813
-
814
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
815
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
816
- else()
817
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
818
- endif()
819
-
614
#
615
# Machete kernels
616
@@ -964,7 +758,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
758
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
759
#
760
set(VLLM_STABLE_EXT_SRC
967
- "csrc/libtorch_stable/torch_bindings.cpp")
761
+ "csrc/libtorch_stable/torch_bindings.cpp"
762
+ "csrc/cutlass_extensions/common.cpp"
763
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
764
765
if(VLLM_GPU_LANG STREQUAL "CUDA")
766
list(APPEND VLLM_STABLE_EXT_SRC
@@ -979,6 +775,209 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
775
CUDA_ARCHS "${CUDA_ARCHS}")
776
endif()
777
778
+ #
779
+ # CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
780
+ #
781
+ set(SCALED_MM_3X_ARCHS)
782
+ # The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
783
+ # CUDA 12.0 or later
784
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
785
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS)
786
+ set(SRCS
787
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu"
788
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu"
789
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu"
790
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu"
791
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu")
792
+ set_gencode_flags_for_srcs(
793
+ SRCS "${SRCS}"
794
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
795
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
796
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
797
+ # Let scaled_mm_c2x know it doesn't need to build these arches
798
+ list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
799
+ message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
800
+ else()
801
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS)
802
+ message(STATUS "Not building scaled_mm_c3x_sm90 as CUDA Compiler version is "
803
+ "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
804
+ "later if you intend on running FP8 quantized models on "
805
+ "Hopper.")
806
+ else()
807
+ message(STATUS "Not building scaled_mm_c3x_sm90 as no compatible archs found "
808
+ "in CUDA target architectures")
809
+ endif()
810
+ endif()
811
+
812
+
813
+ # The cutlass_scaled_mm kernels for Blackwell SM12x (c3x, i.e. CUTLASS 3.x) require
814
+ # CUDA 12.8 or later
815
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
816
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0f" "${CUDA_ARCHS}")
817
+ else()
818
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
819
+ endif()
820
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
821
+ set(SRCS
822
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu"
823
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu"
824
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu"
825
+ )
826
+ set_gencode_flags_for_srcs(
827
+ SRCS "${SRCS}"
828
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
829
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
830
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
831
+ # Let scaled_mm_c2x know it doesn't need to build these arches
832
+ list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
833
+ message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
834
+ else()
835
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
836
+ message(STATUS "Not building scaled_mm_c3x_sm120 as CUDA Compiler version is "
837
+ "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
838
+ "later if you intend on running FP8 quantized models on "
839
+ "Blackwell.")
840
+ else()
841
+ message(STATUS "Not building scaled_mm_c3x_120 as no compatible archs found "
842
+ "in CUDA target architectures")
843
+ endif()
844
+ endif()
845
+
846
+
847
+ # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
848
+ # require CUDA 12.8 or later
849
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
850
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
851
+ else()
852
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
853
+ endif()
854
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
855
+ set(SRCS
856
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu"
857
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu"
858
+ "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu"
859
+ )
860
+ set_gencode_flags_for_srcs(
861
+ SRCS "${SRCS}"
862
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
863
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
864
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
865
+ # Let scaled_mm_c2x know it doesn't need to build these arches
866
+ list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
867
+ message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
868
+ else()
869
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
870
+ message(STATUS "Not building scaled_mm_c3x_sm100 as CUDA Compiler version is "
871
+ "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
872
+ "later if you intend on running FP8 quantized models on "
873
+ "Blackwell.")
874
+ else()
875
+ message(STATUS "Not building scaled_mm_c3x_100 as no compatible archs found "
876
+ "in CUDA target architectures")
877
+ endif()
878
+ endif()
879
+
880
+ #
881
+ # For the cutlass_scaled_mm kernels we want to build the c2x (CUTLASS 2.x)
882
+ # kernels for the remaining archs that are not already built for 3x.
883
+ # (Build 8.9 for FP8)
884
+ cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS
885
+ "7.5;8.0;8.7;8.9+PTX" "${CUDA_ARCHS}")
886
+ # subtract out the archs that are already built for 3x
887
+ list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS})
888
+ if (SCALED_MM_2X_ARCHS)
889
+ set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu")
890
+ set_gencode_flags_for_srcs(
891
+ SRCS "${SRCS}"
892
+ CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
893
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
894
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
895
+ message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
896
+ else()
897
+ if (SCALED_MM_3X_ARCHS)
898
+ message(STATUS "Not building scaled_mm_c2x as all archs are already built"
899
+ " for and covered by scaled_mm_c3x")
900
+ else()
901
+ message(STATUS "Not building scaled_mm_c2x as no compatible archs found "
902
+ "in CUDA target architectures")
903
+ endif()
904
+ endif()
905
+
906
+ #
907
+ # CUTLASS MoE kernels (moved from _C to _C_stable_libtorch)
908
+ #
909
+
910
+ # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works
911
+ # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled
912
+ # if it's possible to compile MoE kernels that use its output.
913
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
914
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
915
+ set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu")
916
+ set_gencode_flags_for_srcs(
917
+ SRCS "${SRCS}"
918
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
919
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
920
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
921
+ message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
922
+ else()
923
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
924
+ message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
925
+ "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
926
+ "if you intend on running FP8 quantized MoE models on Hopper.")
927
+ else()
928
+ message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
929
+ "in CUDA target architectures.")
930
+ endif()
931
+ endif()
932
+
933
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
934
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
935
+ else()
936
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
937
+ endif()
938
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
939
+ set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
940
+ set_gencode_flags_for_srcs(
941
+ SRCS "${SRCS}"
942
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
943
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
944
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
945
+ message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
946
+ else()
947
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
948
+ message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
949
+ "not >= 12.8, we recommend upgrading to CUDA 12.8 or later "
950
+ "if you intend on running FP8 quantized MoE models on Blackwell.")
951
+ else()
952
+ message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
953
+ "in CUDA target architectures.")
954
+ endif()
955
+ endif()
956
+
957
+ # moe_data.cu is used by all CUTLASS MoE kernels.
958
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
959
+ cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
960
+ else()
961
+ cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
962
+ endif()
963
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
964
+ set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu")
965
+ set_gencode_flags_for_srcs(
966
+ SRCS "${SRCS}"
967
+ CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}")
968
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
969
+ message(STATUS "Building moe_data for archs: ${CUTLASS_MOE_DATA_ARCHS}")
970
+ else()
971
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
972
+ message(STATUS "Not building moe_data as CUDA Compiler version is "
973
+ "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
974
+ "if you intend on running FP8 quantized MoE models on Hopper or Blackwell.")
975
+ else()
976
+ message(STATUS "Not building moe_data as no compatible archs found "
977
+ "in CUDA target architectures.")
978
+ endif()
979
+ endif()
980
+
981
message(STATUS "Enabling C_stable extension.")
982
define_extension_target(
983
_C_stable_libtorch
@@ -987,6 +986,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
986
SOURCES ${VLLM_STABLE_EXT_SRC}
987
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
988
ARCHITECTURES ${VLLM_GPU_ARCHES}
989
+ INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
990
USE_SABI 3
991
WITH_SOABI)
992
@@ -1000,6 +1000,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1000
# Needed to use cuda APIs from C-shim
1001
target_compile_definitions(_C_stable_libtorch PRIVATE
1002
USE_CUDA)
1003
+
1004
+ # Needed by CUTLASS kernels
1005
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1006
+ CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
1007
endif()
1008
1009
#
csrc/cutlass_extensions/common.hpp
+7
-5
@@ -6,14 +6,16 @@
6
#include <cstdio>
7
#include <cstdlib>
8
9
+#include <torch/headeronly/util/shim_utils.h>
10
+
11
/**
12
* Helper function for checking CUTLASS errors
13
*/
12
-#define CUTLASS_CHECK(status) \
13
- { \
14
- cutlass::Status error = status; \
15
- TORCH_CHECK(error == cutlass::Status::kSuccess, \
16
- cutlassGetStatusString(error)); \
14
+#define CUTLASS_CHECK(status) \
15
+ { \
16
+ cutlass::Status error = status; \
17
+ STD_TORCH_CHECK(error == cutlass::Status::kSuccess, \
18
+ cutlassGetStatusString(error)); \
19
}
20
21
inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) {
csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+33
-19
@@ -3,6 +3,14 @@
3
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp"
4
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp"
5
6
+// This header is shared by both _C (unstable ABI) and _C_stable_libtorch
7
+// (stable ABI) targets. When compiled under the stable ABI target,
8
+// TORCH_TARGET_VERSION is defined and Tensor is unavailable, so we
9
+// use torch::stable::Tensor instead.
10
+#ifdef TORCH_TARGET_VERSION
11
+ #include <torch/csrc/stable/tensor.h>
12
+#endif
13
+
14
/*
15
This file defines custom epilogues for fusing channel scales, token scales,
16
bias, and activation zero-points onto a GEMM operation using the
@@ -15,6 +23,12 @@
23
24
namespace vllm::c3x {
25
26
+#ifdef TORCH_TARGET_VERSION
27
+using TensorType = torch::stable::Tensor;
28
+#else
29
+using TensorType = torch::Tensor;
30
+#endif
31
+
32
using namespace cute;
33
34
template <typename T>
@@ -84,7 +98,7 @@ struct ScaledEpilogueBase {
98
// from a tensor. It can handle both row and column, as well as row/column or
99
// scalar cases.
100
template <typename Descriptor, typename T>
87
- static auto args_from_tensor(torch::Tensor const& tensor) {
101
+ static auto args_from_tensor(TensorType const& tensor) {
102
using Arguments = typename Descriptor::Arguments;
103
auto* data_ptr = static_cast<T*>(tensor.data_ptr());
104
if constexpr (std::is_same_v<Descriptor, ColOrScalarLoad<T>> ||
@@ -100,7 +114,7 @@ struct ScaledEpilogueBase {
114
// This overload handles the case where there might not be a tensor, in which
115
// case a nullptr is passed and a constant (0) is used.
116
template <typename Descriptor, typename T>
103
- static auto args_from_tensor(std::optional<torch::Tensor> const& tensor) {
117
+ static auto args_from_tensor(std::optional<TensorType> const& tensor) {
118
using Arguments = typename Descriptor::Arguments;
119
auto* data_ptr = tensor ? static_cast<T*>(tensor->data_ptr()) : nullptr;
120
static_assert(std::is_same_v<Descriptor, ColLoad<T, true>> ||
@@ -158,8 +172,8 @@ struct ScaledEpilogue
172
cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0>;
173
using ArgumentType = typename EVTCompute::Arguments;
174
161
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
162
- torch::Tensor const& b_scales) {
175
+ static ArgumentType prepare_args(TensorType const& a_scales,
176
+ TensorType const& b_scales) {
177
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
178
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
179
@@ -203,9 +217,9 @@ struct ScaledEpilogueBias
217
cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0, Bias>;
218
219
using ArgumentType = typename EVTCompute::Arguments;
206
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
207
- torch::Tensor const& b_scales,
208
- torch::Tensor const& bias) {
220
+ static ArgumentType prepare_args(TensorType const& a_scales,
221
+ TensorType const& b_scales,
222
+ TensorType const& bias) {
223
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
224
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
225
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -246,9 +260,9 @@ struct ScaledEpilogueColumnBias
260
cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0, Bias>;
261
262
using ArgumentType = typename EVTCompute::Arguments;
249
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
250
- torch::Tensor const& b_scales,
251
- torch::Tensor const& bias) {
263
+ static ArgumentType prepare_args(TensorType const& a_scales,
264
+ TensorType const& b_scales,
265
+ TensorType const& bias) {
266
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
267
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
268
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -304,10 +318,10 @@ struct ScaledEpilogueBiasAzp
318
EVTComputeScaleB, Bias>;
319
using ArgumentType = typename EVTCompute::Arguments;
320
307
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
308
- torch::Tensor const& b_scales,
309
- torch::Tensor const& azp_adj,
310
- std::optional<torch::Tensor> const& bias) {
321
+ static ArgumentType prepare_args(TensorType const& a_scales,
322
+ TensorType const& b_scales,
323
+ TensorType const& azp_adj,
324
+ std::optional<TensorType> const& bias) {
325
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
326
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
327
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -380,11 +394,11 @@ struct ScaledEpilogueBiasAzpToken
394
EVTComputeScaleB, Bias>;
395
using ArgumentType = typename EVTCompute::Arguments;
396
383
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
384
- torch::Tensor const& b_scales,
385
- torch::Tensor const& azp_adj,
386
- torch::Tensor const& azp,
387
- std::optional<torch::Tensor> const& bias) {
397
+ static ArgumentType prepare_args(TensorType const& a_scales,
398
+ TensorType const& b_scales,
399
+ TensorType const& azp_adj,
400
+ TensorType const& azp,
401
+ std::optional<TensorType> const& bias) {
402
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
403
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
404
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
renamed
+20
-16
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/csrc/stable/tensor.h>
4
+
5
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
6
7
/*
@@ -52,7 +54,7 @@ struct ScaledEpilogueBase {
54
// from a tensor. It can handle both row and column, as well as row/column or
55
// scalar cases.
56
template <typename Descriptor, typename T>
55
- static auto args_from_tensor(torch::Tensor const& tensor) {
57
+ static auto args_from_tensor(torch::stable::Tensor const& tensor) {
58
using Arguments = typename Descriptor::Arguments;
59
auto* data_ptr = static_cast<T*>(tensor.data_ptr());
60
if constexpr (std::is_same_v<Descriptor, ColOrScalarLoad<T>> ||
@@ -68,7 +70,8 @@ struct ScaledEpilogueBase {
70
// This overload handles the case where there might not be a tensor, in which
71
// case a nullptr is passed and a constant (0) is used.
72
template <typename Descriptor, typename T>
71
- static auto args_from_tensor(std::optional<torch::Tensor> const& tensor) {
73
+ static auto args_from_tensor(
74
+ std::optional<torch::stable::Tensor> const& tensor) {
75
static_assert(std::is_same_v<Descriptor, RowOrZeroLoad<T>>);
76
using Arguments = typename Descriptor::Arguments;
77
auto* data_ptr = tensor ? static_cast<T*>(tensor->data_ptr()) : nullptr;
@@ -117,8 +120,8 @@ struct ScaledEpilogue
120
cutlass::epilogue::threadblock::Sm80EVT<Compute1, ScaleA, EVTCompute0>;
121
using ArgumentType = typename EVTCompute::Arguments;
122
120
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
121
- torch::Tensor const& b_scales) {
123
+ static ArgumentType prepare_args(torch::stable::Tensor const& a_scales,
124
+ torch::stable::Tensor const& b_scales) {
125
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
126
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
127
@@ -160,9 +163,9 @@ struct ScaledEpilogueBias
163
using EVTCompute = cutlass::epilogue::threadblock::Sm80EVT<Compute1, ScaleA,
164
EVTCompute0, Bias>;
165
using ArgumentType = typename EVTCompute::Arguments;
163
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
164
- torch::Tensor const& b_scales,
165
- torch::Tensor const& bias) {
166
+ static ArgumentType prepare_args(torch::stable::Tensor const& a_scales,
167
+ torch::stable::Tensor const& b_scales,
168
+ torch::stable::Tensor const& bias) {
169
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
170
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
171
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -220,10 +223,11 @@ struct ScaledEpilogueBiasAzp
223
224
using ArgumentType = typename EVTCompute::Arguments;
225
223
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
224
- torch::Tensor const& b_scales,
225
- torch::Tensor const& azp_adj,
226
- std::optional<torch::Tensor> const& bias) {
226
+ static ArgumentType prepare_args(
227
+ torch::stable::Tensor const& a_scales,
228
+ torch::stable::Tensor const& b_scales,
229
+ torch::stable::Tensor const& azp_adj,
230
+ std::optional<torch::stable::Tensor> const& bias) {
231
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
232
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
233
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -298,11 +302,11 @@ struct ScaledEpilogueBiasAzpToken
302
303
using ArgumentType = typename EVTCompute::Arguments;
304
301
- static ArgumentType prepare_args(torch::Tensor const& a_scales,
302
- torch::Tensor const& b_scales,
303
- torch::Tensor const& azp_adj,
304
- torch::Tensor const& azp,
305
- std::optional<torch::Tensor> const& bias) {
305
+ static ArgumentType prepare_args(
306
+ torch::stable::Tensor const& a_scales,
307
+ torch::stable::Tensor const& b_scales,
308
+ torch::stable::Tensor const& azp_adj, torch::stable::Tensor const& azp,
309
+ std::optional<torch::stable::Tensor> const& bias) {
310
auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
311
auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
312
auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
csrc/libtorch_stable/ops.h
+57
@@ -27,4 +27,61 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
27
torch::stable::Tensor& output_s,
28
int64_t group_size, double eps, double int8_min,
29
double int8_max);
30
+
31
+bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
32
+bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
33
+bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
34
+
35
+void cutlass_scaled_mm(torch::stable::Tensor& out,
36
+ torch::stable::Tensor const& a,
37
+ torch::stable::Tensor const& b,
38
+ torch::stable::Tensor const& a_scales,
39
+ torch::stable::Tensor const& b_scales,
40
+ std::optional<torch::stable::Tensor> const& bias);
41
+
42
+void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
43
+ torch::stable::Tensor const& a_tensors,
44
+ torch::stable::Tensor const& b_tensors,
45
+ torch::stable::Tensor const& a_scales,
46
+ torch::stable::Tensor const& b_scales,
47
+ torch::stable::Tensor const& expert_offsets,
48
+ torch::stable::Tensor const& problem_sizes,
49
+ torch::stable::Tensor const& a_strides,
50
+ torch::stable::Tensor const& b_strides,
51
+ torch::stable::Tensor const& c_strides, bool per_act_token,
52
+ bool per_out_ch);
53
+
54
+void cutlass_scaled_mm_azp(torch::stable::Tensor& out,
55
+ torch::stable::Tensor const& a,
56
+ torch::stable::Tensor const& b,
57
+ torch::stable::Tensor const& a_scales,
58
+ torch::stable::Tensor const& b_scales,
59
+ torch::stable::Tensor const& azp_adj,
60
+ std::optional<torch::stable::Tensor> const& azp,
61
+ std::optional<torch::stable::Tensor> const& bias);
62
+
63
+void get_cutlass_moe_mm_data(
64
+ const torch::stable::Tensor& topk_ids,
65
+ torch::stable::Tensor& expert_offsets,
66
+ torch::stable::Tensor& problem_sizes1,
67
+ torch::stable::Tensor& problem_sizes2,
68
+ torch::stable::Tensor& input_permutation,
69
+ torch::stable::Tensor& output_permutation, const int64_t num_experts,
70
+ const int64_t n, const int64_t k,
71
+ const std::optional<torch::stable::Tensor>& blockscale_offsets,
72
+ const bool is_gated);
73
+
74
+void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
75
+ const torch::stable::Tensor& expert_first_token_offset,
76
+ torch::stable::Tensor& problem_sizes1,
77
+ torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
78
+ const bool swap_ab);
79
+
80
+void get_cutlass_batched_moe_mm_data(
81
+ torch::stable::Tensor& expert_offsets,
82
+ torch::stable::Tensor& problem_sizes1,
83
+ torch::stable::Tensor& problem_sizes2,
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
#endif
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
renamed
+12
-10
@@ -2,9 +2,10 @@
2
3
// clang-format will break include orders
4
// clang-format off
5
-#include <torch/all.h>
5
+#include <torch/csrc/stable/tensor.h>
6
+#include <torch/csrc/stable/ops.h>
7
7
-#include <ATen/cuda/CUDAContext.h>
8
+#include "libtorch_stable/torch_utils.h"
9
10
#include "cutlass/cutlass.h"
11
@@ -25,14 +26,14 @@
26
namespace vllm::c3x {
27
28
static inline cute::Shape<int, int, int, int> get_problem_shape(
28
- torch::Tensor const& a, torch::Tensor const& b) {
29
+ torch::stable::Tensor const& a, torch::stable::Tensor const& b) {
30
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
31
return {m, n, k, 1};
32
}
33
34
template <typename GemmKernel>
35
void cutlass_gemm_caller(
35
- torch::Device device, cute::Shape<int, int, int, int> prob_shape,
36
+ torch::stable::Device device, cute::Shape<int, int, int, int> prob_shape,
37
typename GemmKernel::MainloopArguments mainloop_args,
38
typename GemmKernel::EpilogueArguments epilogue_args,
39
typename GemmKernel::TileSchedulerArguments scheduler = {}) {
@@ -50,19 +51,20 @@ void cutlass_gemm_caller(
51
CUTLASS_CHECK(gemm_op.can_implement(args));
52
53
size_t workspace_size = gemm_op.get_workspace_size(args);
53
- auto const workspace_options =
54
- torch::TensorOptions().dtype(torch::kUInt8).device(device);
55
- auto workspace = torch::empty(workspace_size, workspace_options);
54
+ auto workspace =
55
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
56
+ std::nullopt, device);
57
57
- auto stream = at::cuda::getCurrentCUDAStream(device.index());
58
+ auto stream = get_current_cuda_stream(device.index());
59
60
cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
61
CUTLASS_CHECK(status);
62
}
63
64
template <typename Gemm, typename... EpilogueArgs>
64
-void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
65
- torch::Tensor const& b,
65
+void cutlass_gemm_caller(torch::stable::Tensor& out,
66
+ torch::stable::Tensor const& a,
67
+ torch::stable::Tensor const& b,
68
EpilogueArgs&&... epilogue_params) {
69
using ElementAB = typename Gemm::ElementAB;
70
using ElementC = typename Gemm::ElementC;
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh
renamed
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu
renamed
+6
-7
@@ -4,13 +4,12 @@
4
5
namespace vllm {
6
7
-void cutlass_scaled_mm_azp_sm90_int8(torch::Tensor& out, torch::Tensor const& a,
8
- torch::Tensor const& b,
9
- torch::Tensor const& a_scales,
10
- torch::Tensor const& b_scales,
11
- torch::Tensor const& azp_adj,
12
- std::optional<torch::Tensor> const& azp,
13
- std::optional<torch::Tensor> const& bias) {
7
+void cutlass_scaled_mm_azp_sm90_int8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
11
+ std::optional<torch::stable::Tensor> const& azp,
12
+ std::optional<torch::stable::Tensor> const& bias) {
13
if (azp) {
14
return cutlass_scaled_mm_sm90_int8_epilogue<
15
c3x::ScaledEpilogueBiasAzpToken>(out, a, b, a_scales, b_scales, azp_adj,
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu
new
+22
@@ -0,0 +1,22 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_blockwise_sm100_fp8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales) {
11
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
12
+ cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::bfloat16_t>(
13
+ out, a, b, a_scales, b_scales);
14
+
15
+ } else {
16
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
17
+ cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::half_t>(
18
+ out, a, b, a_scales, b_scales);
19
+ }
20
+}
21
+
22
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
renamed
+11
-9
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "cuda_utils.h"
6
#include "cutlass/cutlass.h"
7
#include "cutlass/numeric_types.h"
@@ -130,10 +132,10 @@ struct cutlass_3x_gemm_fp8_blockwise {
132
};
133
134
template <typename Gemm>
133
-void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
134
- torch::Tensor const& b,
135
- torch::Tensor const& a_scales,
136
- torch::Tensor const& b_scales) {
135
+void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
136
+ torch::stable::Tensor const& b,
137
+ torch::stable::Tensor const& a_scales,
138
+ torch::stable::Tensor const& b_scales) {
139
static constexpr bool swap_ab = Gemm::swap_ab;
140
using GemmKernel = typename Gemm::GemmKernel;
141
using StrideA = typename Gemm::GemmKernel::StrideA;
@@ -200,11 +202,11 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
202
}
203
204
template <typename OutType>
203
-void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::Tensor& out,
204
- torch::Tensor const& a,
205
- torch::Tensor const& b,
206
- torch::Tensor const& a_scales,
207
- torch::Tensor const& b_scales) {
205
+void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out,
206
+ torch::stable::Tensor const& a,
207
+ torch::stable::Tensor const& b,
208
+ torch::stable::Tensor const& a_scales,
209
+ torch::stable::Tensor const& b_scales) {
210
int32_t m = a.size(0), n = b.size(1), k = a.size(1), sms;
211
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
212
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu
new
+22
@@ -0,0 +1,22 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_blockwise_sm120_fp8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales) {
11
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
12
+ cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::bfloat16_t>(
13
+ out, a, b, a_scales, b_scales);
14
+
15
+ } else {
16
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
17
+ cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::half_t>(
18
+ out, a, b, a_scales, b_scales);
19
+ }
20
+}
21
+
22
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
renamed
+11
-9
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "cuda_utils.h"
6
#include "cutlass/cutlass.h"
7
#include "cutlass/numeric_types.h"
@@ -138,10 +140,10 @@ struct sm120_blockwise_fp8_config_M64 {
140
};
141
142
template <typename Gemm>
141
-void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
142
- torch::Tensor const& b,
143
- torch::Tensor const& a_scales,
144
- torch::Tensor const& b_scales) {
143
+void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
144
+ torch::stable::Tensor const& b,
145
+ torch::stable::Tensor const& a_scales,
146
+ torch::stable::Tensor const& b_scales) {
147
using GemmKernel = typename Gemm::GemmKernel;
148
using StrideA = typename Gemm::GemmKernel::StrideA;
149
using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -196,11 +198,11 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
198
}
199
200
template <typename OutType>
199
-void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::Tensor& out,
200
- torch::Tensor const& a,
201
- torch::Tensor const& b,
202
- torch::Tensor const& a_scales,
203
- torch::Tensor const& b_scales) {
201
+void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
202
+ torch::stable::Tensor const& a,
203
+ torch::stable::Tensor const& b,
204
+ torch::stable::Tensor const& a_scales,
205
+ torch::stable::Tensor const& b_scales) {
206
int M = a.size(0);
207
if (M <= 256) {
208
using Gemm = typename sm120_blockwise_fp8_config_M64<OutType>::Gemm;
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu
new
+23
@@ -0,0 +1,23 @@
1
+
2
+#include "scaled_mm_kernels.hpp"
3
+#include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh"
4
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
5
+
6
+namespace vllm {
7
+
8
+void cutlass_scaled_mm_blockwise_sm90_fp8(
9
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
10
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
11
+ torch::stable::Tensor const& b_scales) {
12
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
13
+ cutlass_gemm_blockwise_sm90_fp8_dispatch<cutlass::bfloat16_t>(
14
+ out, a, b, a_scales, b_scales);
15
+
16
+ } else {
17
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
18
+ cutlass_gemm_blockwise_sm90_fp8_dispatch<cutlass::half_t>(
19
+ out, a, b, a_scales, b_scales);
20
+ }
21
+}
22
+
23
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh
renamed
+12
-10
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "cutlass/cutlass.h"
6
#include "cutlass/numeric_types.h"
7
@@ -101,10 +103,10 @@ struct cutlass_3x_gemm_fp8_blockwise {
103
};
104
105
template <typename Gemm>
104
-void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
105
- torch::Tensor const& b,
106
- torch::Tensor const& a_scales,
107
- torch::Tensor const& b_scales) {
106
+void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
107
+ torch::stable::Tensor const& b,
108
+ torch::stable::Tensor const& a_scales,
109
+ torch::stable::Tensor const& b_scales) {
110
using GemmKernel = typename Gemm::GemmKernel;
111
using StrideA = typename Gemm::GemmKernel::StrideA;
112
using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -120,7 +122,7 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
122
123
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
124
123
- TORCH_CHECK(m % 4 == 0, "m must be divisible by 4");
125
+ STD_TORCH_CHECK(m % 4 == 0, "m must be divisible by 4");
126
127
StrideA a_stride;
128
StrideB b_stride;
@@ -161,11 +163,11 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
163
}
164
165
template <typename OutType>
164
-void cutlass_gemm_blockwise_sm90_fp8_dispatch(torch::Tensor& out,
165
- torch::Tensor const& a,
166
- torch::Tensor const& b,
167
- torch::Tensor const& a_scales,
168
- torch::Tensor const& b_scales) {
166
+void cutlass_gemm_blockwise_sm90_fp8_dispatch(torch::stable::Tensor& out,
167
+ torch::stable::Tensor const& a,
168
+ torch::stable::Tensor const& b,
169
+ torch::stable::Tensor const& a_scales,
170
+ torch::stable::Tensor const& b_scales) {
171
// TODO: better heuristics
172
cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
173
OutType, 1, 128, 128, Shape<_128, _128, _128>,
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp
renamed
+21
-16
@@ -1,52 +1,57 @@
1
-#include <torch/all.h>
1
+#include <torch/csrc/stable/tensor.h>
2
+#include <torch/headeronly/core/ScalarType.h>
3
#include "cuda_utils.h"
4
#include "cutlass_extensions/common.hpp"
5
6
template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
6
-void dispatch_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
7
- torch::Tensor const& b, torch::Tensor const& a_scales,
8
- torch::Tensor const& b_scales,
9
- std::optional<torch::Tensor> const& bias,
7
+void dispatch_scaled_mm(torch::stable::Tensor& c,
8
+ torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b,
10
+ torch::stable::Tensor const& a_scales,
11
+ torch::stable::Tensor const& b_scales,
12
+ std::optional<torch::stable::Tensor> const& bias,
13
Fp8Func fp8_func, Int8Func int8_func,
14
BlockwiseFunc blockwise_func) {
12
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
13
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
15
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
16
+ torch::headeronly::ScalarType::Float);
17
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
18
+ torch::headeronly::ScalarType::Float);
19
20
int M = a.size(0), N = b.size(1), K = a.size(1);
21
22
if ((a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
23
(b_scales.numel() == 1 || b_scales.numel() == b.size(1))) {
24
// Standard per-tensor/per-token/per-channel scaling
20
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
21
- if (a.dtype() == torch::kFloat8_e4m3fn) {
25
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
26
+ if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) {
27
fp8_func(c, a, b, a_scales, b_scales, bias);
28
} else {
24
- TORCH_CHECK(a.dtype() == torch::kInt8);
29
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
30
if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
31
int8_func(c, a, b, a_scales, b_scales, bias);
32
} else {
33
int32_t version_num = get_sm_version_num();
29
- TORCH_CHECK(
34
+ STD_TORCH_CHECK(
35
false, "Int8 not supported on SM", version_num,
36
". Use FP8 quantization instead, or run on older arch (SM < 100).");
37
}
38
}
39
} else {
35
- TORCH_CHECK(a_scales.dim() == 2, "a scale must be 2d tensor.");
36
- TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor.");
40
+ STD_TORCH_CHECK(a_scales.dim() == 2, "a scale must be 2d tensor.");
41
+ STD_TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor.");
42
int32_t version_num = get_sm_version_num();
43
if (version_num >= 90) {
39
- TORCH_CHECK(
44
+ STD_TORCH_CHECK(
45
a.size(0) == a_scales.size(0) &&
46
cuda_utils::ceil_div(a.size(1), int64_t(128)) == a_scales.size(1),
47
"a_scale_group_shape must be [1, 128].");
43
- TORCH_CHECK(
48
+ STD_TORCH_CHECK(
49
cuda_utils::ceil_div(b.size(0), int64_t(128)) == b_scales.size(0) &&
50
cuda_utils::ceil_div(b.size(1), int64_t(128)) == b_scales.size(1),
51
"b_scale_group_shape must be [128, 128].");
52
}
53
49
- TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm");
54
+ STD_TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm");
55
blockwise_func(c, a, b, a_scales, b_scales);
56
}
57
}
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp
new
+52
@@ -0,0 +1,52 @@
1
+#pragma once
2
+
3
+#include <torch/csrc/stable/tensor.h>
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_sm90_fp8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales,
11
+ std::optional<torch::stable::Tensor> const& bias);
12
+
13
+void cutlass_scaled_mm_sm90_int8(
14
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
15
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
16
+ torch::stable::Tensor const& b_scales,
17
+ std::optional<torch::stable::Tensor> const& bias);
18
+
19
+void cutlass_scaled_mm_azp_sm90_int8(
20
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
21
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
22
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
23
+ std::optional<torch::stable::Tensor> const& azp,
24
+ std::optional<torch::stable::Tensor> const& bias);
25
+
26
+void cutlass_scaled_mm_blockwise_sm90_fp8(
27
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
28
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
29
+ torch::stable::Tensor const& b_scales);
30
+
31
+void cutlass_scaled_mm_sm100_fp8(
32
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
33
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
34
+ torch::stable::Tensor const& b_scales,
35
+ std::optional<torch::stable::Tensor> const& bias);
36
+
37
+void cutlass_scaled_mm_sm120_fp8(
38
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
39
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
40
+ torch::stable::Tensor const& b_scales,
41
+ std::optional<torch::stable::Tensor> const& bias);
42
+
43
+void cutlass_scaled_mm_blockwise_sm100_fp8(
44
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
45
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
46
+ torch::stable::Tensor const& b_scales);
47
+
48
+void cutlass_scaled_mm_blockwise_sm120_fp8(
49
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
50
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
51
+ torch::stable::Tensor const& b_scales);
52
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu
new
+24
@@ -0,0 +1,24 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm100_fp8_dispatch.cuh"
3
+
4
+namespace vllm {
5
+
6
+void cutlass_scaled_mm_sm100_fp8(
7
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
8
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
9
+ torch::stable::Tensor const& b_scales,
10
+ std::optional<torch::stable::Tensor> const& bias) {
11
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
12
+ if (bias) {
13
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
14
+ "currently bias dtype must match output dtype ",
15
+ out.scalar_type());
16
+ return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
17
+ b_scales, *bias);
18
+ } else {
19
+ return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
20
+ b_scales);
21
+ }
22
+}
23
+
24
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh
renamed
+24
-19
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm.cuh"
6
#include "cutlass_gemm_caller.cuh"
7
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
@@ -192,8 +194,9 @@ struct sm100_fp8_config_M16_swap_ab {
194
};
195
196
template <typename Gemm, typename... EpilogueArgs>
195
-void cutlass_gemm_caller_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
196
- torch::Tensor const& b,
197
+void cutlass_gemm_caller_sm100_fp8(torch::stable::Tensor& out,
198
+ torch::stable::Tensor const& a,
199
+ torch::stable::Tensor const& b,
200
EpilogueArgs&&... epilogue_params) {
201
static constexpr bool swap_ab = Gemm::swap_ab;
202
using ElementAB = typename Gemm::ElementAB;
@@ -237,15 +240,15 @@ void cutlass_gemm_caller_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
240
241
template <typename InType, typename OutType, bool EnableBias,
242
typename... EpilogueArgs>
240
-inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
241
- torch::Tensor const& a,
242
- torch::Tensor const& b,
243
- torch::Tensor const& a_scales,
244
- torch::Tensor const& b_scales,
245
- EpilogueArgs&&... args) {
243
+inline void cutlass_gemm_sm100_fp8_dispatch(
244
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
245
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
246
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
247
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
247
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
248
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
248
+ STD_TORCH_CHECK(a.scalar_type() ==
249
+ torch::headeronly::ScalarType::Float8_e4m3fn);
250
+ STD_TORCH_CHECK(b.scalar_type() ==
251
+ torch::headeronly::ScalarType::Float8_e4m3fn);
252
253
using Cutlass3xGemmDefault =
254
typename sm100_fp8_config_default<InType, OutType,
@@ -292,22 +295,24 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
295
}
296
297
template <bool EnableBias, typename... EpilogueArgs>
295
-void cutlass_scaled_mm_sm100_fp8_epilogue(torch::Tensor& out,
296
- torch::Tensor const& a,
297
- torch::Tensor const& b,
298
- torch::Tensor const& a_scales,
299
- torch::Tensor const& b_scales,
298
+void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
299
+ torch::stable::Tensor const& a,
300
+ torch::stable::Tensor const& b,
301
+ torch::stable::Tensor const& a_scales,
302
+ torch::stable::Tensor const& b_scales,
303
EpilogueArgs&&... epilogue_args) {
301
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
302
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
304
+ STD_TORCH_CHECK(a.scalar_type() ==
305
+ torch::headeronly::ScalarType::Float8_e4m3fn);
306
+ STD_TORCH_CHECK(b.scalar_type() ==
307
+ torch::headeronly::ScalarType::Float8_e4m3fn);
308
304
- if (out.dtype() == torch::kBFloat16) {
309
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
310
return cutlass_gemm_sm100_fp8_dispatch<cutlass::float_e4m3_t,
311
cutlass::bfloat16_t, EnableBias>(
312
out, a, b, a_scales, b_scales,
313
std::forward<EpilogueArgs>(epilogue_args)...);
314
} else {
310
- TORCH_CHECK(out.dtype() == torch::kFloat16);
315
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
316
return cutlass_gemm_sm100_fp8_dispatch<cutlass::float_e4m3_t,
317
cutlass::half_t, EnableBias>(
318
out, a, b, a_scales, b_scales,
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu
new
+25
@@ -0,0 +1,25 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm120_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_sm120_fp8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales,
11
+ std::optional<torch::stable::Tensor> const& bias) {
12
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
13
+ if (bias) {
14
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
15
+ "currently bias dtype must match output dtype ",
16
+ out.scalar_type());
17
+ return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
18
+ out, a, b, a_scales, b_scales, *bias);
19
+ } else {
20
+ return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
21
+ out, a, b, a_scales, b_scales);
22
+ }
23
+}
24
+
25
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh
renamed
+18
-12
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm.cuh"
6
#include "cutlass_gemm_caller.cuh"
7
@@ -138,13 +140,15 @@ struct sm120_fp8_config_M16 {
140
template <typename InType, typename OutType,
141
template <typename, typename, typename> typename Epilogue,
142
typename... EpilogueArgs>
141
-inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
142
- torch::Tensor const& a,
143
- torch::Tensor const& b,
143
+inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
144
+ torch::stable::Tensor const& a,
145
+ torch::stable::Tensor const& b,
146
EpilogueArgs&&... args) {
147
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
146
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
147
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
148
+ STD_TORCH_CHECK(a.scalar_type() ==
149
+ torch::headeronly::ScalarType::Float8_e4m3fn);
150
+ STD_TORCH_CHECK(b.scalar_type() ==
151
+ torch::headeronly::ScalarType::Float8_e4m3fn);
152
153
int M = a.size(0);
154
@@ -177,19 +181,21 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
181
182
template <template <typename, typename, typename> typename Epilogue,
183
typename... EpilogueArgs>
180
-void cutlass_scaled_mm_sm120_fp8_epilogue(torch::Tensor& out,
181
- torch::Tensor const& a,
182
- torch::Tensor const& b,
184
+void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
185
+ torch::stable::Tensor const& a,
186
+ torch::stable::Tensor const& b,
187
EpilogueArgs&&... epilogue_args) {
184
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
185
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
188
+ STD_TORCH_CHECK(a.scalar_type() ==
189
+ torch::headeronly::ScalarType::Float8_e4m3fn);
190
+ STD_TORCH_CHECK(b.scalar_type() ==
191
+ torch::headeronly::ScalarType::Float8_e4m3fn);
192
187
- if (out.dtype() == torch::kBFloat16) {
193
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
194
return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
195
cutlass::bfloat16_t, Epilogue>(
196
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
197
} else {
192
- TORCH_CHECK(out.dtype() == torch::kFloat16);
198
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
199
return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
200
cutlass::half_t, Epilogue>(
201
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu
new
+24
@@ -0,0 +1,24 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm90_fp8_dispatch.cuh"
3
+
4
+namespace vllm {
5
+
6
+void cutlass_scaled_mm_sm90_fp8(
7
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
8
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
9
+ torch::stable::Tensor const& b_scales,
10
+ std::optional<torch::stable::Tensor> const& bias) {
11
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
12
+ if (bias) {
13
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
14
+ "currently bias dtype must match output dtype ",
15
+ out.scalar_type());
16
+ return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
17
+ b_scales, *bias);
18
+ } else {
19
+ return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
20
+ b_scales);
21
+ }
22
+}
23
+
24
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh
renamed
+24
-19
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm.cuh"
6
#include "cutlass_gemm_caller.cuh"
7
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
@@ -235,8 +237,9 @@ struct sm90_fp8_config_M16_N8192 {
237
};
238
239
template <typename Gemm, typename... EpilogueArgs>
238
-void cutlass_gemm_caller_sm90_fp8(torch::Tensor& out, torch::Tensor const& a,
239
- torch::Tensor const& b,
240
+void cutlass_gemm_caller_sm90_fp8(torch::stable::Tensor& out,
241
+ torch::stable::Tensor const& a,
242
+ torch::stable::Tensor const& b,
243
EpilogueArgs&&... epilogue_params) {
244
static constexpr bool swap_ab = Gemm::swap_ab;
245
using ElementAB = typename Gemm::ElementAB;
@@ -280,15 +283,15 @@ void cutlass_gemm_caller_sm90_fp8(torch::Tensor& out, torch::Tensor const& a,
283
284
template <typename InType, typename OutType, bool EnableBias,
285
typename... EpilogueArgs>
283
-inline void cutlass_gemm_sm90_fp8_dispatch(torch::Tensor& out,
284
- torch::Tensor const& a,
285
- torch::Tensor const& b,
286
- torch::Tensor const& a_scales,
287
- torch::Tensor const& b_scales,
288
- EpilogueArgs&&... args) {
286
+inline void cutlass_gemm_sm90_fp8_dispatch(
287
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
288
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
289
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
290
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
290
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
291
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
291
+ STD_TORCH_CHECK(a.scalar_type() ==
292
+ torch::headeronly::ScalarType::Float8_e4m3fn);
293
+ STD_TORCH_CHECK(b.scalar_type() ==
294
+ torch::headeronly::ScalarType::Float8_e4m3fn);
295
296
using Cutlass3xGemmDefault =
297
typename sm90_fp8_config_default<InType, OutType,
@@ -347,22 +350,24 @@ inline void cutlass_gemm_sm90_fp8_dispatch(torch::Tensor& out,
350
}
351
352
template <bool EnableBias, typename... EpilogueArgs>
350
-void cutlass_scaled_mm_sm90_fp8_epilogue(torch::Tensor& out,
351
- torch::Tensor const& a,
352
- torch::Tensor const& b,
353
- torch::Tensor const& a_scales,
354
- torch::Tensor const& b_scales,
353
+void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
354
+ torch::stable::Tensor const& a,
355
+ torch::stable::Tensor const& b,
356
+ torch::stable::Tensor const& a_scales,
357
+ torch::stable::Tensor const& b_scales,
358
EpilogueArgs&&... epilogue_args) {
356
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
357
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
359
+ STD_TORCH_CHECK(a.scalar_type() ==
360
+ torch::headeronly::ScalarType::Float8_e4m3fn);
361
+ STD_TORCH_CHECK(b.scalar_type() ==
362
+ torch::headeronly::ScalarType::Float8_e4m3fn);
363
359
- if (out.dtype() == torch::kBFloat16) {
364
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
365
return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
366
cutlass::bfloat16_t, EnableBias>(
367
out, a, b, a_scales, b_scales,
368
std::forward<EpilogueArgs>(epilogue_args)...);
369
} else {
365
- TORCH_CHECK(out.dtype() == torch::kFloat16);
370
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
371
return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
372
cutlass::half_t, EnableBias>(
373
out, a, b, a_scales, b_scales,
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu
new
+25
@@ -0,0 +1,25 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm90_int8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_sm90_int8(
8
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
9
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
10
+ torch::stable::Tensor const& b_scales,
11
+ std::optional<torch::stable::Tensor> const& bias) {
12
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
13
+ if (bias) {
14
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
15
+ "currently bias dtype must match output dtype ",
16
+ out.scalar_type());
17
+ return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBias>(
18
+ out, a, b, a_scales, b_scales, *bias);
19
+ } else {
20
+ return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogue>(
21
+ out, a, b, a_scales, b_scales);
22
+ }
23
+}
24
+
25
+} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh
renamed
+14
-12
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm.cuh"
6
#include "cutlass_gemm_caller.cuh"
7
@@ -87,13 +89,13 @@ struct sm90_int8_config_M32_NSmall {
89
template <typename InType, typename OutType,
90
template <typename, typename, typename> typename Epilogue,
91
typename... EpilogueArgs>
90
-inline void cutlass_gemm_sm90_int8_dispatch(torch::Tensor& out,
91
- torch::Tensor const& a,
92
- torch::Tensor const& b,
92
+inline void cutlass_gemm_sm90_int8_dispatch(torch::stable::Tensor& out,
93
+ torch::stable::Tensor const& a,
94
+ torch::stable::Tensor const& b,
95
EpilogueArgs&&... args) {
96
static_assert(std::is_same<InType, int8_t>());
95
- TORCH_CHECK(a.dtype() == torch::kInt8);
96
- TORCH_CHECK(b.dtype() == torch::kInt8);
97
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
98
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
99
100
using Cutlass3xGemmDefault =
101
typename sm90_int8_config_default<InType, OutType,
@@ -142,19 +144,19 @@ inline void cutlass_gemm_sm90_int8_dispatch(torch::Tensor& out,
144
145
template <template <typename, typename, typename> typename Epilogue,
146
typename... EpilogueArgs>
145
-void cutlass_scaled_mm_sm90_int8_epilogue(torch::Tensor& out,
146
- torch::Tensor const& a,
147
- torch::Tensor const& b,
147
+void cutlass_scaled_mm_sm90_int8_epilogue(torch::stable::Tensor& out,
148
+ torch::stable::Tensor const& a,
149
+ torch::stable::Tensor const& b,
150
EpilogueArgs&&... epilogue_args) {
149
- TORCH_CHECK(a.dtype() == torch::kInt8);
150
- TORCH_CHECK(b.dtype() == torch::kInt8);
151
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
152
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
153
152
- if (out.dtype() == torch::kBFloat16) {
154
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
155
return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
156
Epilogue>(
157
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
158
} else {
157
- TORCH_CHECK(out.dtype() == torch::kFloat16);
159
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
160
return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::half_t, Epilogue>(
161
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
162
}
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh
renamed
+27
-20
@@ -1,10 +1,10 @@
1
#pragma once
2
3
#include <cuda.h>
4
-#include <torch/all.h>
5
-#include <c10/cuda/CUDAStream.h>
4
+#include <torch/csrc/stable/tensor.h>
5
+#include <torch/headeronly/core/ScalarType.h>
6
+#include "libtorch_stable/torch_utils.h"
7
7
-#include "core/scalar_type.hpp"
8
#include "cutlass/bfloat16.h"
9
#include "cutlass/float8.h"
10
@@ -31,7 +31,7 @@ __global__ void get_group_gemm_starts(
31
}
32
33
#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
34
- else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
34
+ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
35
get_group_gemm_starts<cutlass::float_e4m3_t, C_TYPE, float> \
36
<<<1, num_experts, 0, stream>>>( \
37
static_cast<int64_t*>(expert_offsets.data_ptr()), \
@@ -51,32 +51,39 @@ __global__ void get_group_gemm_starts(
51
namespace {
52
53
void run_get_group_gemm_starts(
54
- torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
55
- torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
56
- torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
57
- torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
58
- torch::Tensor& out_tensors, torch::Tensor const& a_scales,
59
- torch::Tensor const& b_scales) {
60
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
61
- TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);
62
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
63
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
54
+ torch::stable::Tensor const& expert_offsets, torch::stable::Tensor& a_ptrs,
55
+ torch::stable::Tensor& b_ptrs, torch::stable::Tensor& out_ptrs,
56
+ torch::stable::Tensor& a_scales_ptrs, torch::stable::Tensor& b_scales_ptrs,
57
+ torch::stable::Tensor const& a_tensors,
58
+ torch::stable::Tensor const& b_tensors, torch::stable::Tensor& out_tensors,
59
+ torch::stable::Tensor const& a_scales,
60
+ torch::stable::Tensor const& b_scales) {
61
+ STD_TORCH_CHECK(a_tensors.scalar_type() ==
62
+ torch::headeronly::ScalarType::Float8_e4m3fn);
63
+ STD_TORCH_CHECK(b_tensors.scalar_type() ==
64
+ torch::headeronly::ScalarType::Float8_e4m3fn);
65
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
66
+ torch::headeronly::ScalarType::Float);
67
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
68
+ torch::headeronly::ScalarType::Float);
69
// expect int64_t to avoid overflow during offset calculations
65
- TORCH_CHECK(expert_offsets.dtype() == torch::kInt64);
70
+ STD_TORCH_CHECK(expert_offsets.scalar_type() ==
71
+ torch::headeronly::ScalarType::Long);
72
73
int num_experts = static_cast<int>(expert_offsets.size(0));
74
bool per_act_token = a_scales.numel() != 1;
75
bool per_out_ch = b_scales.numel() != num_experts;
76
71
- auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
77
+ auto stream = get_current_cuda_stream(a_tensors.get_device_index());
78
79
if (false) {
80
}
75
- __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
76
- __CALL_GET_STARTS_KERNEL(torch::kFloat16, half)
81
+ __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::BFloat16,
82
+ cutlass::bfloat16_t)
83
+ __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::Half, half)
84
else {
78
- TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
85
+ STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
86
}
87
}
88
82
-} // namespace
\ No newline at end of file
89
+} // namespace
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh
renamed
+28
-19
@@ -6,6 +6,7 @@
6
#include "cutlass/epilogue/collective/collective_builder.hpp"
7
#include "cutlass/gemm/device/gemm_universal_adapter.h"
8
9
+#include <torch/csrc/stable/ops.h>
10
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
11
#include "cutlass_extensions/common.hpp"
12
#include "get_group_starts.cuh"
@@ -84,13 +85,17 @@ struct cutlass_3x_group_gemm {
85
};
86
87
template <typename Gemm>
87
-void cutlass_group_gemm_caller(
88
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
89
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
90
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
91
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
92
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
93
- bool per_act_token, bool per_out_ch) {
88
+void cutlass_group_gemm_caller(torch::stable::Tensor& out_tensors,
89
+ torch::stable::Tensor const& a_tensors,
90
+ torch::stable::Tensor const& b_tensors,
91
+ torch::stable::Tensor const& a_scales,
92
+ torch::stable::Tensor const& b_scales,
93
+ torch::stable::Tensor const& expert_offsets,
94
+ torch::stable::Tensor const& problem_sizes,
95
+ torch::stable::Tensor const& a_strides,
96
+ torch::stable::Tensor const& b_strides,
97
+ torch::stable::Tensor const& c_strides,
98
+ bool per_act_token, bool per_out_ch) {
99
static constexpr bool swap_ab = Gemm::swap_ab;
100
101
using ElementAB = typename Gemm::ElementAB;
@@ -98,16 +103,20 @@ void cutlass_group_gemm_caller(
103
104
int num_experts = static_cast<int>(expert_offsets.size(0));
105
101
- auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
106
+ auto stream = get_current_cuda_stream(a_tensors.get_device_index());
107
103
- auto options_int =
104
- torch::TensorOptions().dtype(torch::kInt64).device(a_tensors.device());
108
+ auto device = a_tensors.device();
109
106
- torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
107
- torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
108
- torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
109
- torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
110
- torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
110
+ torch::stable::Tensor a_ptrs = torch::stable::empty(
111
+ {num_experts}, torch::headeronly::ScalarType::Long, std::nullopt, device);
112
+ torch::stable::Tensor b_ptrs = torch::stable::empty(
113
+ {num_experts}, torch::headeronly::ScalarType::Long, std::nullopt, device);
114
+ torch::stable::Tensor out_ptrs = torch::stable::empty(
115
+ {num_experts}, torch::headeronly::ScalarType::Long, std::nullopt, device);
116
+ torch::stable::Tensor a_scales_ptrs = torch::stable::empty(
117
+ {num_experts}, torch::headeronly::ScalarType::Long, std::nullopt, device);
118
+ torch::stable::Tensor b_scales_ptrs = torch::stable::empty(
119
+ {num_experts}, torch::headeronly::ScalarType::Long, std::nullopt, device);
120
121
run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
122
a_scales_ptrs, b_scales_ptrs, a_tensors, b_tensors,
@@ -156,7 +165,7 @@ void cutlass_group_gemm_caller(
165
static_cast<ElementD**>(out_ptrs.data_ptr()),
166
static_cast<StrideC*>(c_strides.data_ptr())};
167
159
- int device_id = a_tensors.device().index();
168
+ int device_id = a_tensors.get_device_index();
169
static const cutlass::KernelHardwareInfo hw_info{
170
device_id, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
171
device_id)};
@@ -170,9 +179,9 @@ void cutlass_group_gemm_caller(
179
CUTLASS_CHECK(gemm_op.can_implement(args));
180
181
size_t workspace_size = gemm_op.get_workspace_size(args);
173
- auto const workspace_options =
174
- torch::TensorOptions().dtype(torch::kUInt8).device(a_tensors.device());
175
- auto workspace = torch::empty(workspace_size, workspace_options);
182
+ auto workspace =
183
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
184
+ std::nullopt, device);
185
186
cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
187
CUTLASS_CHECK(status);
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu
renamed
+47
-32
@@ -1,7 +1,8 @@
1
#include <cudaTypedefs.h>
2
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <torch/all.h>
3
+#include "libtorch_stable/torch_utils.h"
4
+#include <torch/csrc/stable/tensor.h>
5
+#include <torch/headeronly/core/ScalarType.h>
6
7
#include "cutlass/cutlass.h"
8
#include "grouped_mm_c3x.cuh"
@@ -62,21 +63,27 @@ struct sm100_fp8_config_N8192 {
63
};
64
65
template <typename InType, typename OutType>
65
-void run_cutlass_moe_mm_sm100(
66
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
67
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
68
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
69
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
70
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
71
- bool per_act_token, bool per_out_ch) {
72
- TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
73
- TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
74
- TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
75
-
76
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn,
77
- "A tensors must be of type float8_e4m3fn.");
78
- TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn,
79
- "B tensors must be of type float8_e4m3fn.");
66
+void run_cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
67
+ torch::stable::Tensor const& a_tensors,
68
+ torch::stable::Tensor const& b_tensors,
69
+ torch::stable::Tensor const& a_scales,
70
+ torch::stable::Tensor const& b_scales,
71
+ torch::stable::Tensor const& expert_offsets,
72
+ torch::stable::Tensor const& problem_sizes,
73
+ torch::stable::Tensor const& a_strides,
74
+ torch::stable::Tensor const& b_strides,
75
+ torch::stable::Tensor const& c_strides,
76
+ bool per_act_token, bool per_out_ch) {
77
+ STD_TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
78
+ STD_TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
79
+ STD_TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
80
+
81
+ STD_TORCH_CHECK(
82
+ a_tensors.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn,
83
+ "A tensors must be of type float8_e4m3fn.");
84
+ STD_TORCH_CHECK(
85
+ b_tensors.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn,
86
+ "B tensors must be of type float8_e4m3fn.");
87
88
using Cutlass3xGemmDefault = typename sm100_fp8_config_default<
89
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
@@ -107,14 +114,18 @@ void run_cutlass_moe_mm_sm100(
114
}
115
} // namespace
116
110
-void dispatch_moe_mm_sm100(
111
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
112
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
113
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
114
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
115
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
116
- bool per_act_token, bool per_out_ch) {
117
- if (out_tensors.dtype() == torch::kBFloat16) {
117
+void dispatch_moe_mm_sm100(torch::stable::Tensor& out_tensors,
118
+ torch::stable::Tensor const& a_tensors,
119
+ torch::stable::Tensor const& b_tensors,
120
+ torch::stable::Tensor const& a_scales,
121
+ torch::stable::Tensor const& b_scales,
122
+ torch::stable::Tensor const& expert_offsets,
123
+ torch::stable::Tensor const& problem_sizes,
124
+ torch::stable::Tensor const& a_strides,
125
+ torch::stable::Tensor const& b_strides,
126
+ torch::stable::Tensor const& c_strides,
127
+ bool per_act_token, bool per_out_ch) {
128
+ if (out_tensors.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
129
run_cutlass_moe_mm_sm100<cutlass::float_e4m3_t, cutlass::bfloat16_t>(
130
out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
131
problem_sizes, a_strides, b_strides, c_strides, per_act_token,
@@ -127,13 +138,17 @@ void dispatch_moe_mm_sm100(
138
}
139
}
140
130
-void cutlass_moe_mm_sm100(
131
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
132
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
133
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
134
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
135
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
136
- bool per_act_token, bool per_out_ch) {
141
+void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
142
+ torch::stable::Tensor const& a_tensors,
143
+ torch::stable::Tensor const& b_tensors,
144
+ torch::stable::Tensor const& a_scales,
145
+ torch::stable::Tensor const& b_scales,
146
+ torch::stable::Tensor const& expert_offsets,
147
+ torch::stable::Tensor const& problem_sizes,
148
+ torch::stable::Tensor const& a_strides,
149
+ torch::stable::Tensor const& b_strides,
150
+ torch::stable::Tensor const& c_strides,
151
+ bool per_act_token, bool per_out_ch) {
152
dispatch_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
153
expert_offsets, problem_sizes, a_strides, b_strides,
154
c_strides, per_act_token, per_out_ch);
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu
renamed
+47
-32
@@ -1,7 +1,8 @@
1
#include <cudaTypedefs.h>
2
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <torch/all.h>
3
+#include "libtorch_stable/torch_utils.h"
4
+#include <torch/csrc/stable/tensor.h>
5
+#include <torch/headeronly/core/ScalarType.h>
6
7
#include "cutlass/cutlass.h"
8
#include "grouped_mm_c3x.cuh"
@@ -103,21 +104,27 @@ struct sm90_fp8_config_N8192 {
104
};
105
106
template <typename InType, typename OutType>
106
-void run_cutlass_moe_mm_sm90(
107
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
108
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
109
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
110
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
111
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
112
- bool per_act_token, bool per_out_ch) {
113
- TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
114
- TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
115
- TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
116
-
117
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn,
118
- "A tensors must be of type float8_e4m3fn.");
119
- TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn,
120
- "B tensors must be of type float8_e4m3fn.");
107
+void run_cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
108
+ torch::stable::Tensor const& a_tensors,
109
+ torch::stable::Tensor const& b_tensors,
110
+ torch::stable::Tensor const& a_scales,
111
+ torch::stable::Tensor const& b_scales,
112
+ torch::stable::Tensor const& expert_offsets,
113
+ torch::stable::Tensor const& problem_sizes,
114
+ torch::stable::Tensor const& a_strides,
115
+ torch::stable::Tensor const& b_strides,
116
+ torch::stable::Tensor const& c_strides,
117
+ bool per_act_token, bool per_out_ch) {
118
+ STD_TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
119
+ STD_TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
120
+ STD_TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
121
+
122
+ STD_TORCH_CHECK(
123
+ a_tensors.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn,
124
+ "A tensors must be of type float8_e4m3fn.");
125
+ STD_TORCH_CHECK(
126
+ b_tensors.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn,
127
+ "B tensors must be of type float8_e4m3fn.");
128
129
using Cutlass3xGemmN8192 = typename sm90_fp8_config_N8192<
130
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
@@ -163,14 +170,18 @@ void run_cutlass_moe_mm_sm90(
170
}
171
}
172
166
-void dispatch_moe_mm_sm90(
167
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
168
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
169
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
170
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
171
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
172
- bool per_act_token, bool per_out_ch) {
173
- if (out_tensors.dtype() == torch::kBFloat16) {
173
+void dispatch_moe_mm_sm90(torch::stable::Tensor& out_tensors,
174
+ torch::stable::Tensor const& a_tensors,
175
+ torch::stable::Tensor const& b_tensors,
176
+ torch::stable::Tensor const& a_scales,
177
+ torch::stable::Tensor const& b_scales,
178
+ torch::stable::Tensor const& expert_offsets,
179
+ torch::stable::Tensor const& problem_sizes,
180
+ torch::stable::Tensor const& a_strides,
181
+ torch::stable::Tensor const& b_strides,
182
+ torch::stable::Tensor const& c_strides,
183
+ bool per_act_token, bool per_out_ch) {
184
+ if (out_tensors.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
185
run_cutlass_moe_mm_sm90<cutlass::float_e4m3_t, cutlass::bfloat16_t>(
186
out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
187
problem_sizes, a_strides, b_strides, c_strides, per_act_token,
@@ -185,13 +196,17 @@ void dispatch_moe_mm_sm90(
196
197
} // namespace
198
188
-void cutlass_moe_mm_sm90(
189
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
190
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
191
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
192
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
193
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
194
- bool per_act_token, bool per_out_ch) {
199
+void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
200
+ torch::stable::Tensor const& a_tensors,
201
+ torch::stable::Tensor const& b_tensors,
202
+ torch::stable::Tensor const& a_scales,
203
+ torch::stable::Tensor const& b_scales,
204
+ torch::stable::Tensor const& expert_offsets,
205
+ torch::stable::Tensor const& problem_sizes,
206
+ torch::stable::Tensor const& a_strides,
207
+ torch::stable::Tensor const& b_strides,
208
+ torch::stable::Tensor const& c_strides,
209
+ bool per_act_token, bool per_out_ch) {
210
dispatch_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
211
expert_offsets, problem_sizes, a_strides, b_strides,
212
c_strides, per_act_token, per_out_ch);
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu
renamed
+74
-57
@@ -1,9 +1,11 @@
1
#include <cudaTypedefs.h>
2
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <torch/all.h>
3
+#include "libtorch_stable/torch_utils.h"
4
+#include <torch/csrc/stable/tensor.h>
5
+#include <torch/csrc/stable/ops.h>
6
+#include <torch/headeronly/core/ScalarType.h>
7
6
-#include "dispatch_utils.h"
8
+#include "libtorch_stable/dispatch_utils.h"
9
10
#include <iostream>
11
@@ -110,19 +112,22 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
112
}
113
114
namespace {
113
-inline void launch_compute_problem_sizes(
114
- const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
115
- torch::Tensor& problem_sizes2, torch::Tensor& atomic_buffer,
116
- int64_t num_experts, int64_t n, int64_t k, cudaStream_t stream,
117
- const bool swap_ab, const bool is_gated) {
115
+inline void launch_compute_problem_sizes(const torch::stable::Tensor& topk_ids,
116
+ torch::stable::Tensor& problem_sizes1,
117
+ torch::stable::Tensor& problem_sizes2,
118
+ torch::stable::Tensor& atomic_buffer,
119
+ int64_t num_experts, int64_t n,
120
+ int64_t k, cudaStream_t stream,
121
+ const bool swap_ab,
122
+ const bool is_gated) {
123
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
124
120
- auto const* topk_ptr = topk_ids.data_ptr<int32_t>();
121
- auto* ps1_ptr = problem_sizes1.data_ptr<int32_t>();
122
- auto* ps2_ptr = problem_sizes2.data_ptr<int32_t>();
123
- auto* atomic_ptr = atomic_buffer.data_ptr<int32_t>();
125
+ auto const* topk_ptr = topk_ids.const_data_ptr<int32_t>();
126
+ auto* ps1_ptr = problem_sizes1.mutable_data_ptr<int32_t>();
127
+ auto* ps2_ptr = problem_sizes2.mutable_data_ptr<int32_t>();
128
+ auto* atomic_ptr = atomic_buffer.mutable_data_ptr<int32_t>();
129
125
- VLLM_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
130
+ VLLM_STABLE_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
131
compute_problem_sizes<SwapAB><<<num_experts, num_threads, 0, stream>>>(
132
topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,
133
static_cast<int>(topk_ids.numel()), static_cast<int>(n),
@@ -171,46 +176,53 @@ __global__ void compute_problem_sizes_from_expert_offsets(
176
}
177
178
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
174
- const torch::Tensor& expert_first_token_offset,
175
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
176
- const int64_t n, const int64_t k, const bool swap_ab) {
177
- TORCH_CHECK(expert_first_token_offset.is_cuda(),
178
- "expert_first_token_offset must be a CUDA tensor");
179
- TORCH_CHECK(expert_first_token_offset.dtype() == torch::kInt64,
180
- "expert_first_token_offset must be int64");
181
-
182
- TORCH_CHECK(problem_sizes1.is_cuda() && problem_sizes2.is_cuda(),
183
- "problem_sizes must be CUDA tensors");
184
- TORCH_CHECK(problem_sizes1.dtype() == torch::kInt32 &&
185
- problem_sizes2.dtype() == torch::kInt32,
186
- "problem_sizes must be int32");
187
- TORCH_CHECK(problem_sizes1.is_contiguous() && problem_sizes2.is_contiguous(),
188
- "problem_sizes must be contiguous");
189
- TORCH_CHECK(problem_sizes1.dim() == 2 && problem_sizes2.dim() == 2,
190
- "problem_sizes must be 2D tensors");
191
- TORCH_CHECK(problem_sizes1.size(1) == 3 && problem_sizes2.size(1) == 3,
192
- "problem_sizes second dim must be 3");
193
- TORCH_CHECK(problem_sizes1.sizes() == problem_sizes2.sizes(),
194
- "problem_sizes1 and problem_sizes2 must have same shape");
179
+ const torch::stable::Tensor& expert_first_token_offset,
180
+ torch::stable::Tensor& problem_sizes1,
181
+ torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
182
+ const bool swap_ab) {
183
+ STD_TORCH_CHECK(expert_first_token_offset.is_cuda(),
184
+ "expert_first_token_offset must be a CUDA tensor");
185
+ STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
186
+ torch::headeronly::ScalarType::Long,
187
+ "expert_first_token_offset must be int64");
188
+
189
+ STD_TORCH_CHECK(problem_sizes1.is_cuda() && problem_sizes2.is_cuda(),
190
+ "problem_sizes must be CUDA tensors");
191
+ STD_TORCH_CHECK(
192
+ problem_sizes1.scalar_type() == torch::headeronly::ScalarType::Int &&
193
+ problem_sizes2.scalar_type() == torch::headeronly::ScalarType::Int,
194
+ "problem_sizes must be int32");
195
+ STD_TORCH_CHECK(
196
+ problem_sizes1.is_contiguous() && problem_sizes2.is_contiguous(),
197
+ "problem_sizes must be contiguous");
198
+ STD_TORCH_CHECK(problem_sizes1.dim() == 2 && problem_sizes2.dim() == 2,
199
+ "problem_sizes must be 2D tensors");
200
+ STD_TORCH_CHECK(problem_sizes1.size(1) == 3 && problem_sizes2.size(1) == 3,
201
+ "problem_sizes second dim must be 3");
202
+ STD_TORCH_CHECK(problem_sizes1.size(0) == problem_sizes2.size(0) &&
203
+ problem_sizes1.size(1) == problem_sizes2.size(1),
204
+ "problem_sizes1 and problem_sizes2 must have same shape");
205
206
int64_t const num_experts64 = problem_sizes1.size(0);
197
- TORCH_CHECK(expert_first_token_offset.numel() == num_experts64 + 1,
198
- "expert_first_token_offset must have num_experts + 1 elements");
199
- TORCH_CHECK(num_experts64 <= INT32_MAX, "num_experts must fit in int32");
200
- TORCH_CHECK(n <= INT32_MAX && k <= INT32_MAX, "n and k must fit in int32");
207
+ STD_TORCH_CHECK(
208
+ expert_first_token_offset.numel() == num_experts64 + 1,
209
+ "expert_first_token_offset must have num_experts + 1 elements");
210
+ STD_TORCH_CHECK(num_experts64 <= INT32_MAX, "num_experts must fit in int32");
211
+ STD_TORCH_CHECK(n <= INT32_MAX && k <= INT32_MAX,
212
+ "n and k must fit in int32");
213
214
int const num_experts = static_cast<int>(num_experts64);
203
- auto stream = at::cuda::getCurrentCUDAStream(
204
- expert_first_token_offset.device().index());
215
+ auto stream =
216
+ get_current_cuda_stream(expert_first_token_offset.get_device_index());
217
218
int const threads = (num_experts < 256) ? num_experts : 256;
219
int const blocks = (num_experts + threads - 1) / threads;
220
209
- auto const* offsets_ptr = expert_first_token_offset.data_ptr<int64_t>();
210
- auto* ps1_ptr = problem_sizes1.data_ptr<int32_t>();
211
- auto* ps2_ptr = problem_sizes2.data_ptr<int32_t>();
221
+ auto const* offsets_ptr = expert_first_token_offset.const_data_ptr<int64_t>();
222
+ auto* ps1_ptr = problem_sizes1.mutable_data_ptr<int32_t>();
223
+ auto* ps2_ptr = problem_sizes2.mutable_data_ptr<int32_t>();
224
213
- VLLM_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
225
+ VLLM_STABLE_DISPATCH_BOOL(swap_ab, SwapAB, [&] {
226
compute_problem_sizes_from_expert_offsets<SwapAB>
227
<<<blocks, threads, 0, stream>>>(offsets_ptr, ps1_ptr, ps2_ptr,
228
num_experts, static_cast<int>(n),
@@ -219,16 +231,19 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
231
}
232
233
void get_cutlass_moe_mm_data_caller(
222
- const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
223
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
224
- torch::Tensor& input_permutation, torch::Tensor& output_permutation,
225
- const int64_t num_experts, const int64_t n, const int64_t k,
226
- const std::optional<torch::Tensor>& blockscale_offsets,
234
+ const torch::stable::Tensor& topk_ids,
235
+ torch::stable::Tensor& expert_offsets,
236
+ torch::stable::Tensor& problem_sizes1,
237
+ torch::stable::Tensor& problem_sizes2,
238
+ torch::stable::Tensor& input_permutation,
239
+ torch::stable::Tensor& output_permutation, const int64_t num_experts,
240
+ const int64_t n, const int64_t k,
241
+ const std::optional<torch::stable::Tensor>& blockscale_offsets,
242
const bool is_gated) {
228
- auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
229
- auto options_int32 =
230
- torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
231
- torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);
243
+ auto device = topk_ids.device();
244
+ auto stream = get_current_cuda_stream(device.index());
245
+ torch::stable::Tensor atomic_buffer = torch::stable::new_zeros(
246
+ topk_ids, {num_experts}, torch::headeronly::ScalarType::Int);
247
248
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
249
@@ -290,11 +305,13 @@ __global__ void compute_batched_moe_data(
305
}
306
307
void get_cutlass_batched_moe_mm_data_caller(
293
- torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1,
294
- torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens,
308
+ torch::stable::Tensor& expert_offsets,
309
+ torch::stable::Tensor& problem_sizes1,
310
+ torch::stable::Tensor& problem_sizes2,
311
+ const torch::stable::Tensor& expert_num_tokens,
312
const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
313
const int64_t k) {
297
- auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index());
314
+ auto stream = get_current_cuda_stream(expert_offsets.get_device_index());
315
316
if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) {
317
compute_batched_moe_data<false><<<1, num_local_experts, 0, stream>>>(
@@ -311,4 +328,4 @@ void get_cutlass_batched_moe_mm_data_caller(
328
static_cast<const int32_t*>(expert_num_tokens.data_ptr()), padded_m, n,
329
k);
330
}
314
-}
\ No newline at end of file
331
+}
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu
new
+220
@@ -0,0 +1,220 @@
1
+#include <stddef.h>
2
+#include <torch/csrc/stable/tensor.h>
3
+#include <torch/headeronly/core/ScalarType.h>
4
+#include "cutlass/cutlass.h"
5
+
6
+#include "scaled_mm_c2x.cuh"
7
+#include "scaled_mm_c2x_sm75_dispatch.cuh"
8
+#include "scaled_mm_c2x_sm80_dispatch.cuh"
9
+#include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
10
+#include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
11
+
12
+#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
13
+
14
+using namespace vllm;
15
+
16
+/*
17
+ This file defines quantized GEMM operations using the CUTLASS 2.x API, for
18
+ NVIDIA GPUs with SM versions prior to sm90 (Hopper).
19
+*/
20
+
21
+template <template <typename, typename> typename Epilogue,
22
+ typename... EpilogueArgs>
23
+void cutlass_scaled_mm_sm75_epilogue(torch::stable::Tensor& out,
24
+ torch::stable::Tensor const& a,
25
+ torch::stable::Tensor const& b,
26
+ EpilogueArgs&&... epilogue_args) {
27
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
28
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
29
+
30
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
31
+ return cutlass_gemm_sm75_dispatch<int8_t, cutlass::bfloat16_t, Epilogue>(
32
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
33
+ } else {
34
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
35
+ return cutlass_gemm_sm75_dispatch<int8_t, cutlass::half_t, Epilogue>(
36
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
37
+ }
38
+}
39
+
40
+void cutlass_scaled_mm_sm75(torch::stable::Tensor& out,
41
+ torch::stable::Tensor const& a,
42
+ torch::stable::Tensor const& b,
43
+ torch::stable::Tensor const& a_scales,
44
+ torch::stable::Tensor const& b_scales,
45
+ std::optional<torch::stable::Tensor> const& bias) {
46
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
47
+ torch::headeronly::ScalarType::Float);
48
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
49
+ torch::headeronly::ScalarType::Float);
50
+ if (bias) {
51
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
52
+ "currently bias dtype must match output dtype ",
53
+ out.scalar_type());
54
+ return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBias>(
55
+ out, a, b, a_scales, b_scales, *bias);
56
+ } else {
57
+ return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogue>(
58
+ out, a, b, a_scales, b_scales);
59
+ }
60
+}
61
+
62
+void cutlass_scaled_mm_azp_sm75(
63
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
64
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
65
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
66
+ std::optional<torch::stable::Tensor> const& azp,
67
+ std::optional<torch::stable::Tensor> const& bias) {
68
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
69
+ torch::headeronly::ScalarType::Float);
70
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
71
+ torch::headeronly::ScalarType::Float);
72
+
73
+ if (azp) {
74
+ return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
75
+ out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
76
+ } else {
77
+ return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBiasAzp>(
78
+ out, a, b, a_scales, b_scales, azp_adj, bias);
79
+ }
80
+}
81
+
82
+template <template <typename, typename> typename Epilogue,
83
+ typename... EpilogueArgs>
84
+void cutlass_scaled_mm_sm80_epilogue(torch::stable::Tensor& out,
85
+ torch::stable::Tensor const& a,
86
+ torch::stable::Tensor const& b,
87
+ EpilogueArgs&&... epilogue_args) {
88
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
89
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
90
+
91
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
92
+ return cutlass_gemm_sm80_dispatch<int8_t, cutlass::bfloat16_t, Epilogue>(
93
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
94
+ } else {
95
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
96
+ return cutlass_gemm_sm80_dispatch<int8_t, cutlass::half_t, Epilogue>(
97
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
98
+ }
99
+}
100
+
101
+void cutlass_scaled_mm_sm80(torch::stable::Tensor& out,
102
+ torch::stable::Tensor const& a,
103
+ torch::stable::Tensor const& b,
104
+ torch::stable::Tensor const& a_scales,
105
+ torch::stable::Tensor const& b_scales,
106
+ std::optional<torch::stable::Tensor> const& bias) {
107
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
108
+ torch::headeronly::ScalarType::Float);
109
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
110
+ torch::headeronly::ScalarType::Float);
111
+ if (bias) {
112
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
113
+ "currently bias dtype must match output dtype ",
114
+ out.scalar_type());
115
+ return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBias>(
116
+ out, a, b, a_scales, b_scales, *bias);
117
+ } else {
118
+ return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogue>(
119
+ out, a, b, a_scales, b_scales);
120
+ }
121
+}
122
+
123
+void cutlass_scaled_mm_azp_sm80(
124
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
125
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
126
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
127
+ std::optional<torch::stable::Tensor> const& azp,
128
+ std::optional<torch::stable::Tensor> const& bias) {
129
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
130
+ torch::headeronly::ScalarType::Float);
131
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
132
+ torch::headeronly::ScalarType::Float);
133
+
134
+ if (azp) {
135
+ return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
136
+ out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
137
+ } else {
138
+ return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzp>(
139
+ out, a, b, a_scales, b_scales, azp_adj, bias);
140
+ }
141
+}
142
+
143
+template <template <typename, typename> typename Epilogue,
144
+ typename... EpilogueArgs>
145
+void cutlass_scaled_mm_sm89_epilogue(torch::stable::Tensor& out,
146
+ torch::stable::Tensor const& a,
147
+ torch::stable::Tensor const& b,
148
+ EpilogueArgs&&... epilogue_args) {
149
+ if (a.scalar_type() == torch::headeronly::ScalarType::Char) {
150
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
151
+
152
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
153
+ return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::bfloat16_t,
154
+ Epilogue>(
155
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
156
+ } else {
157
+ assert(out.scalar_type() == torch::headeronly::ScalarType::Half);
158
+ return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::half_t, Epilogue>(
159
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
160
+ }
161
+ } else {
162
+ STD_TORCH_CHECK(a.scalar_type() ==
163
+ torch::headeronly::ScalarType::Float8_e4m3fn);
164
+ STD_TORCH_CHECK(b.scalar_type() ==
165
+ torch::headeronly::ScalarType::Float8_e4m3fn);
166
+
167
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
168
+ return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
169
+ cutlass::bfloat16_t, Epilogue>(
170
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
171
+ } else {
172
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
173
+ return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
174
+ cutlass::half_t, Epilogue>(
175
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
176
+ }
177
+ }
178
+}
179
+
180
+void cutlass_scaled_mm_sm89(torch::stable::Tensor& out,
181
+ torch::stable::Tensor const& a,
182
+ torch::stable::Tensor const& b,
183
+ torch::stable::Tensor const& a_scales,
184
+ torch::stable::Tensor const& b_scales,
185
+ std::optional<torch::stable::Tensor> const& bias) {
186
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
187
+ torch::headeronly::ScalarType::Float);
188
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
189
+ torch::headeronly::ScalarType::Float);
190
+ if (bias) {
191
+ STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
192
+ "currently bias dtype must match output dtype ",
193
+ out.scalar_type());
194
+ return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
195
+ out, a, b, a_scales, b_scales, *bias);
196
+ } else {
197
+ return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
198
+ out, a, b, a_scales, b_scales);
199
+ }
200
+}
201
+
202
+void cutlass_scaled_mm_azp_sm89(
203
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
204
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
205
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
206
+ std::optional<torch::stable::Tensor> const& azp,
207
+ std::optional<torch::stable::Tensor> const& bias) {
208
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
209
+ torch::headeronly::ScalarType::Float);
210
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
211
+ torch::headeronly::ScalarType::Float);
212
+
213
+ if (azp) {
214
+ return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
215
+ out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
216
+ } else {
217
+ return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBiasAzp>(
218
+ out, a, b, a_scales, b_scales, azp_adj, bias);
219
+ }
220
+}
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh
renamed
+16
-13
@@ -1,8 +1,9 @@
1
#pragma once
2
#include <stddef.h>
3
-#include <torch/all.h>
3
+#include <torch/csrc/stable/tensor.h>
4
+#include <torch/csrc/stable/ops.h>
5
5
-#include <ATen/cuda/CUDAContext.h>
6
+#include "libtorch_stable/torch_utils.h"
7
8
// clang-format will break include orders
9
// clang-format off
@@ -95,8 +96,9 @@ struct cutlass_2x_gemm {
96
};
97
98
template <typename Gemm, typename... EpilogueArgs>
98
-inline void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
99
- torch::Tensor const& b,
99
+inline void cutlass_gemm_caller(torch::stable::Tensor& out,
100
+ torch::stable::Tensor const& a,
101
+ torch::stable::Tensor const& b,
102
EpilogueArgs&&... epilogue_params) {
103
using ElementAB = typename Gemm::ElementAB;
104
using ElementD = typename Gemm::ElementD;
@@ -149,11 +151,12 @@ inline void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
151
// Launch the CUTLASS GEMM kernel.
152
typename Gemm::Op gemm_op;
153
size_t workspace_size = gemm_op.get_workspace_size(args);
152
- auto const workspace_options =
153
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
154
- auto workspace = torch::empty(workspace_size, workspace_options);
154
+ auto device = a.device();
155
+ auto workspace =
156
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
157
+ std::nullopt, device);
158
156
- auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
159
+ auto stream = get_current_cuda_stream(device.index());
160
161
CUTLASS_CHECK(gemm_op.can_implement(args));
162
cutlass::Status status = gemm_op(args, workspace.data_ptr(), stream);
@@ -161,9 +164,9 @@ inline void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
164
}
165
166
template <typename Gemm, typename FallbackGemm, typename... EpilogueArgs>
164
-inline void fallback_cutlass_gemm_caller(torch::Tensor& out,
165
- torch::Tensor const& a,
166
- torch::Tensor const& b,
167
+inline void fallback_cutlass_gemm_caller(torch::stable::Tensor& out,
168
+ torch::stable::Tensor const& a,
169
+ torch::stable::Tensor const& b,
170
EpilogueArgs&&... args) {
171
// In some cases, the GPU isn't able to accommodate the
172
// shared memory requirements of the Gemm. In such cases, use
@@ -180,8 +183,8 @@ inline void fallback_cutlass_gemm_caller(torch::Tensor& out,
183
return cutlass_gemm_caller<Gemm>(out, a, b,
184
std::forward<EpilogueArgs>(args)...);
185
} else {
183
- TORCH_CHECK(fallback_gemm_shared_mem_size <=
184
- max_shared_mem_per_block_opt_in);
186
+ STD_TORCH_CHECK(fallback_gemm_shared_mem_size <=
187
+ max_shared_mem_per_block_opt_in);
188
return cutlass_gemm_caller<FallbackGemm>(
189
out, a, b, std::forward<EpilogueArgs>(args)...);
190
}
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh
renamed
+7
-5
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm_c2x.cuh"
6
7
/**
@@ -70,13 +72,13 @@ struct sm75_config_M32 {
72
template <typename InType, typename OutType,
73
template <typename, typename> typename Epilogue,
74
typename... EpilogueArgs>
73
-inline void cutlass_gemm_sm75_dispatch(torch::Tensor& out,
74
- torch::Tensor const& a,
75
- torch::Tensor const& b,
75
+inline void cutlass_gemm_sm75_dispatch(torch::stable::Tensor& out,
76
+ torch::stable::Tensor const& a,
77
+ torch::stable::Tensor const& b,
78
EpilogueArgs&&... args) {
79
static_assert(std::is_same<InType, int8_t>());
78
- TORCH_CHECK(a.dtype() == torch::kInt8);
79
- TORCH_CHECK(b.dtype() == torch::kInt8);
80
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
81
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
82
83
using Cutlass2xGemmDefault =
84
typename sm75_config_default<InType, OutType, Epilogue>::Cutlass2xGemm;
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh
renamed
+7
-5
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm_c2x.cuh"
6
7
/**
@@ -72,13 +74,13 @@ struct sm80_config_M16 {
74
template <typename InType, typename OutType,
75
template <typename, typename> typename Epilogue,
76
typename... EpilogueArgs>
75
-inline void cutlass_gemm_sm80_dispatch(torch::Tensor& out,
76
- torch::Tensor const& a,
77
- torch::Tensor const& b,
77
+inline void cutlass_gemm_sm80_dispatch(torch::stable::Tensor& out,
78
+ torch::stable::Tensor const& a,
79
+ torch::stable::Tensor const& b,
80
EpilogueArgs&&... args) {
81
static_assert(std::is_same<InType, int8_t>());
80
- TORCH_CHECK(a.dtype() == torch::kInt8);
81
- TORCH_CHECK(b.dtype() == torch::kInt8);
82
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
83
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
84
85
using Cutlass2xGemmDefault =
86
typename sm80_config_default<InType, OutType, Epilogue>::Cutlass2xGemm;
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh
renamed
+39
-23
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm_c2x.cuh"
6
#include "cutlass/float8.h"
7
@@ -34,10 +36,12 @@ struct sm89_fp8_config_default {
36
template <typename InType, typename OutType,
37
template <typename, typename> typename Epilogue,
38
typename... EpilogueArgs>
37
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
38
- torch::Tensor const& b, EpilogueArgs&&... args) {
39
+ static void dispatch(torch::stable::Tensor& out,
40
+ torch::stable::Tensor const& a,
41
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
42
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
40
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
43
+ STD_TORCH_CHECK(a.scalar_type() ==
44
+ torch::headeronly::ScalarType::Float8_e4m3fn);
45
46
using FallbackGemm =
47
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -84,10 +88,12 @@ struct sm89_fp8_config_M256 {
88
template <typename InType, typename OutType,
89
template <typename, typename> typename Epilogue,
90
typename... EpilogueArgs>
87
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
88
- torch::Tensor const& b, EpilogueArgs&&... args) {
91
+ static void dispatch(torch::stable::Tensor& out,
92
+ torch::stable::Tensor const& a,
93
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
94
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
90
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
95
+ STD_TORCH_CHECK(a.scalar_type() ==
96
+ torch::headeronly::ScalarType::Float8_e4m3fn);
97
98
using FallbackGemm =
99
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -125,10 +131,12 @@ struct sm89_fp8_config_M128 {
131
template <typename InType, typename OutType,
132
template <typename, typename> typename Epilogue,
133
typename... EpilogueArgs>
128
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
129
- torch::Tensor const& b, EpilogueArgs&&... args) {
134
+ static void dispatch(torch::stable::Tensor& out,
135
+ torch::stable::Tensor const& a,
136
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
137
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
131
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
138
+ STD_TORCH_CHECK(a.scalar_type() ==
139
+ torch::headeronly::ScalarType::Float8_e4m3fn);
140
141
using FallbackGemm =
142
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -173,10 +181,12 @@ struct sm89_fp8_config_M64 {
181
template <typename InType, typename OutType,
182
template <typename, typename> typename Epilogue,
183
typename... EpilogueArgs>
176
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
177
- torch::Tensor const& b, EpilogueArgs&&... args) {
184
+ static void dispatch(torch::stable::Tensor& out,
185
+ torch::stable::Tensor const& a,
186
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
187
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
179
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
188
+ STD_TORCH_CHECK(a.scalar_type() ==
189
+ torch::headeronly::ScalarType::Float8_e4m3fn);
190
191
using FallbackGemm =
192
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -227,10 +237,12 @@ struct sm89_fp8_config_M32 {
237
template <typename InType, typename OutType,
238
template <typename, typename> typename Epilogue,
239
typename... EpilogueArgs>
230
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
231
- torch::Tensor const& b, EpilogueArgs&&... args) {
240
+ static void dispatch(torch::stable::Tensor& out,
241
+ torch::stable::Tensor const& a,
242
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
243
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
233
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
244
+ STD_TORCH_CHECK(a.scalar_type() ==
245
+ torch::headeronly::ScalarType::Float8_e4m3fn);
246
247
using FallbackGemm =
248
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -280,10 +292,12 @@ struct sm89_fp8_config_M16 {
292
template <typename InType, typename OutType,
293
template <typename, typename> typename Epilogue,
294
typename... EpilogueArgs>
283
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
284
- torch::Tensor const& b, EpilogueArgs&&... args) {
295
+ static void dispatch(torch::stable::Tensor& out,
296
+ torch::stable::Tensor const& a,
297
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
298
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
286
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
299
+ STD_TORCH_CHECK(a.scalar_type() ==
300
+ torch::headeronly::ScalarType::Float8_e4m3fn);
301
302
using FallbackGemm =
303
typename sm89_fp8_fallback_gemm<InType, OutType,
@@ -326,13 +340,15 @@ struct sm89_fp8_config_M16 {
340
template <typename InType, typename OutType,
341
template <typename, typename> typename Epilogue,
342
typename... EpilogueArgs>
329
-inline void cutlass_gemm_sm89_fp8_dispatch(torch::Tensor& out,
330
- torch::Tensor const& a,
331
- torch::Tensor const& b,
343
+inline void cutlass_gemm_sm89_fp8_dispatch(torch::stable::Tensor& out,
344
+ torch::stable::Tensor const& a,
345
+ torch::stable::Tensor const& b,
346
EpilogueArgs&&... args) {
347
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
334
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
335
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
348
+ STD_TORCH_CHECK(a.scalar_type() ==
349
+ torch::headeronly::ScalarType::Float8_e4m3fn);
350
+ STD_TORCH_CHECK(b.scalar_type() ==
351
+ torch::headeronly::ScalarType::Float8_e4m3fn);
352
353
uint32_t const m = a.size(0);
354
uint32_t const mp2 =
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh
renamed
+31
-23
@@ -1,5 +1,7 @@
1
#pragma once
2
3
+#include <torch/headeronly/util/shim_utils.h>
4
+
5
#include "scaled_mm_c2x.cuh"
6
7
/**
@@ -32,10 +34,11 @@ struct sm89_int8_config_default {
34
template <typename InType, typename OutType,
35
template <typename, typename> typename Epilogue,
36
typename... EpilogueArgs>
35
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
36
- torch::Tensor const& b, EpilogueArgs&&... args) {
37
+ static void dispatch(torch::stable::Tensor& out,
38
+ torch::stable::Tensor const& a,
39
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
40
static_assert(std::is_same<InType, int8_t>());
38
- TORCH_CHECK(a.dtype() == torch::kInt8);
41
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
42
43
using FallbackGemm =
44
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -88,10 +91,11 @@ struct sm89_int8_config_M256 {
91
template <typename InType, typename OutType,
92
template <typename, typename> typename Epilogue,
93
typename... EpilogueArgs>
91
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
92
- torch::Tensor const& b, EpilogueArgs&&... args) {
94
+ static void dispatch(torch::stable::Tensor& out,
95
+ torch::stable::Tensor const& a,
96
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
97
static_assert(std::is_same<InType, int8_t>());
94
- TORCH_CHECK(a.dtype() == torch::kInt8);
98
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
99
100
using FallbackGemm =
101
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -143,10 +147,11 @@ struct sm89_int8_config_M128 {
147
template <typename InType, typename OutType,
148
template <typename, typename> typename Epilogue,
149
typename... EpilogueArgs>
146
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
147
- torch::Tensor const& b, EpilogueArgs&&... args) {
150
+ static void dispatch(torch::stable::Tensor& out,
151
+ torch::stable::Tensor const& a,
152
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
153
static_assert(std::is_same<InType, int8_t>());
149
- TORCH_CHECK(a.dtype() == torch::kInt8);
154
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
155
156
using FallbackGemm =
157
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -193,10 +198,11 @@ struct sm89_int8_config_M64 {
198
template <typename InType, typename OutType,
199
template <typename, typename> typename Epilogue,
200
typename... EpilogueArgs>
196
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
197
- torch::Tensor const& b, EpilogueArgs&&... args) {
201
+ static void dispatch(torch::stable::Tensor& out,
202
+ torch::stable::Tensor const& a,
203
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
204
static_assert(std::is_same<InType, int8_t>());
199
- TORCH_CHECK(a.dtype() == torch::kInt8);
205
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
206
207
using FallbackGemm =
208
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -234,10 +240,11 @@ struct sm89_int8_config_M32 {
240
template <typename InType, typename OutType,
241
template <typename, typename> typename Epilogue,
242
typename... EpilogueArgs>
237
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
238
- torch::Tensor const& b, EpilogueArgs&&... args) {
243
+ static void dispatch(torch::stable::Tensor& out,
244
+ torch::stable::Tensor const& a,
245
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
246
static_assert(std::is_same<InType, int8_t>());
240
- TORCH_CHECK(a.dtype() == torch::kInt8);
247
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
248
249
using FallbackGemm =
250
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -276,10 +283,11 @@ struct sm89_int8_config_M16 {
283
template <typename InType, typename OutType,
284
template <typename, typename> typename Epilogue,
285
typename... EpilogueArgs>
279
- static void dispatch(torch::Tensor& out, torch::Tensor const& a,
280
- torch::Tensor const& b, EpilogueArgs&&... args) {
286
+ static void dispatch(torch::stable::Tensor& out,
287
+ torch::stable::Tensor const& a,
288
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
289
static_assert(std::is_same<InType, int8_t>());
282
- TORCH_CHECK(a.dtype() == torch::kInt8);
290
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
291
292
using FallbackGemm =
293
typename sm89_int8_fallback_gemm<InType, OutType,
@@ -311,13 +319,13 @@ struct sm89_int8_config_M16 {
319
template <typename InType, typename OutType,
320
template <typename, typename> typename Epilogue,
321
typename... EpilogueArgs>
314
-inline void cutlass_gemm_sm89_int8_dispatch(torch::Tensor& out,
315
- torch::Tensor const& a,
316
- torch::Tensor const& b,
322
+inline void cutlass_gemm_sm89_int8_dispatch(torch::stable::Tensor& out,
323
+ torch::stable::Tensor const& a,
324
+ torch::stable::Tensor const& b,
325
EpilogueArgs&&... args) {
326
static_assert(std::is_same<InType, int8_t>());
319
- TORCH_CHECK(a.dtype() == torch::kInt8);
320
- TORCH_CHECK(b.dtype() == torch::kInt8);
327
+ STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Char);
328
+ STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Char);
329
330
uint32_t const m = a.size(0);
331
uint32_t const mp2 =
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu
renamed
+6
-5
@@ -8,11 +8,12 @@
8
9
#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
10
11
-void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
12
- torch::Tensor const& b,
13
- torch::Tensor const& a_scales,
14
- torch::Tensor const& b_scales,
15
- std::optional<torch::Tensor> const& bias) {
11
+void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
12
+ torch::stable::Tensor const& a,
13
+ torch::stable::Tensor const& b,
14
+ torch::stable::Tensor const& a_scales,
15
+ torch::stable::Tensor const& b_scales,
16
+ std::optional<torch::stable::Tensor> const& bias) {
17
dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
18
vllm::cutlass_scaled_mm_sm100_fp8,
19
nullptr, // int8 not supported on SM100
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu
renamed
+6
-5
@@ -8,11 +8,12 @@
8
9
#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
10
11
-void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
12
- torch::Tensor const& b,
13
- torch::Tensor const& a_scales,
14
- torch::Tensor const& b_scales,
15
- std::optional<torch::Tensor> const& bias) {
11
+void cutlass_scaled_mm_sm120(torch::stable::Tensor& c,
12
+ torch::stable::Tensor const& a,
13
+ torch::stable::Tensor const& b,
14
+ torch::stable::Tensor const& a_scales,
15
+ torch::stable::Tensor const& b_scales,
16
+ std::optional<torch::stable::Tensor> const& bias) {
17
dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
18
vllm::cutlass_scaled_mm_sm120_fp8,
19
nullptr, // int8 not supported on SM120
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu
new
+38
@@ -0,0 +1,38 @@
1
+#include "c3x/scaled_mm_helper.hpp"
2
+#include "c3x/scaled_mm_kernels.hpp"
3
+
4
+/*
5
+ This file defines quantized GEMM operations using the CUTLASS 3.x API, for
6
+ NVIDIA GPUs with sm90a (Hopper).
7
+*/
8
+
9
+#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
10
+
11
+void cutlass_scaled_mm_sm90(torch::stable::Tensor& c,
12
+ torch::stable::Tensor const& a,
13
+ torch::stable::Tensor const& b,
14
+ torch::stable::Tensor const& a_scales,
15
+ torch::stable::Tensor const& b_scales,
16
+ std::optional<torch::stable::Tensor> const& bias) {
17
+ dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
18
+ vllm::cutlass_scaled_mm_sm90_fp8,
19
+ vllm::cutlass_scaled_mm_sm90_int8,
20
+ vllm::cutlass_scaled_mm_blockwise_sm90_fp8);
21
+}
22
+
23
+void cutlass_scaled_mm_azp_sm90(
24
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
25
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
26
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
27
+ std::optional<torch::stable::Tensor> const& azp,
28
+ std::optional<torch::stable::Tensor> const& bias) {
29
+ STD_TORCH_CHECK(a_scales.scalar_type() ==
30
+ torch::headeronly::ScalarType::Float);
31
+ STD_TORCH_CHECK(b_scales.scalar_type() ==
32
+ torch::headeronly::ScalarType::Float);
33
+
34
+ vllm::cutlass_scaled_mm_azp_sm90_int8(out, a, b, a_scales, b_scales, azp_adj,
35
+ azp, bias);
36
+}
37
+
38
+#endif
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu
new
+451
@@ -0,0 +1,451 @@
1
+#include <cudaTypedefs.h>
2
+
3
+#include <torch/csrc/stable/tensor.h>
4
+
5
+#include "libtorch_stable/torch_utils.h"
6
+
7
+#include "cutlass_extensions/common.hpp"
8
+
9
+void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
10
+ torch::stable::Tensor const& a,
11
+ torch::stable::Tensor const& b,
12
+ torch::stable::Tensor const& a_scales,
13
+ torch::stable::Tensor const& b_scales,
14
+ std::optional<torch::stable::Tensor> const& bias);
15
+
16
+void cutlass_scaled_mm_sm80(torch::stable::Tensor& c,
17
+ torch::stable::Tensor const& a,
18
+ torch::stable::Tensor const& b,
19
+ torch::stable::Tensor const& a_scales,
20
+ torch::stable::Tensor const& b_scales,
21
+ std::optional<torch::stable::Tensor> const& bias);
22
+
23
+void cutlass_scaled_mm_sm89(torch::stable::Tensor& c,
24
+ torch::stable::Tensor const& a,
25
+ torch::stable::Tensor const& b,
26
+ torch::stable::Tensor const& a_scales,
27
+ torch::stable::Tensor const& b_scales,
28
+ std::optional<torch::stable::Tensor> const& bias);
29
+
30
+#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
31
+void cutlass_scaled_mm_sm90(torch::stable::Tensor& c,
32
+ torch::stable::Tensor const& a,
33
+ torch::stable::Tensor const& b,
34
+ torch::stable::Tensor const& a_scales,
35
+ torch::stable::Tensor const& b_scales,
36
+ std::optional<torch::stable::Tensor> const& bias);
37
+#endif
38
+#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
39
+void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
40
+ torch::stable::Tensor const& a_tensors,
41
+ torch::stable::Tensor const& b_tensors,
42
+ torch::stable::Tensor const& a_scales,
43
+ torch::stable::Tensor const& b_scales,
44
+ torch::stable::Tensor const& expert_offsets,
45
+ torch::stable::Tensor const& problem_sizes,
46
+ torch::stable::Tensor const& a_strides,
47
+ torch::stable::Tensor const& b_strides,
48
+ torch::stable::Tensor const& c_strides,
49
+ bool per_act_token, bool per_out_ch);
50
+
51
+#endif
52
+
53
+#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
54
+void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
55
+ torch::stable::Tensor const& a_tensors,
56
+ torch::stable::Tensor const& b_tensors,
57
+ torch::stable::Tensor const& a_scales,
58
+ torch::stable::Tensor const& b_scales,
59
+ torch::stable::Tensor const& expert_offsets,
60
+ torch::stable::Tensor const& problem_sizes,
61
+ torch::stable::Tensor const& a_strides,
62
+ torch::stable::Tensor const& b_strides,
63
+ torch::stable::Tensor const& c_strides,
64
+ bool per_act_token, bool per_out_ch);
65
+#endif
66
+
67
+#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
68
+void cutlass_scaled_mm_sm120(torch::stable::Tensor& c,
69
+ torch::stable::Tensor const& a,
70
+ torch::stable::Tensor const& b,
71
+ torch::stable::Tensor const& a_scales,
72
+ torch::stable::Tensor const& b_scales,
73
+ std::optional<torch::stable::Tensor> const& bias);
74
+#endif
75
+
76
+#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
77
+void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
78
+ torch::stable::Tensor const& a,
79
+ torch::stable::Tensor const& b,
80
+ torch::stable::Tensor const& a_scales,
81
+ torch::stable::Tensor const& b_scales,
82
+ std::optional<torch::stable::Tensor> const& bias);
83
+#endif
84
+
85
+#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
86
+ (defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
87
+ (defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
88
+void get_cutlass_moe_mm_data_caller(
89
+ const torch::stable::Tensor& topk_ids,
90
+ torch::stable::Tensor& expert_offsets,
91
+ torch::stable::Tensor& problem_sizes1,
92
+ torch::stable::Tensor& problem_sizes2,
93
+ torch::stable::Tensor& input_permutation,
94
+ torch::stable::Tensor& output_permutation, const int64_t num_experts,
95
+ const int64_t n, const int64_t k,
96
+ const std::optional<torch::stable::Tensor>& blockscale_offsets,
97
+ const bool is_gated);
98
+
99
+void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
100
+ const torch::stable::Tensor& expert_first_token_offset,
101
+ torch::stable::Tensor& problem_sizes1,
102
+ torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
103
+ const bool swap_ab);
104
+
105
+void get_cutlass_batched_moe_mm_data_caller(
106
+ torch::stable::Tensor& expert_offsets,
107
+ torch::stable::Tensor& problem_sizes1,
108
+ torch::stable::Tensor& problem_sizes2,
109
+ const torch::stable::Tensor& expert_num_tokens,
110
+ const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
111
+ const int64_t k);
112
+#endif
113
+
114
+void cutlass_scaled_mm_azp_sm75(
115
+ torch::stable::Tensor& c, torch::stable::Tensor const& a,
116
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
117
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
118
+ std::optional<torch::stable::Tensor> const& azp,
119
+ std::optional<torch::stable::Tensor> const& bias);
120
+
121
+void cutlass_scaled_mm_azp_sm80(
122
+ torch::stable::Tensor& c, torch::stable::Tensor const& a,
123
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
124
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
125
+ std::optional<torch::stable::Tensor> const& azp,
126
+ std::optional<torch::stable::Tensor> const& bias);
127
+
128
+void cutlass_scaled_mm_azp_sm89(
129
+ torch::stable::Tensor& c, torch::stable::Tensor const& a,
130
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
131
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
132
+ std::optional<torch::stable::Tensor> const& azp,
133
+ std::optional<torch::stable::Tensor> const& bias);
134
+
135
+#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
136
+void cutlass_scaled_mm_azp_sm90(
137
+ torch::stable::Tensor& c, torch::stable::Tensor const& a,
138
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
139
+ torch::stable::Tensor const& b_scales, torch::stable::Tensor const& azp_adj,
140
+ std::optional<torch::stable::Tensor> const& azp,
141
+ std::optional<torch::stable::Tensor> const& bias);
142
+#endif
143
+
144
+bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability) {
145
+ // CUTLASS FP8 kernels need at least
146
+ // CUDA 12.0 on SM90 systems (Hopper)
147
+ // CUDA 12.4 on SM89 systems (Lovelace)
148
+
149
+#if defined CUDA_VERSION
150
+ if (cuda_device_capability >= 90) {
151
+ return CUDA_VERSION >= 12000;
152
+ } else if (cuda_device_capability >= 89) {
153
+ return CUDA_VERSION >= 12040;
154
+ }
155
+#endif
156
+
157
+ return false;
158
+}
159
+
160
+bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
161
+ // CUTLASS block-quantized FP8 kernels need at least CUDA 12.0
162
+ // and at least SM90 (Hopper)
163
+
164
+#if defined CUDA_VERSION
165
+ if (cuda_device_capability >= 100) {
166
+ return CUDA_VERSION >= 12080;
167
+ } else if (cuda_device_capability >= 90) {
168
+ return CUDA_VERSION >= 12000;
169
+ }
170
+#endif
171
+
172
+ return false;
173
+}
174
+
175
+bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
176
+ // CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
177
+ // or CUDA 12.8 and SM100 (Blackwell)
178
+
179
+#if defined CUDA_VERSION
180
+ if (cuda_device_capability >= 100) {
181
+ return CUDA_VERSION >= 12080;
182
+ }
183
+ if (cuda_device_capability >= 90) {
184
+ return CUDA_VERSION >= 12030;
185
+ }
186
+#endif
187
+
188
+ return false;
189
+}
190
+
191
+void cutlass_scaled_mm(torch::stable::Tensor& c, torch::stable::Tensor const& a,
192
+ torch::stable::Tensor const& b,
193
+ torch::stable::Tensor const& a_scales,
194
+ torch::stable::Tensor const& b_scales,
195
+ std::optional<torch::stable::Tensor> const& bias) {
196
+ // Checks for conformality
197
+ STD_TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
198
+ STD_TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
199
+ b.size(1) == c.size(1));
200
+
201
+ // Check for strides and alignment
202
+ STD_TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
203
+ STD_TORCH_CHECK(b.stride(0) == 1); // Column-major
204
+ STD_TORCH_CHECK(c.stride(0) % 16 == 0 &&
205
+ b.stride(1) % 16 == 0); // 16 Byte Alignment
206
+
207
+ if (bias) {
208
+ STD_TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
209
+ bias->dim() == 1);
210
+ }
211
+
212
+ const torch::stable::accelerator::DeviceGuard device_guard(
213
+ a.get_device_index());
214
+ int32_t version_num = get_sm_version_num();
215
+
216
+#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
217
+ if (version_num >= 120) {
218
+ cutlass_scaled_mm_sm120(c, a, b, a_scales, b_scales, bias);
219
+ return;
220
+ }
221
+#endif
222
+
223
+#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
224
+ if (version_num >= 100 && version_num < 120) {
225
+ cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
226
+ return;
227
+ }
228
+#endif
229
+
230
+ // Guard against compilation issues for sm90 kernels
231
+#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
232
+ if (version_num >= 90 && version_num < 100) {
233
+ // Hopper
234
+ cutlass_scaled_mm_sm90(c, a, b, a_scales, b_scales, bias);
235
+ return;
236
+ }
237
+#endif
238
+
239
+#if defined ENABLE_SCALED_MM_C2X && ENABLE_SCALED_MM_C2X
240
+ if (version_num == 89) {
241
+ // Ada Lovelace
242
+ cutlass_scaled_mm_sm89(c, a, b, a_scales, b_scales, bias);
243
+ return;
244
+ }
245
+
246
+ if (version_num >= 80) {
247
+ // Ampere
248
+ cutlass_scaled_mm_sm80(c, a, b, a_scales, b_scales, bias);
249
+ return;
250
+ }
251
+
252
+ if (version_num >= 75) {
253
+ // Turing
254
+ cutlass_scaled_mm_sm75(c, a, b, a_scales, b_scales, bias);
255
+ return;
256
+ }
257
+#endif
258
+
259
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
260
+ false,
261
+ "No compiled cutlass_scaled_mm for a compute capability less than "
262
+ "CUDA device capability: ",
263
+ version_num);
264
+}
265
+
266
+void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
267
+ torch::stable::Tensor const& a_tensors,
268
+ torch::stable::Tensor const& b_tensors,
269
+ torch::stable::Tensor const& a_scales,
270
+ torch::stable::Tensor const& b_scales,
271
+ torch::stable::Tensor const& expert_offsets,
272
+ torch::stable::Tensor const& problem_sizes,
273
+ torch::stable::Tensor const& a_strides,
274
+ torch::stable::Tensor const& b_strides,
275
+ torch::stable::Tensor const& c_strides, bool per_act_token,
276
+ bool per_out_ch) {
277
+ int32_t version_num = get_sm_version_num();
278
+#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
279
+ if (version_num >= 100 && version_num < 110) {
280
+ cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
281
+ expert_offsets, problem_sizes, a_strides, b_strides,
282
+ c_strides, per_act_token, per_out_ch);
283
+ return;
284
+ }
285
+#endif
286
+#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
287
+ if (version_num >= 90 && version_num < 100) {
288
+ cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
289
+ expert_offsets, problem_sizes, a_strides, b_strides,
290
+ c_strides, per_act_token, per_out_ch);
291
+ return;
292
+ }
293
+#endif
294
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
295
+ false,
296
+ "No compiled cutlass_scaled_mm for CUDA device capability: ", version_num,
297
+ ". Required capability: 90 or 100");
298
+}
299
+
300
+void get_cutlass_moe_mm_data(
301
+ const torch::stable::Tensor& topk_ids,
302
+ torch::stable::Tensor& expert_offsets,
303
+ torch::stable::Tensor& problem_sizes1,
304
+ torch::stable::Tensor& problem_sizes2,
305
+ torch::stable::Tensor& input_permutation,
306
+ torch::stable::Tensor& output_permutation, const int64_t num_experts,
307
+ const int64_t n, const int64_t k,
308
+ const std::optional<torch::stable::Tensor>& blockscale_offsets,
309
+ const bool is_gated) {
310
+ // This function currently gets compiled only if we have a valid cutlass moe
311
+ // mm to run it for.
312
+ int32_t version_num = get_sm_version_num();
313
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
314
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
315
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
316
+ get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
317
+ problem_sizes2, input_permutation,
318
+ output_permutation, num_experts, n, k,
319
+ blockscale_offsets, is_gated);
320
+ return;
321
+#endif
322
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
323
+ false,
324
+ "No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
325
+ "CUDA device capability: ",
326
+ version_num, ". Required capability: 90, 100, or 120");
327
+}
328
+
329
+void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
330
+ const torch::stable::Tensor& expert_first_token_offset,
331
+ torch::stable::Tensor& problem_sizes1,
332
+ torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
333
+ const bool swap_ab) {
334
+ int32_t version_num = get_sm_version_num();
335
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
336
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
337
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
338
+ get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
339
+ expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
340
+ return;
341
+#endif
342
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
343
+ false,
344
+ "No compiled get_cutlass_moe_mm_problem_sizes_from_expert_offsets: "
345
+ "no cutlass_scaled_mm kernel for CUDA device capability: ",
346
+ version_num, ". Required capability: 90, 100, or 120");
347
+}
348
+
349
+void get_cutlass_batched_moe_mm_data(
350
+ torch::stable::Tensor& expert_offsets,
351
+ torch::stable::Tensor& problem_sizes1,
352
+ torch::stable::Tensor& problem_sizes2,
353
+ const torch::stable::Tensor& expert_num_tokens,
354
+ const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
355
+ const int64_t k) {
356
+ // This function currently gets compiled only if we have a valid cutlass moe
357
+ // mm to run it for.
358
+ int32_t version_num = get_sm_version_num();
359
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
360
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
361
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
362
+ get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
363
+ problem_sizes2, expert_num_tokens,
364
+ num_local_experts, padded_m, n, k);
365
+ return;
366
+#endif
367
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
368
+ false,
369
+ "No compiled get_cutlass_batched_moe_mm_data: no "
370
+ "cutlass_scaled_mm kernel "
371
+ "for CUDA device capability: ",
372
+ version_num, ". Required capability: 90, 100, or 120");
373
+}
374
+
375
+void cutlass_scaled_mm_azp(torch::stable::Tensor& c,
376
+ torch::stable::Tensor const& a,
377
+ torch::stable::Tensor const& b,
378
+ torch::stable::Tensor const& a_scales,
379
+ torch::stable::Tensor const& b_scales,
380
+ torch::stable::Tensor const& azp_adj,
381
+ std::optional<torch::stable::Tensor> const& azp,
382
+ std::optional<torch::stable::Tensor> const& bias) {
383
+ // Checks for conformality
384
+ STD_TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
385
+ STD_TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
386
+ b.size(1) == c.size(1));
387
+ STD_TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
388
+ STD_TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
389
+
390
+ // Check for strides and alignment
391
+ STD_TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
392
+ STD_TORCH_CHECK(b.stride(0) == 1); // Column-major
393
+ STD_TORCH_CHECK(c.stride(0) % 16 == 0 &&
394
+ b.stride(1) % 16 == 0); // 16 Byte Alignment
395
+ STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
396
+
397
+ // bias, azp, azp_adj are all 1d
398
+ // bias and azp_adj have n elements, azp has m elements
399
+ if (bias) {
400
+ STD_TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous());
401
+ }
402
+ if (azp) {
403
+ STD_TORCH_CHECK(azp->numel() == a.size(0) && azp->is_contiguous());
404
+ }
405
+ STD_TORCH_CHECK(azp_adj.numel() == b.size(1) && azp_adj.is_contiguous());
406
+
407
+ // azp & bias types
408
+ STD_TORCH_CHECK(azp_adj.scalar_type() == torch::headeronly::ScalarType::Int);
409
+ STD_TORCH_CHECK(!azp ||
410
+ azp->scalar_type() == torch::headeronly::ScalarType::Int);
411
+ STD_TORCH_CHECK(!bias || bias->scalar_type() == c.scalar_type(),
412
+ "currently bias dtype must match output dtype ",
413
+ c.scalar_type());
414
+
415
+ const torch::stable::accelerator::DeviceGuard device_guard(
416
+ a.get_device_index());
417
+
418
+ int32_t version_num = get_sm_version_num();
419
+
420
+#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
421
+ if (version_num >= 90) {
422
+ cutlass_scaled_mm_azp_sm90(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
423
+ return;
424
+ }
425
+#endif
426
+
427
+#if defined ENABLE_SCALED_MM_C2X && ENABLE_SCALED_MM_C2X
428
+ if (version_num == 89) {
429
+ // Ada Lovelace
430
+ cutlass_scaled_mm_azp_sm89(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
431
+ return;
432
+ }
433
+
434
+ if (version_num >= 80) {
435
+ // Ampere
436
+ cutlass_scaled_mm_azp_sm80(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
437
+ return;
438
+ }
439
+
440
+ // Turing
441
+ STD_TORCH_CHECK(version_num >= 75);
442
+ cutlass_scaled_mm_azp_sm75(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
443
+ return;
444
+#endif
445
+
446
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
447
+ false,
448
+ "No compiled cutlass_scaled_mm_azp for a compute capability less than "
449
+ "CUDA device capability: ",
450
+ version_num);
451
+}
csrc/libtorch_stable/torch_bindings.cpp
+97
@@ -31,6 +31,78 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
31
"per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
32
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
33
"()");
34
+
35
+ // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
36
+ // quantization, as well as bias
37
+ ops.def(
38
+ "cutlass_scaled_mm(Tensor! out, Tensor a,"
39
+ " Tensor b, Tensor a_scales,"
40
+ " Tensor b_scales, Tensor? bias) -> ()");
41
+
42
+ // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
43
+ // quantization.
44
+ ops.def(
45
+ "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
46
+ " Tensor b, Tensor a_scales,"
47
+ " Tensor b_scales, Tensor azp_adj,"
48
+ " Tensor? azp, Tensor? bias) -> ()");
49
+
50
+ // Check if cutlass scaled_mm is supported for CUDA devices of the given
51
+ // capability
52
+ ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool");
53
+
54
+ // Check if cutlass grouped gemm is supported for CUDA devices of the given
55
+ // capability
56
+ ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool");
57
+
58
+ // CUTLASS w8a8 grouped GEMM
59
+ ops.def(
60
+ "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, "
61
+ " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
62
+ " Tensor problem_sizes, Tensor a_strides, "
63
+ " Tensor b_strides, Tensor c_strides, bool per_act_token, "
64
+ " bool per_out_ch) -> ()");
65
+
66
+ // A function that computes data required to run fused MoE with w8a8 grouped
67
+ // GEMM. It takes topk_ids as an input, and computes expert_offsets
68
+ // (token start indices of each expert). In addition to this, it computes
69
+ // problem sizes for each expert's multiplication used by the two mms called
70
+ // from fused MoE operation, and arrays with permutations required to shuffle
71
+ // and de-shuffle the input/output of the fused operation.
72
+ ops.def(
73
+ "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, "
74
+ " Tensor! problem_sizes1, Tensor! problem_sizes2, "
75
+ " Tensor! input_permutation, "
76
+ " Tensor! output_permutation, int num_experts, "
77
+ " int n, int k, Tensor? blockscale_offsets, "
78
+ " bool is_gated) -> ()");
79
+
80
+ // compute per-expert problem sizes from expert_first_token_offset
81
+ // produced by vLLM's moe_permute kernel
82
+ ops.def(
83
+ "get_cutlass_moe_mm_problem_sizes_from_expert_offsets("
84
+ " Tensor expert_first_token_offset, "
85
+ " Tensor! problem_sizes1, "
86
+ " Tensor! problem_sizes2, "
87
+ " int n, int k, bool swap_ab) -> ()");
88
+
89
+ // A function that computes data required to run fused MoE with w8a8 grouped
90
+ // GEMM in batched expert format. It takes expert_num_tokens
91
+ // as an input, and computes expert_offsets (token start indices of each
92
+ // expert). In addition to this, it computes problem sizes for each expert's
93
+ // multiplication used by the two mms called from fused MoE operation.
94
+ ops.def(
95
+ "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, "
96
+ " Tensor! problem_sizes1, "
97
+ " Tensor! problem_sizes2, "
98
+ " Tensor expert_num_tokens, "
99
+ " int num_local_experts, int padded_m, "
100
+ " int n, int k) -> ()");
101
+
102
+ // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
103
+ ops.def(
104
+ "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
105
+ "bool");
106
#endif
107
}
108
@@ -46,6 +118,31 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
118
TORCH_BOX(&per_token_group_quant_8bit_packed));
119
ops.impl("per_token_group_quant_int8",
120
TORCH_BOX(&per_token_group_quant_int8));
121
+
122
+ // CUTLASS scaled_mm ops
123
+ ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));
124
+ ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp));
125
+ ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm));
126
+ ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data));
127
+ ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets",
128
+ TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets));
129
+ ops.impl("get_cutlass_batched_moe_mm_data",
130
+ TORCH_BOX(&get_cutlass_batched_moe_mm_data));
131
+#endif
132
+}
133
+
134
+// These capability-check functions take only primitive args (no tensors), so
135
+// there is no device to dispatch on. CompositeExplicitAutograd makes them
136
+// available for all backends. This is the stable ABI equivalent of calling
137
+// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
138
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
139
+#ifndef USE_ROCM
140
+ ops.impl("cutlass_scaled_mm_supports_fp8",
141
+ TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
142
+ ops.impl("cutlass_group_gemm_supported",
143
+ TORCH_BOX(&cutlass_group_gemm_supported));
144
+ ops.impl("cutlass_scaled_mm_supports_block_fp8",
145
+ TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8));
146
#endif
147
}
148
csrc/libtorch_stable/torch_utils.h
+6
@@ -1,10 +1,16 @@
1
#pragma once
2
3
#include <torch/csrc/inductor/aoti_torch/c/shim.h>
4
+#include <torch/csrc/stable/accelerator.h>
5
+#include <torch/csrc/stable/tensor.h>
6
#include <torch/headeronly/util/shim_utils.h>
7
8
#include <cuda_runtime.h>
9
10
+// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED.
11
+#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \
12
+ STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__)
13
+
14
// Utility to get the current CUDA stream for a given device using stable APIs.
15
// Returns a cudaStream_t for use in kernel launches.
16
inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
csrc/ops.h
-45
@@ -228,63 +228,18 @@ int64_t ggml_moe_get_block_size(int64_t type);
228
#ifndef USE_ROCM
229
230
bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
231
-bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
232
-bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
233
-bool cutlass_group_gemm_supported(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
240
-void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
241
- torch::Tensor const& b, torch::Tensor const& a_scales,
242
- torch::Tensor const& b_scales,
243
- std::optional<torch::Tensor> const& bias);
244
-
245
-void cutlass_moe_mm(
246
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
247
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
248
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
249
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
250
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
251
- bool per_act_token, bool per_out_ch);
252
-
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
259
-void get_cutlass_moe_mm_data(
260
- const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
261
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
262
- torch::Tensor& input_permutation, torch::Tensor& output_permutation,
263
- const int64_t num_experts, const int64_t n, const int64_t k,
264
- const std::optional<torch::Tensor>& blockscale_offsets,
265
- const bool is_gated);
266
-
267
-void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
268
- const torch::Tensor& expert_first_token_offset,
269
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
270
- const int64_t n, const int64_t k, const bool swap_ab);
271
-
272
-void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets,
273
- torch::Tensor& problem_sizes1,
274
- torch::Tensor& problem_sizes2,
275
- const torch::Tensor& expert_num_tokens,
276
- const int64_t num_local_experts,
277
- const int64_t padded_m, const int64_t n,
278
- const int64_t k);
279
-
280
-void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
281
- torch::Tensor const& b,
282
- torch::Tensor const& a_scales,
283
- torch::Tensor const& b_scales,
284
- torch::Tensor const& azp_adj,
285
- std::optional<torch::Tensor> const& azp,
286
- std::optional<torch::Tensor> const& bias);
287
-
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);
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu
deleted
-23
@@ -1,23 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh"
3
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
-
5
-namespace vllm {
6
-
7
-void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out,
8
- torch::Tensor const& a,
9
- torch::Tensor const& b,
10
- torch::Tensor const& a_scales,
11
- torch::Tensor const& b_scales) {
12
- if (out.dtype() == torch::kBFloat16) {
13
- cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::bfloat16_t>(
14
- out, a, b, a_scales, b_scales);
15
-
16
- } else {
17
- TORCH_CHECK(out.dtype() == torch::kFloat16);
18
- cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::half_t>(
19
- out, a, b, a_scales, b_scales);
20
- }
21
-}
22
-
23
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu
deleted
-23
@@ -1,23 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh"
3
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
-
5
-namespace vllm {
6
-
7
-void cutlass_scaled_mm_blockwise_sm120_fp8(torch::Tensor& out,
8
- torch::Tensor const& a,
9
- torch::Tensor const& b,
10
- torch::Tensor const& a_scales,
11
- torch::Tensor const& b_scales) {
12
- if (out.dtype() == torch::kBFloat16) {
13
- cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::bfloat16_t>(
14
- out, a, b, a_scales, b_scales);
15
-
16
- } else {
17
- TORCH_CHECK(out.dtype() == torch::kFloat16);
18
- cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::half_t>(
19
- out, a, b, a_scales, b_scales);
20
- }
21
-}
22
-
23
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu
deleted
-24
@@ -1,24 +0,0 @@
1
-
2
-#include "scaled_mm_kernels.hpp"
3
-#include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh"
4
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
5
-
6
-namespace vllm {
7
-
8
-void cutlass_scaled_mm_blockwise_sm90_fp8(torch::Tensor& out,
9
- torch::Tensor const& a,
10
- torch::Tensor const& b,
11
- torch::Tensor const& a_scales,
12
- torch::Tensor const& b_scales) {
13
- if (out.dtype() == torch::kBFloat16) {
14
- cutlass_gemm_blockwise_sm90_fp8_dispatch<cutlass::bfloat16_t>(
15
- out, a, b, a_scales, b_scales);
16
-
17
- } else {
18
- TORCH_CHECK(out.dtype() == torch::kFloat16);
19
- cutlass_gemm_blockwise_sm90_fp8_dispatch<cutlass::half_t>(
20
- out, a, b, a_scales, b_scales);
21
- }
22
-}
23
-
24
-} // namespace vllm
\ No newline at end of file
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp
deleted
-56
@@ -1,56 +0,0 @@
1
-#pragma once
2
-
3
-#include <torch/all.h>
4
-
5
-namespace vllm {
6
-
7
-void cutlass_scaled_mm_sm90_fp8(torch::Tensor& out, torch::Tensor const& a,
8
- torch::Tensor const& b,
9
- torch::Tensor const& a_scales,
10
- torch::Tensor const& b_scales,
11
- std::optional<torch::Tensor> const& bias);
12
-
13
-void cutlass_scaled_mm_sm90_int8(torch::Tensor& out, torch::Tensor const& a,
14
- torch::Tensor const& b,
15
- torch::Tensor const& a_scales,
16
- torch::Tensor const& b_scales,
17
- std::optional<torch::Tensor> const& bias);
18
-
19
-void cutlass_scaled_mm_azp_sm90_int8(torch::Tensor& out, torch::Tensor const& a,
20
- torch::Tensor const& b,
21
- torch::Tensor const& a_scales,
22
- torch::Tensor const& b_scales,
23
- torch::Tensor const& azp_adj,
24
- std::optional<torch::Tensor> const& azp,
25
- std::optional<torch::Tensor> const& bias);
26
-
27
-void cutlass_scaled_mm_blockwise_sm90_fp8(torch::Tensor& out,
28
- torch::Tensor const& a,
29
- torch::Tensor const& b,
30
- torch::Tensor const& a_scales,
31
- torch::Tensor const& b_scales);
32
-
33
-void cutlass_scaled_mm_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
34
- torch::Tensor const& b,
35
- torch::Tensor const& a_scales,
36
- torch::Tensor const& b_scales,
37
- std::optional<torch::Tensor> const& bias);
38
-
39
-void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
40
- torch::Tensor const& b,
41
- torch::Tensor const& a_scales,
42
- torch::Tensor const& b_scales,
43
- std::optional<torch::Tensor> const& bias);
44
-
45
-void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out,
46
- torch::Tensor const& a,
47
- torch::Tensor const& b,
48
- torch::Tensor const& a_scales,
49
- torch::Tensor const& b_scales);
50
-
51
-void cutlass_scaled_mm_blockwise_sm120_fp8(torch::Tensor& out,
52
- torch::Tensor const& a,
53
- torch::Tensor const& b,
54
- torch::Tensor const& a_scales,
55
- torch::Tensor const& b_scales);
56
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu
deleted
-23
@@ -1,23 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_sm100_fp8_dispatch.cuh"
3
-
4
-namespace vllm {
5
-
6
-void cutlass_scaled_mm_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
7
- torch::Tensor const& b,
8
- torch::Tensor const& a_scales,
9
- torch::Tensor const& b_scales,
10
- std::optional<torch::Tensor> const& bias) {
11
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
12
- if (bias) {
13
- TORCH_CHECK(bias->dtype() == out.dtype(),
14
- "currently bias dtype must match output dtype ", out.dtype());
15
- return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
16
- b_scales, *bias);
17
- } else {
18
- return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
19
- b_scales);
20
- }
21
-}
22
-
23
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu
deleted
-24
@@ -1,24 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_sm120_fp8_dispatch.cuh"
3
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
-
5
-namespace vllm {
6
-
7
-void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
8
- torch::Tensor const& b,
9
- torch::Tensor const& a_scales,
10
- torch::Tensor const& b_scales,
11
- std::optional<torch::Tensor> const& bias) {
12
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
13
- if (bias) {
14
- TORCH_CHECK(bias->dtype() == out.dtype(),
15
- "currently bias dtype must match output dtype ", out.dtype());
16
- return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
17
- out, a, b, a_scales, b_scales, *bias);
18
- } else {
19
- return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
20
- out, a, b, a_scales, b_scales);
21
- }
22
-}
23
-
24
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu
deleted
-23
@@ -1,23 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_sm90_fp8_dispatch.cuh"
3
-
4
-namespace vllm {
5
-
6
-void cutlass_scaled_mm_sm90_fp8(torch::Tensor& out, torch::Tensor const& a,
7
- torch::Tensor const& b,
8
- torch::Tensor const& a_scales,
9
- torch::Tensor const& b_scales,
10
- std::optional<torch::Tensor> const& bias) {
11
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
12
- if (bias) {
13
- TORCH_CHECK(bias->dtype() == out.dtype(),
14
- "currently bias dtype must match output dtype ", out.dtype());
15
- return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
16
- b_scales, *bias);
17
- } else {
18
- return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
19
- b_scales);
20
- }
21
-}
22
-
23
-} // namespace vllm
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu
deleted
-24
@@ -1,24 +0,0 @@
1
-#include "scaled_mm_kernels.hpp"
2
-#include "scaled_mm_sm90_int8_dispatch.cuh"
3
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
-
5
-namespace vllm {
6
-
7
-void cutlass_scaled_mm_sm90_int8(torch::Tensor& out, torch::Tensor const& a,
8
- torch::Tensor const& b,
9
- torch::Tensor const& a_scales,
10
- torch::Tensor const& b_scales,
11
- std::optional<torch::Tensor> const& bias) {
12
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
13
- if (bias) {
14
- TORCH_CHECK(bias->dtype() == out.dtype(),
15
- "currently bias dtype must match output dtype ", out.dtype());
16
- return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBias>(
17
- out, a, b, a_scales, b_scales, *bias);
18
- } else {
19
- return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogue>(
20
- out, a, b, a_scales, b_scales);
21
- }
22
-}
23
-
24
-} // namespace vllm
csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu
deleted
-199
@@ -1,199 +0,0 @@
1
-#include <stddef.h>
2
-#include <torch/all.h>
3
-#include "cutlass/cutlass.h"
4
-
5
-#include "scaled_mm_c2x.cuh"
6
-#include "scaled_mm_c2x_sm75_dispatch.cuh"
7
-#include "scaled_mm_c2x_sm80_dispatch.cuh"
8
-#include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
9
-#include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
10
-
11
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
12
-
13
-using namespace vllm;
14
-
15
-/*
16
- This file defines quantized GEMM operations using the CUTLASS 2.x API, for
17
- NVIDIA GPUs with SM versions prior to sm90 (Hopper).
18
-*/
19
-
20
-template <template <typename, typename> typename Epilogue,
21
- typename... EpilogueArgs>
22
-void cutlass_scaled_mm_sm75_epilogue(torch::Tensor& out, torch::Tensor const& a,
23
- torch::Tensor const& b,
24
- EpilogueArgs&&... epilogue_args) {
25
- TORCH_CHECK(a.dtype() == torch::kInt8);
26
- TORCH_CHECK(b.dtype() == torch::kInt8);
27
-
28
- if (out.dtype() == torch::kBFloat16) {
29
- return cutlass_gemm_sm75_dispatch<int8_t, cutlass::bfloat16_t, Epilogue>(
30
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
31
- } else {
32
- TORCH_CHECK(out.dtype() == torch::kFloat16);
33
- return cutlass_gemm_sm75_dispatch<int8_t, cutlass::half_t, Epilogue>(
34
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
35
- }
36
-}
37
-
38
-void cutlass_scaled_mm_sm75(torch::Tensor& out, torch::Tensor const& a,
39
- torch::Tensor const& b,
40
- torch::Tensor const& a_scales,
41
- torch::Tensor const& b_scales,
42
- std::optional<torch::Tensor> const& bias) {
43
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
44
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
45
- if (bias) {
46
- TORCH_CHECK(bias->dtype() == out.dtype(),
47
- "currently bias dtype must match output dtype ", out.dtype());
48
- return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBias>(
49
- out, a, b, a_scales, b_scales, *bias);
50
- } else {
51
- return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogue>(
52
- out, a, b, a_scales, b_scales);
53
- }
54
-}
55
-
56
-void cutlass_scaled_mm_azp_sm75(torch::Tensor& out, torch::Tensor const& a,
57
- torch::Tensor const& b,
58
- torch::Tensor const& a_scales,
59
- torch::Tensor const& b_scales,
60
- torch::Tensor const& azp_adj,
61
- std::optional<torch::Tensor> const& azp,
62
- std::optional<torch::Tensor> const& bias) {
63
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
64
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
65
-
66
- if (azp) {
67
- return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
68
- out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
69
- } else {
70
- return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBiasAzp>(
71
- out, a, b, a_scales, b_scales, azp_adj, bias);
72
- }
73
-}
74
-
75
-template <template <typename, typename> typename Epilogue,
76
- typename... EpilogueArgs>
77
-void cutlass_scaled_mm_sm80_epilogue(torch::Tensor& out, torch::Tensor const& a,
78
- torch::Tensor const& b,
79
- EpilogueArgs&&... epilogue_args) {
80
- TORCH_CHECK(a.dtype() == torch::kInt8);
81
- TORCH_CHECK(b.dtype() == torch::kInt8);
82
-
83
- if (out.dtype() == torch::kBFloat16) {
84
- return cutlass_gemm_sm80_dispatch<int8_t, cutlass::bfloat16_t, Epilogue>(
85
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
86
- } else {
87
- TORCH_CHECK(out.dtype() == torch::kFloat16);
88
- return cutlass_gemm_sm80_dispatch<int8_t, cutlass::half_t, Epilogue>(
89
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
90
- }
91
-}
92
-
93
-void cutlass_scaled_mm_sm80(torch::Tensor& out, torch::Tensor const& a,
94
- torch::Tensor const& b,
95
- torch::Tensor const& a_scales,
96
- torch::Tensor const& b_scales,
97
- std::optional<torch::Tensor> const& bias) {
98
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
99
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
100
- if (bias) {
101
- TORCH_CHECK(bias->dtype() == out.dtype(),
102
- "currently bias dtype must match output dtype ", out.dtype());
103
- return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBias>(
104
- out, a, b, a_scales, b_scales, *bias);
105
- } else {
106
- return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogue>(
107
- out, a, b, a_scales, b_scales);
108
- }
109
-}
110
-
111
-void cutlass_scaled_mm_azp_sm80(torch::Tensor& out, torch::Tensor const& a,
112
- torch::Tensor const& b,
113
- torch::Tensor const& a_scales,
114
- torch::Tensor const& b_scales,
115
- torch::Tensor const& azp_adj,
116
- std::optional<torch::Tensor> const& azp,
117
- std::optional<torch::Tensor> const& bias) {
118
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
119
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
120
-
121
- if (azp) {
122
- return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
123
- out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
124
- } else {
125
- return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBiasAzp>(
126
- out, a, b, a_scales, b_scales, azp_adj, bias);
127
- }
128
-}
129
-
130
-template <template <typename, typename> typename Epilogue,
131
- typename... EpilogueArgs>
132
-void cutlass_scaled_mm_sm89_epilogue(torch::Tensor& out, torch::Tensor const& a,
133
- torch::Tensor const& b,
134
- EpilogueArgs&&... epilogue_args) {
135
- if (a.dtype() == torch::kInt8) {
136
- TORCH_CHECK(b.dtype() == torch::kInt8);
137
-
138
- if (out.dtype() == torch::kBFloat16) {
139
- return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::bfloat16_t,
140
- Epilogue>(
141
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
142
- } else {
143
- assert(out.dtype() == torch::kFloat16);
144
- return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::half_t, Epilogue>(
145
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
146
- }
147
- } else {
148
- TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
149
- TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
150
-
151
- if (out.dtype() == torch::kBFloat16) {
152
- return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
153
- cutlass::bfloat16_t, Epilogue>(
154
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
155
- } else {
156
- TORCH_CHECK(out.dtype() == torch::kFloat16);
157
- return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
158
- cutlass::half_t, Epilogue>(
159
- out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
160
- }
161
- }
162
-}
163
-
164
-void cutlass_scaled_mm_sm89(torch::Tensor& out, torch::Tensor const& a,
165
- torch::Tensor const& b,
166
- torch::Tensor const& a_scales,
167
- torch::Tensor const& b_scales,
168
- std::optional<torch::Tensor> const& bias) {
169
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
170
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
171
- if (bias) {
172
- TORCH_CHECK(bias->dtype() == out.dtype(),
173
- "currently bias dtype must match output dtype ", out.dtype());
174
- return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
175
- out, a, b, a_scales, b_scales, *bias);
176
- } else {
177
- return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
178
- out, a, b, a_scales, b_scales);
179
- }
180
-}
181
-
182
-void cutlass_scaled_mm_azp_sm89(torch::Tensor& out, torch::Tensor const& a,
183
- torch::Tensor const& b,
184
- torch::Tensor const& a_scales,
185
- torch::Tensor const& b_scales,
186
- torch::Tensor const& azp_adj,
187
- std::optional<torch::Tensor> const& azp,
188
- std::optional<torch::Tensor> const& bias) {
189
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
190
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
191
-
192
- if (azp) {
193
- return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBiasAzpToken>(
194
- out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
195
- } else {
196
- return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBiasAzp>(
197
- out, a, b, a_scales, b_scales, azp_adj, bias);
198
- }
199
-}
csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu
deleted
-36
@@ -1,36 +0,0 @@
1
-#include "c3x/scaled_mm_helper.hpp"
2
-#include "c3x/scaled_mm_kernels.hpp"
3
-
4
-/*
5
- This file defines quantized GEMM operations using the CUTLASS 3.x API, for
6
- NVIDIA GPUs with sm90a (Hopper).
7
-*/
8
-
9
-#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
10
-
11
-void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
12
- torch::Tensor const& b,
13
- torch::Tensor const& a_scales,
14
- torch::Tensor const& b_scales,
15
- std::optional<torch::Tensor> const& bias) {
16
- dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
17
- vllm::cutlass_scaled_mm_sm90_fp8,
18
- vllm::cutlass_scaled_mm_sm90_int8,
19
- vllm::cutlass_scaled_mm_blockwise_sm90_fp8);
20
-}
21
-
22
-void cutlass_scaled_mm_azp_sm90(torch::Tensor& out, torch::Tensor const& a,
23
- torch::Tensor const& b,
24
- torch::Tensor const& a_scales,
25
- torch::Tensor const& b_scales,
26
- torch::Tensor const& azp_adj,
27
- std::optional<torch::Tensor> const& azp,
28
- std::optional<torch::Tensor> const& bias) {
29
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
30
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
31
-
32
- vllm::cutlass_scaled_mm_azp_sm90_int8(out, a, b, a_scales, b_scales, azp_adj,
33
- azp, bias);
34
-}
35
-
36
-#endif
csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu
deleted
-420
@@ -1,420 +0,0 @@
1
-#include <cudaTypedefs.h>
2
-
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <torch/all.h>
5
-
6
-#include "cutlass_extensions/common.hpp"
7
-
8
-void cutlass_scaled_mm_sm75(torch::Tensor& c, torch::Tensor const& a,
9
- torch::Tensor const& b,
10
- torch::Tensor const& a_scales,
11
- torch::Tensor const& b_scales,
12
- std::optional<torch::Tensor> const& bias);
13
-
14
-void cutlass_scaled_mm_sm80(torch::Tensor& c, torch::Tensor const& a,
15
- torch::Tensor const& b,
16
- torch::Tensor const& a_scales,
17
- torch::Tensor const& b_scales,
18
- std::optional<torch::Tensor> const& bias);
19
-
20
-void cutlass_scaled_mm_sm89(torch::Tensor& c, torch::Tensor const& a,
21
- torch::Tensor const& b,
22
- torch::Tensor const& a_scales,
23
- torch::Tensor const& b_scales,
24
- std::optional<torch::Tensor> const& bias);
25
-
26
-#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
27
-void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
28
- torch::Tensor const& b,
29
- torch::Tensor const& a_scales,
30
- torch::Tensor const& b_scales,
31
- std::optional<torch::Tensor> const& bias);
32
-#endif
33
-#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
34
-void cutlass_moe_mm_sm90(
35
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
36
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
37
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
38
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
39
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
40
- bool per_act_token, bool per_out_ch);
41
-
42
-#endif
43
-
44
-#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
45
-void cutlass_moe_mm_sm100(
46
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
47
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
48
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
49
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
50
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
51
- bool per_act_token, bool per_out_ch);
52
-#endif
53
-
54
-#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
55
-void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
56
- torch::Tensor const& b,
57
- torch::Tensor const& a_scales,
58
- torch::Tensor const& b_scales,
59
- std::optional<torch::Tensor> const& bias);
60
-#endif
61
-
62
-#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
63
-void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
64
- torch::Tensor const& b,
65
- torch::Tensor const& a_scales,
66
- torch::Tensor const& b_scales,
67
- std::optional<torch::Tensor> const& bias);
68
-#endif
69
-
70
-#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
71
- (defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
72
- (defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
73
-void get_cutlass_moe_mm_data_caller(
74
- const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
75
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
76
- torch::Tensor& input_permutation, torch::Tensor& output_permutation,
77
- const int64_t num_experts, const int64_t n, const int64_t k,
78
- const std::optional<torch::Tensor>& blockscale_offsets,
79
- const bool is_gated);
80
-
81
-void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
82
- const torch::Tensor& expert_first_token_offset,
83
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
84
- const int64_t n, const int64_t k, const bool swap_ab);
85
-
86
-void get_cutlass_batched_moe_mm_data_caller(
87
- torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1,
88
- torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens,
89
- const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
90
- const int64_t k);
91
-#endif
92
-
93
-void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a,
94
- torch::Tensor const& b,
95
- torch::Tensor const& a_scales,
96
- torch::Tensor const& b_scales,
97
- torch::Tensor const& azp_adj,
98
- std::optional<torch::Tensor> const& azp,
99
- std::optional<torch::Tensor> const& bias);
100
-
101
-void cutlass_scaled_mm_azp_sm80(torch::Tensor& c, torch::Tensor const& a,
102
- torch::Tensor const& b,
103
- torch::Tensor const& a_scales,
104
- torch::Tensor const& b_scales,
105
- torch::Tensor const& azp_adj,
106
- std::optional<torch::Tensor> const& azp,
107
- std::optional<torch::Tensor> const& bias);
108
-
109
-void cutlass_scaled_mm_azp_sm89(torch::Tensor& c, torch::Tensor const& a,
110
- torch::Tensor const& b,
111
- torch::Tensor const& a_scales,
112
- torch::Tensor const& b_scales,
113
- torch::Tensor const& azp_adj,
114
- std::optional<torch::Tensor> const& azp,
115
- std::optional<torch::Tensor> const& bias);
116
-
117
-#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
118
-void cutlass_scaled_mm_azp_sm90(torch::Tensor& c, torch::Tensor const& a,
119
- torch::Tensor const& b,
120
- torch::Tensor const& a_scales,
121
- torch::Tensor const& b_scales,
122
- torch::Tensor const& azp_adj,
123
- std::optional<torch::Tensor> const& azp,
124
- std::optional<torch::Tensor> const& bias);
125
-#endif
126
-
127
-bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability) {
128
- // CUTLASS FP8 kernels need at least
129
- // CUDA 12.0 on SM90 systems (Hopper)
130
- // CUDA 12.4 on SM89 systems (Lovelace)
131
-
132
-#if defined CUDA_VERSION
133
- if (cuda_device_capability >= 90) {
134
- return CUDA_VERSION >= 12000;
135
- } else if (cuda_device_capability >= 89) {
136
- return CUDA_VERSION >= 12040;
137
- }
138
-#endif
139
-
140
- return false;
141
-}
142
-
143
-bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
144
- // CUTLASS block-quantized FP8 kernels need at least CUDA 12.0
145
- // and at least SM90 (Hopper)
146
-
147
-#if defined CUDA_VERSION
148
- if (cuda_device_capability >= 100) {
149
- return CUDA_VERSION >= 12080;
150
- } else if (cuda_device_capability >= 90) {
151
- return CUDA_VERSION >= 12000;
152
- }
153
-#endif
154
-
155
- return false;
156
-}
157
-
158
-bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
159
- // CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
160
- // or CUDA 12.8 and SM100 (Blackwell)
161
-
162
-#if defined CUDA_VERSION
163
- if (cuda_device_capability >= 100) {
164
- return CUDA_VERSION >= 12080;
165
- }
166
- if (cuda_device_capability >= 90) {
167
- return CUDA_VERSION >= 12030;
168
- }
169
-#endif
170
-
171
- return false;
172
-}
173
-
174
-void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
175
- torch::Tensor const& b, torch::Tensor const& a_scales,
176
- torch::Tensor const& b_scales,
177
- std::optional<torch::Tensor> const& bias) {
178
- // Checks for conformality
179
- TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
180
- TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
181
- b.size(1) == c.size(1));
182
-
183
- // Check for strides and alignment
184
- TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
185
- TORCH_CHECK(b.stride(0) == 1); // Column-major
186
- TORCH_CHECK(c.stride(0) % 16 == 0 &&
187
- b.stride(1) % 16 == 0); // 16 Byte Alignment
188
-
189
- if (bias) {
190
- TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
191
- bias->dim() == 1);
192
- }
193
-
194
- at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
195
- int32_t version_num = get_sm_version_num();
196
-
197
-#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
198
- if (version_num >= 120) {
199
- cutlass_scaled_mm_sm120(c, a, b, a_scales, b_scales, bias);
200
- return;
201
- }
202
-#endif
203
-
204
-#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
205
- if (version_num >= 100 && version_num < 120) {
206
- cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
207
- return;
208
- }
209
-#endif
210
-
211
- // Guard against compilation issues for sm90 kernels
212
-#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
213
- if (version_num >= 90 && version_num < 100) {
214
- // Hopper
215
- cutlass_scaled_mm_sm90(c, a, b, a_scales, b_scales, bias);
216
- return;
217
- }
218
-#endif
219
-
220
-#if defined ENABLE_SCALED_MM_C2X && ENABLE_SCALED_MM_C2X
221
- if (version_num == 89) {
222
- // Ada Lovelace
223
- cutlass_scaled_mm_sm89(c, a, b, a_scales, b_scales, bias);
224
- return;
225
- }
226
-
227
- if (version_num >= 80) {
228
- // Ampere
229
- cutlass_scaled_mm_sm80(c, a, b, a_scales, b_scales, bias);
230
- return;
231
- }
232
-
233
- if (version_num >= 75) {
234
- // Turing
235
- cutlass_scaled_mm_sm75(c, a, b, a_scales, b_scales, bias);
236
- return;
237
- }
238
-#endif
239
-
240
- TORCH_CHECK_NOT_IMPLEMENTED(
241
- false,
242
- "No compiled cutlass_scaled_mm for a compute capability less than "
243
- "CUDA device capability: ",
244
- version_num);
245
-}
246
-
247
-void cutlass_moe_mm(
248
- torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
249
- torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
250
- torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
251
- torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
252
- torch::Tensor const& b_strides, torch::Tensor const& c_strides,
253
- bool per_act_token, bool per_out_ch) {
254
- int32_t version_num = get_sm_version_num();
255
-#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
256
- if (version_num >= 100 && version_num < 110) {
257
- cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
258
- expert_offsets, problem_sizes, a_strides, b_strides,
259
- c_strides, per_act_token, per_out_ch);
260
- return;
261
- }
262
-#endif
263
-#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
264
- if (version_num >= 90 && version_num < 100) {
265
- cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
266
- expert_offsets, problem_sizes, a_strides, b_strides,
267
- c_strides, per_act_token, per_out_ch);
268
- return;
269
- }
270
-#endif
271
- TORCH_CHECK_NOT_IMPLEMENTED(
272
- false,
273
- "No compiled cutlass_scaled_mm for CUDA device capability: ", version_num,
274
- ". Required capability: 90 or 100");
275
-}
276
-
277
-void get_cutlass_moe_mm_data(
278
- const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
279
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
280
- torch::Tensor& input_permutation, torch::Tensor& output_permutation,
281
- const int64_t num_experts, const int64_t n, const int64_t k,
282
- const std::optional<torch::Tensor>& blockscale_offsets,
283
- const bool is_gated) {
284
- // This function currently gets compiled only if we have a valid cutlass moe
285
- // mm to run it for.
286
- int32_t version_num = get_sm_version_num();
287
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
288
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
289
- (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
290
- get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
291
- problem_sizes2, input_permutation,
292
- output_permutation, num_experts, n, k,
293
- blockscale_offsets, is_gated);
294
- return;
295
-#endif
296
- TORCH_CHECK_NOT_IMPLEMENTED(
297
- false,
298
- "No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
299
- "CUDA device capability: ",
300
- version_num, ". Required capability: 90, 100, or 120");
301
-}
302
-
303
-void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
304
- const torch::Tensor& expert_first_token_offset,
305
- torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
306
- const int64_t n, const int64_t k, const bool swap_ab) {
307
- int32_t version_num = get_sm_version_num();
308
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
309
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
310
- (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
311
- get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
312
- expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
313
- return;
314
-#endif
315
- TORCH_CHECK_NOT_IMPLEMENTED(
316
- false,
317
- "No compiled get_cutlass_moe_mm_problem_sizes_from_expert_offsets: "
318
- "no cutlass_scaled_mm kernel for CUDA device capability: ",
319
- version_num, ". Required capability: 90, 100, or 120");
320
-}
321
-
322
-void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets,
323
- torch::Tensor& problem_sizes1,
324
- torch::Tensor& problem_sizes2,
325
- const torch::Tensor& expert_num_tokens,
326
- const int64_t num_local_experts,
327
- const int64_t padded_m, const int64_t n,
328
- const int64_t k) {
329
- // This function currently gets compiled only if we have a valid cutlass moe
330
- // mm to run it for.
331
- int32_t version_num = get_sm_version_num();
332
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
333
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
334
- (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
335
- get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
336
- problem_sizes2, expert_num_tokens,
337
- num_local_experts, padded_m, n, k);
338
- return;
339
-#endif
340
- TORCH_CHECK_NOT_IMPLEMENTED(false,
341
- "No compiled get_cutlass_batched_moe_mm_data: no "
342
- "cutlass_scaled_mm kernel "
343
- "for CUDA device capability: ",
344
- version_num,
345
- ". Required capability: 90, 100, or 120");
346
-}
347
-
348
-void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
349
- torch::Tensor const& b,
350
- torch::Tensor const& a_scales,
351
- torch::Tensor const& b_scales,
352
- torch::Tensor const& azp_adj,
353
- std::optional<torch::Tensor> const& azp,
354
- std::optional<torch::Tensor> const& bias) {
355
- // Checks for conformality
356
- TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
357
- TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
358
- b.size(1) == c.size(1));
359
- TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
360
- TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
361
-
362
- // Check for strides and alignment
363
- TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
364
- TORCH_CHECK(b.stride(0) == 1); // Column-major
365
- TORCH_CHECK(c.stride(0) % 16 == 0 &&
366
- b.stride(1) % 16 == 0); // 16 Byte Alignment
367
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
368
-
369
- // bias, azp, azp_adj are all 1d
370
- // bias and azp_adj have n elements, azp has m elements
371
- if (bias) {
372
- TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous());
373
- }
374
- if (azp) {
375
- TORCH_CHECK(azp->numel() == a.size(0) && azp->is_contiguous());
376
- }
377
- TORCH_CHECK(azp_adj.numel() == b.size(1) && azp_adj.is_contiguous());
378
-
379
- // azp & bias types
380
- TORCH_CHECK(azp_adj.dtype() == torch::kInt32);
381
- TORCH_CHECK(!azp || azp->dtype() == torch::kInt32);
382
- TORCH_CHECK(!bias || bias->dtype() == c.dtype(),
383
- "currently bias dtype must match output dtype ", c.dtype());
384
-
385
- at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
386
-
387
- int32_t version_num = get_sm_version_num();
388
-
389
-#if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
390
- if (version_num >= 90) {
391
- cutlass_scaled_mm_azp_sm90(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
392
- return;
393
- }
394
-#endif
395
-
396
-#if defined ENABLE_SCALED_MM_C2X && ENABLE_SCALED_MM_C2X
397
- if (version_num == 89) {
398
- // Ada Lovelace
399
- cutlass_scaled_mm_azp_sm89(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
400
- return;
401
- }
402
-
403
- if (version_num >= 80) {
404
- // Ampere
405
- cutlass_scaled_mm_azp_sm80(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
406
- return;
407
- }
408
-
409
- // Turing
410
- TORCH_CHECK(version_num >= 75);
411
- cutlass_scaled_mm_azp_sm75(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
412
- return;
413
-#endif
414
-
415
- TORCH_CHECK_NOT_IMPLEMENTED(
416
- false,
417
- "No compiled cutlass_scaled_mm_azp for a compute capability less than "
418
- "CUDA device capability: ",
419
- version_num);
420
-}
csrc/torch_bindings.cpp
-84
@@ -439,90 +439,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
439
" -> ()");
440
// conditionally compiled so impl registration is in source file
441
442
- // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
443
- // quantization, as well as bias
444
- ops.def(
445
- "cutlass_scaled_mm(Tensor! out, Tensor a,"
446
- " Tensor b, Tensor a_scales,"
447
- " Tensor b_scales, Tensor? bias) -> ()");
448
- ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);
449
-
450
- // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
451
- // quantization.
452
- ops.def(
453
- "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
454
- " Tensor b, Tensor a_scales,"
455
- " Tensor b_scales, Tensor azp_adj,"
456
- " Tensor? azp, Tensor? bias) -> ()");
457
- ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);
458
-
459
- // Check if cutlass scaled_mm is supported for CUDA devices of the given
460
- // capability
461
- ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool");
462
- ops.impl("cutlass_scaled_mm_supports_fp8", &cutlass_scaled_mm_supports_fp8);
463
-
464
- // Check if cutlass grouped gemm is supported for CUDA devices of the given
465
- // capability
466
- ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool");
467
- ops.impl("cutlass_group_gemm_supported", &cutlass_group_gemm_supported);
468
-
469
- // CUTLASS w8a8 grouped GEMM
470
- ops.def(
471
- "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, "
472
- " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
473
- " Tensor problem_sizes, Tensor a_strides, "
474
- " Tensor b_strides, Tensor c_strides, bool per_act_token, "
475
- " bool per_out_ch) -> ()");
476
- ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);
477
-
478
- // A function that computes data required to run fused MoE with w8a8 grouped
479
- // GEMM. It takes topk_ids as an input, and computes expert_offsets
480
- // (token start indices of each expert). In addition to this, it computes
481
- // problem sizes for each expert's multiplication used by the two mms called
482
- // from fused MoE operation, and arrays with permutations required to shuffle
483
- // and de-shuffle the input/output of the fused operation.
484
- ops.def(
485
- "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, "
486
- " Tensor! problem_sizes1, Tensor! problem_sizes2, "
487
- " Tensor! input_permutation, "
488
- " Tensor! output_permutation, int num_experts, "
489
- " int n, int k, Tensor? blockscale_offsets, "
490
- " bool is_gated) -> ()");
491
- ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
492
-
493
- // compute per-expert problem sizes from expert_first_token_offset
494
- // produced by vLLM's moe_permute kernel
495
- ops.def(
496
- "get_cutlass_moe_mm_problem_sizes_from_expert_offsets("
497
- " Tensor expert_first_token_offset, "
498
- " Tensor! problem_sizes1, "
499
- " Tensor! problem_sizes2, "
500
- " int n, int k, bool swap_ab) -> ()");
501
- ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", torch::kCUDA,
502
- &get_cutlass_moe_mm_problem_sizes_from_expert_offsets);
503
-
504
- // A function that computes data required to run fused MoE with w8a8 grouped
505
- // GEMM in batched expert format. It takes expert_num_tokens
506
- // as an input, and computes expert_offsets (token start indices of each
507
- // expert). In addition to this, it computes problem sizes for each expert's
508
- // multiplication used by the two mms called from fused MoE operation.
509
- ops.def(
510
- "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, "
511
- " Tensor! problem_sizes1, "
512
- " Tensor! problem_sizes2, "
513
- " Tensor expert_num_tokens, "
514
- " int num_local_experts, int padded_m, "
515
- " int n, int k) -> ()");
516
- ops.impl("get_cutlass_batched_moe_mm_data", torch::kCUDA,
517
- &get_cutlass_batched_moe_mm_data);
518
-
519
- // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
520
- ops.def(
521
- "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
522
- "bool");
523
- ops.impl("cutlass_scaled_mm_supports_block_fp8",
524
- &cutlass_scaled_mm_supports_block_fp8);
525
-
442
// SM100 CUTLASS MLA decode
443
ops.def(
444
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"