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 3, 19:00 UTC
[5/n] Migrate CUTLASS MLA, hadamard, awq, allspark and DSV3 fused a gemm to torch stable ABI (continued) (#42339)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
Chris Leonard committed
May 13, 2026 at 03:24 UTC
85b2fecab7776cb7526856803c661c82d1fa8433
20 files changed
+530
-412
CMakeLists.txt
+71
-70
@@ -369,7 +369,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
369
FetchContent_MakeAvailable(cutlass)
370
371
list(APPEND VLLM_EXT_SRC
372
- "csrc/quantization/awq/gemm_kernels.cu"
372
"csrc/cutlass_extensions/common.cpp")
373
374
set_gencode_flags_for_srcs(
@@ -501,46 +500,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
500
" in CUDA target architectures")
501
endif()
502
504
- # Only build AllSpark kernels if we are building for at least some compatible archs.
505
- cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
506
- if (ALLSPARK_ARCHS)
507
- set(ALLSPARK_SRCS
508
- "csrc/quantization/gptq_allspark/allspark_repack.cu"
509
- "csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
510
- set_gencode_flags_for_srcs(
511
- SRCS "${ALLSPARK_SRCS}"
512
- CUDA_ARCHS "${ALLSPARK_ARCHS}")
513
- list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
514
- message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
515
- else()
516
- message(STATUS "Not building AllSpark kernels as no compatible archs found"
517
- " in CUDA target architectures")
518
- endif()
519
-
520
- # CUTLASS MLA Archs and flags
521
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
522
- cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
523
- else()
524
- cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
525
- endif()
526
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
527
- set(SRCS
528
- "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
529
- set_gencode_flags_for_srcs(
530
- SRCS "${SRCS}"
531
- CUDA_ARCHS "${MLA_ARCHS}")
532
- list(APPEND VLLM_EXT_SRC "${SRCS}")
533
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
534
- # Add MLA-specific include directories only to MLA source files
535
- set_source_files_properties(${SRCS}
536
- PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
537
- message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
538
- else()
539
- message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
540
- # clear MLA_ARCHS
541
- set(MLA_ARCHS)
542
- endif()
543
-
503
# Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
504
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
505
cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -568,24 +527,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
527
endif()
528
endif()
529
571
- # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
572
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
573
- cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
574
- else()
575
- cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
576
- endif()
577
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
578
- set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
579
- set_gencode_flags_for_srcs(
580
- SRCS "${DSV3_FUSED_A_GEMM_SRC}"
581
- CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
582
- list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
583
- message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
584
- else()
585
- message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
586
- "in CUDA target architectures.")
587
- endif()
588
-
530
#
531
# Machete kernels
532
@@ -657,16 +598,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
598
endif()
599
600
660
- # Hadacore kernels
661
- cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
662
- if(HADACORE_ARCHS)
663
- set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
664
- set_gencode_flags_for_srcs(
665
- SRCS "${SRCS}"
666
- CUDA_ARCHS "${HADACORE_ARCHS}")
667
- list(APPEND VLLM_EXT_SRC "${SRCS}")
668
- message(STATUS "Building hadacore")
669
- endif()
601
602
# if CUDA endif
603
endif()
@@ -716,7 +647,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
647
list(APPEND VLLM_STABLE_EXT_SRC
648
"csrc/libtorch_stable/permute_cols.cu"
649
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
719
- "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
650
+ "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
651
+ "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
652
endif()
653
654
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -725,6 +657,40 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
657
CUDA_ARCHS "${CUDA_ARCHS}")
658
endif()
659
660
+ # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
661
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
662
+ cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
663
+ else()
664
+ cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
665
+ endif()
666
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
667
+ set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu")
668
+ set_gencode_flags_for_srcs(
669
+ SRCS "${SRCS}"
670
+ CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
671
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
672
+ message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
673
+ else()
674
+ message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
675
+ "in CUDA target architectures.")
676
+ endif()
677
+
678
+ # Only build AllSpark kernels if we are building for at least some compatible archs.
679
+ cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
680
+ if (ALLSPARK_ARCHS)
681
+ set(SRCS
682
+ "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
683
+ "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
684
+ set_gencode_flags_for_srcs(
685
+ SRCS "${SRCS}"
686
+ CUDA_ARCHS "${ALLSPARK_ARCHS}")
687
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
688
+ message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
689
+ else()
690
+ message(STATUS "Not building AllSpark kernels as no compatible archs found"
691
+ " in CUDA target architectures")
692
+ endif()
693
+
694
#
695
# CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
696
#
@@ -1034,6 +1000,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1000
endif()
1001
endif()
1002
1003
+ # CUTLASS MLA Archs and flags
1004
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1005
+ cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
1006
+ else()
1007
+ cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
1008
+ endif()
1009
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
1010
+ set(SRCS
1011
+ "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu")
1012
+ set_gencode_flags_for_srcs(
1013
+ SRCS "${SRCS}"
1014
+ CUDA_ARCHS "${MLA_ARCHS}")
1015
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
1016
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
1017
+ # Add MLA-specific include directories only to MLA source files
1018
+ set_source_files_properties(${SRCS}
1019
+ PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
1020
+ message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
1021
+ else()
1022
+ message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
1023
+ # clear MLA_ARCHS
1024
+ set(MLA_ARCHS)
1025
+ endif()
1026
+
1027
+ # Hadacore kernels
1028
+ cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
1029
+ if(HADACORE_ARCHS)
1030
+ set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
1031
+ set_gencode_flags_for_srcs(
1032
+ SRCS "${SRCS}"
1033
+ CUDA_ARCHS "${HADACORE_ARCHS}")
1034
+ list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
1035
+ message(STATUS "Building hadacore")
1036
+ endif()
1037
+
1038
message(STATUS "Enabling C_stable extension.")
1039
define_extension_target(
1040
_C_stable_libtorch
csrc/core/scalar_type.hpp
+26
-18
@@ -1,7 +1,13 @@
1
#pragma once
2
3
-// For TORCH_CHECK
4
-#include <torch/library.h>
3
+#include <cstdint>
4
+#include <string>
5
+#include <tuple>
6
+#include <utility>
7
+#include <variant>
8
+
9
+// For STD_TORCH_CHECK
10
+#include <torch/headeronly/util/Exception.h>
11
12
namespace vllm {
13
@@ -45,7 +51,7 @@ class ScalarType {
51
// IEEE 754 compliant floating point type
52
static constexpr ScalarType float_IEEE754(uint8_t exponent,
53
uint8_t mantissa) {
48
- TORCH_CHECK(mantissa > 0 && exponent > 0);
54
+ STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
55
return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
56
}
57
@@ -53,11 +59,12 @@ class ScalarType {
59
static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
60
bool finite_values_only,
61
NanRepr nan_repr) {
56
- TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
57
- TORCH_CHECK(mantissa > 0 && exponent > 0);
58
- TORCH_CHECK(nan_repr != NAN_IEEE_754,
59
- "use `float_IEEE754` constructor for floating point types that "
60
- "follow IEEE 754 conventions");
62
+ STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
63
+ STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
64
+ STD_TORCH_CHECK(
65
+ nan_repr != NAN_IEEE_754,
66
+ "use `float_IEEE754` constructor for floating point types that "
67
+ "follow IEEE 754 conventions");
68
return ScalarType(exponent, mantissa, true, 0, finite_values_only,
69
nan_repr);
70
}
@@ -176,8 +183,8 @@ class ScalarType {
183
184
private:
185
double _floating_point_max() const {
179
- TORCH_CHECK(mantissa <= 52 && exponent <= 11,
180
- "Cannot represent max/min as a double for type ", str());
186
+ STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
187
+ "Cannot represent max/min as a double for type ", str());
188
189
uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
190
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
@@ -186,8 +193,8 @@ class ScalarType {
193
194
uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
195
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
189
- TORCH_CHECK(exponent < 11,
190
- "Cannot represent max/min as a double for type ", str());
196
+ STD_TORCH_CHECK(exponent < 11,
197
+ "Cannot represent max/min as a double for type ", str());
198
max_exponent += 1;
199
}
200
@@ -216,16 +223,17 @@ class ScalarType {
223
if (is_floating_point()) {
224
return {_floating_point_max()};
225
} else {
219
- TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
220
- "Cannot represent max as a int64_t");
226
+ STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
227
+ "Cannot represent max as a int64_t");
228
return {(int64_t(1) << mantissa) - 1};
229
}
230
}
231
232
constexpr std::variant<int64_t, double> _raw_min() const {
233
if (is_floating_point()) {
227
- TORCH_CHECK(is_signed(),
228
- "We currently assume all floating point types are signed");
234
+ STD_TORCH_CHECK(
235
+ is_signed(),
236
+ "We currently assume all floating point types are signed");
237
constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
238
239
double max = _floating_point_max();
@@ -233,8 +241,8 @@ class ScalarType {
241
uint64_t min_raw = max_raw | sign_bit_double;
242
return {*reinterpret_cast<double*>(&min_raw)};
243
} else {
236
- TORCH_CHECK(!is_signed() || size_bits() <= 64,
237
- "Cannot represent min as a int64_t");
244
+ STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
245
+ "Cannot represent min as a int64_t");
246
if (is_signed()) {
247
// set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
248
// then perform an arithmetic shift right to set all the bits above
csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp
renamed
csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp
renamed
csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
renamed
csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp
renamed
csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu
renamed
+55
-56
@@ -18,13 +18,12 @@ limitations under the License.
18
* Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
19
* by Alcanderian JieXin Liang
20
*/
21
-#include "core/registration.h"
21
+#include "libtorch_stable/torch_utils.h"
22
+
23
+#include <torch/csrc/stable/library.h>
24
23
-#include <ATen/cuda/CUDAContext.h>
24
-#include <c10/cuda/CUDAGuard.h>
25
#include <cutlass/cutlass.h>
26
#include <cutlass/kernel_hardware_info.h>
27
-#include <torch/all.h>
27
28
#include <cute/tensor.hpp>
29
#include <iostream>
@@ -35,27 +34,27 @@ limitations under the License.
34
// clang-format off
35
#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
36
void sm100_cutlass_mla_decode(
38
- torch::Tensor const& out,
39
- torch::Tensor const& lse,
40
- torch::Tensor const& q_nope,
41
- torch::Tensor const& q_pe,
42
- torch::Tensor const& kv_c_and_k_pe_cache,
43
- torch::Tensor const& seq_lens,
44
- torch::Tensor const& page_table,
45
- torch::Tensor const& workspace,
37
+ torch::stable::Tensor const& out,
38
+ torch::stable::Tensor const& lse,
39
+ torch::stable::Tensor const& q_nope,
40
+ torch::stable::Tensor const& q_pe,
41
+ torch::stable::Tensor const& kv_c_and_k_pe_cache,
42
+ torch::stable::Tensor const& seq_lens,
43
+ torch::stable::Tensor const& page_table,
44
+ torch::stable::Tensor const& workspace,
45
double sm_scale,
46
int64_t num_kv_splits) {
48
- TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
47
+ STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
48
}
49
int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
51
- TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
50
+ STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
51
}
52
#else
53
54
#define CUTLASS_CHECK(status) \
55
{ \
56
cutlass::Status error = status; \
58
- TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
57
+ STD_TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
58
}
59
60
using namespace cute;
@@ -100,23 +99,23 @@ struct MlaSm100 {
99
100
template <typename T>
101
typename T::Fmha::Arguments args_from_options(
103
- at::Tensor const& out,
104
- at::Tensor const& lse,
105
- at::Tensor const& q_nope,
106
- at::Tensor const& q_pe,
107
- at::Tensor const& kv_c_and_k_pe_cache,
108
- at::Tensor const& seq_lens,
109
- at::Tensor const& page_table,
102
+ torch::stable::Tensor const& out,
103
+ torch::stable::Tensor const& lse,
104
+ torch::stable::Tensor const& q_nope,
105
+ torch::stable::Tensor const& q_pe,
106
+ torch::stable::Tensor const& kv_c_and_k_pe_cache,
107
+ torch::stable::Tensor const& seq_lens,
108
+ torch::stable::Tensor const& page_table,
109
double sm_scale,
110
int64_t num_kv_splits) {
111
cutlass::KernelHardwareInfo hw_info;
113
- hw_info.device_id = q_nope.device().index();
112
+ hw_info.device_id = q_nope.get_device_index();
113
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
114
116
- int batches = q_nope.sizes()[0];
117
- int page_count_per_seq = page_table.sizes()[1];
118
- int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
119
- int page_size = kv_c_and_k_pe_cache.sizes()[1];
115
+ int batches = q_nope.size(0);
116
+ int page_count_per_seq = page_table.size(1);
117
+ int page_count_total = kv_c_and_k_pe_cache.size(0);
118
+ int page_size = kv_c_and_k_pe_cache.size(1);
119
int max_seq_len = page_size * page_count_per_seq;
120
using TileShapeH = typename T::TileShapeH;
121
using TileShapeD = typename T::TileShapeD;
@@ -186,14 +185,14 @@ typename T::Fmha::Arguments args_from_options(
185
186
template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
187
void runMla(
189
- at::Tensor const& out,
190
- at::Tensor const& lse,
191
- at::Tensor const& q_nope,
192
- at::Tensor const& q_pe,
193
- at::Tensor const& kv_c_and_k_pe_cache,
194
- at::Tensor const& seq_lens,
195
- at::Tensor const& page_table,
196
- at::Tensor const& workspace,
188
+ torch::stable::Tensor const& out,
189
+ torch::stable::Tensor const& lse,
190
+ torch::stable::Tensor const& q_nope,
191
+ torch::stable::Tensor const& q_pe,
192
+ torch::stable::Tensor const& kv_c_and_k_pe_cache,
193
+ torch::stable::Tensor const& seq_lens,
194
+ torch::stable::Tensor const& page_table,
195
+ torch::stable::Tensor const& workspace,
196
double sm_scale,
197
int64_t num_kv_splits,
198
cudaStream_t stream) {
@@ -220,37 +219,37 @@ void runMla(
219
}()
220
221
void sm100_cutlass_mla_decode(
223
- torch::Tensor const& out,
224
- torch::Tensor const& lse,
225
- torch::Tensor const& q_nope,
226
- torch::Tensor const& q_pe,
227
- torch::Tensor const& kv_c_and_k_pe_cache,
228
- torch::Tensor const& seq_lens,
229
- torch::Tensor const& page_table,
230
- torch::Tensor const& workspace,
222
+ torch::stable::Tensor const& out,
223
+ torch::stable::Tensor const& lse,
224
+ torch::stable::Tensor const& q_nope,
225
+ torch::stable::Tensor const& q_pe,
226
+ torch::stable::Tensor const& kv_c_and_k_pe_cache,
227
+ torch::stable::Tensor const& seq_lens,
228
+ torch::stable::Tensor const& page_table,
229
+ torch::stable::Tensor const& workspace,
230
double sm_scale,
231
int64_t num_kv_splits) {
233
- auto in_dtype = q_nope.dtype();
234
- at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
235
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
236
- const int page_size = kv_c_and_k_pe_cache.sizes()[1];
237
-
232
+ auto in_dtype = q_nope.scalar_type();
233
+ torch::stable::accelerator::DeviceGuard device_guard(q_nope.get_device_index());
234
+ const cudaStream_t stream = get_current_cuda_stream(q_nope.get_device_index());
235
+ const int page_size = kv_c_and_k_pe_cache.size(1);
236
+
237
// NOTE(alcanderian): IsPersistent has bug with manual split_kv.
238
// Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
239
// Maybe per batch split kv will fix this.
240
DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
241
DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
243
- if (in_dtype == at::ScalarType::Half) {
242
+ if (in_dtype == torch::headeronly::ScalarType::Half) {
243
runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
244
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
246
- } else if (in_dtype == at::ScalarType::BFloat16) {
245
+ } else if (in_dtype == torch::headeronly::ScalarType::BFloat16) {
246
runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
247
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
249
- } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
248
+ } else if (in_dtype == torch::headeronly::ScalarType::Float8_e4m3fn) {
249
runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
250
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
251
} else {
253
- TORCH_CHECK(false, "Unsupported input data type of MLA");
252
+ STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
253
}
254
return true;
255
});
@@ -280,12 +279,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
279
280
#endif
281
283
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
284
- m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
282
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
283
+ m.impl("sm100_cutlass_mla_decode", TORCH_BOX(&sm100_cutlass_mla_decode));
284
}
285
287
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
288
- m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
286
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
287
+ m.impl("sm100_cutlass_mla_get_workspace_size", TORCH_BOX(&sm100_cutlass_mla_get_workspace_size));
288
}
289
290
// clang-format on
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
renamed
+38
-33
@@ -20,13 +20,15 @@
20
* limitations under the License.
21
*/
22
23
-#include <ATen/ATen.h>
24
-#include <ATen/cuda/CUDAContext.h>
25
-#include <cuda_bf16.h>
26
-#include <cuda_runtime.h>
27
-#include <torch/all.h>
23
+#include <torch/csrc/stable/library.h>
24
+#include <torch/csrc/stable/tensor.h>
25
+#include <torch/headeronly/core/ScalarType.h>
26
27
#include "core/registration.h"
28
+#include "libtorch_stable/torch_utils.h"
29
+
30
+#include <cuda_bf16.h>
31
+#include <cuda_runtime.h>
32
33
#include <cstdlib>
34
#include <mutex>
@@ -34,7 +36,7 @@
36
namespace {
37
38
inline int getSMVersion() {
37
- auto* props = at::cuda::getCurrentDeviceProperties();
39
+ auto* props = get_device_prop();
40
return props->major * 10 + props->minor;
41
}
42
@@ -700,37 +702,40 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
702
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
703
cudaStream_t);
704
703
-void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
704
- torch::Tensor const& mat_b) {
705
- TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
705
+void dsv3_fused_a_gemm(torch::stable::Tensor& output,
706
+ torch::stable::Tensor const& mat_a,
707
+ torch::stable::Tensor const& mat_b) {
708
+ STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
709
int const num_tokens = mat_a.size(0);
710
int const hd_in = mat_a.size(1);
711
int const hd_out = mat_b.size(1);
712
713
constexpr int kHdIn = 7168;
714
constexpr int kHdOut = 2112;
712
- TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
713
- "required 1 <= mat_a.shape[0] <= 16")
714
- TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
715
- TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
716
- TORCH_CHECK(output.size(0) == num_tokens,
717
- "required output.shape[0] == mat_a.shape[0]")
718
- TORCH_CHECK(output.size(1) == hd_out,
719
- "required output.shape[1] == mat_b.shape[1]")
720
-
721
- TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
722
- TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
723
- TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
724
-
725
- TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
726
- mat_b.scalar_type() == torch::kBFloat16,
727
- "Only BFloat16 input dtype is supported")
728
- TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
729
- "Only BFloat16 output dtype is supported")
730
-
731
- TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
732
-
733
- auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
715
+ STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
716
+ "required 1 <= mat_a.shape[0] <= 16");
717
+ STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
718
+ STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
719
+ STD_TORCH_CHECK(output.size(0) == num_tokens,
720
+ "required output.shape[0] == mat_a.shape[0]");
721
+ STD_TORCH_CHECK(output.size(1) == hd_out,
722
+ "required output.shape[1] == mat_b.shape[1]");
723
+
724
+ STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
725
+ STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
726
+ STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
727
+
728
+ STD_TORCH_CHECK(
729
+ mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
730
+ mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
731
+ "Only BFloat16 input dtype is supported");
732
+ STD_TORCH_CHECK(
733
+ output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
734
+ "Only BFloat16 output dtype is supported");
735
+
736
+ STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
737
+
738
+ auto stream = get_current_cuda_stream(mat_a.get_device_index());
739
if (num_tokens <= 8) {
740
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
741
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
@@ -746,6 +751,6 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
751
}
752
}
753
749
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
750
- m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
754
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
755
+ m.impl("dsv3_fused_a_gemm", TORCH_BOX(&dsv3_fused_a_gemm));
756
}
csrc/libtorch_stable/ops.h
+22
@@ -143,4 +143,26 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
143
const torch::stable::Tensor& expert_offsets,
144
const torch::stable::Tensor& sf_offsets);
145
146
+// AWQ ops
147
+torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
148
+ torch::stable::Tensor _kernel,
149
+ torch::stable::Tensor _scaling_factors,
150
+ torch::stable::Tensor _zeros,
151
+ int64_t split_k_iters);
152
+
153
+torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
154
+ torch::stable::Tensor _scaling_factors,
155
+ torch::stable::Tensor _zeros,
156
+ int64_t split_k_iters, int64_t thx,
157
+ int64_t thy);
158
+
159
+// DSV3 fused A GEMM: conditionally compiled so declaration and impl
160
+// registration are in the source file (dsv3_fused_a_gemm.cu)
161
+
162
+// AllSpark ops: declarations are in the source files
163
+// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
164
+
165
#endif
166
+
167
+torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
168
+ bool inplace);
csrc/libtorch_stable/quantization/awq/dequantize.cuh
renamed
csrc/libtorch_stable/quantization/awq/gemm_kernels.cu
renamed
+41
-35
@@ -7,10 +7,11 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
7
}
8
*/
9
10
-#include <torch/all.h>
11
-#include <c10/cuda/CUDAGuard.h>
10
+#include <torch/csrc/stable/tensor.h>
11
+#include <torch/csrc/stable/ops.h>
12
+#include "libtorch_stable/torch_utils.h"
13
13
-#include "dequantize.cuh"
14
+#include "libtorch_stable/quantization/awq/dequantize.cuh"
15
16
#include <cuda_fp16.h>
17
@@ -410,10 +411,11 @@ __global__ void __launch_bounds__(64)
411
} // namespace awq
412
} // namespace vllm
413
413
-torch::Tensor awq_dequantize(torch::Tensor _kernel,
414
- torch::Tensor _scaling_factors,
415
- torch::Tensor _zeros, int64_t split_k_iters,
416
- int64_t thx, int64_t thy) {
414
+torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
415
+ torch::stable::Tensor _scaling_factors,
416
+ torch::stable::Tensor _zeros,
417
+ int64_t split_k_iters, int64_t thx,
418
+ int64_t thy) {
419
int in_c = _kernel.size(0);
420
int qout_c = _kernel.size(1);
421
int out_c = qout_c * 8;
@@ -437,23 +439,24 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
439
y_blocks = (int)(in_c / 8);
440
}
441
440
- const at::cuda::OptionalCUDAGuard device_guard(device_of(_scaling_factors));
442
+ const torch::stable::accelerator::DeviceGuard device_guard(
443
+ _scaling_factors.get_device_index());
444
442
- auto options = torch::TensorOptions()
443
- .dtype(_scaling_factors.dtype())
444
- .device(_scaling_factors.device());
445
- at::Tensor _de_kernel = torch::empty({in_c, out_c}, options);
445
+ auto _de_kernel =
446
+ torch::stable::empty({in_c, out_c}, _scaling_factors.scalar_type(),
447
+ std::nullopt, _scaling_factors.device());
448
447
- auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
448
- auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
449
- auto scaling_factors =
450
- reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
451
- auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
449
+ auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
450
+ auto de_kernel = reinterpret_cast<half*>(
451
+ _de_kernel.mutable_data_ptr<torch::headeronly::Half>());
452
+ auto scaling_factors = reinterpret_cast<half*>(
453
+ _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
454
+ auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
455
456
dim3 num_blocks(x_blocks, y_blocks);
457
dim3 threads_per_block(x_thread, y_thread);
458
456
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
459
+ const cudaStream_t stream = get_current_cuda_stream();
460
vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
461
kernel, scaling_factors, zeros, de_kernel, G);
462
@@ -466,27 +469,30 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
469
// zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b]
470
// assume that batch_size < 16 for now
471
469
-torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
470
- torch::Tensor _scaling_factors, torch::Tensor _zeros,
471
- int64_t split_k_iters) {
472
+torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
473
+ torch::stable::Tensor _kernel,
474
+ torch::stable::Tensor _scaling_factors,
475
+ torch::stable::Tensor _zeros,
476
+ int64_t split_k_iters) {
477
int num_in_feats = _in_feats.size(0);
478
int num_in_channels = _in_feats.size(1);
474
- const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats));
479
+ const torch::stable::accelerator::DeviceGuard device_guard(
480
+ _in_feats.get_device_index());
481
476
- auto options = torch::TensorOptions()
477
- .dtype(_in_feats.dtype())
478
- .device(_in_feats.device());
479
- at::Tensor _out_feats =
480
- torch::empty({split_k_iters, num_in_feats, _kernel.size(1) * 8}, options);
482
+ auto _out_feats = torch::stable::empty(
483
+ {split_k_iters, num_in_feats, _kernel.size(1) * 8},
484
+ _in_feats.scalar_type(), std::nullopt, _in_feats.device());
485
int num_out_feats = _out_feats.size(-2);
486
int num_out_channels = _out_feats.size(-1);
487
484
- auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
485
- auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
486
- auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
487
- auto scaling_factors =
488
- reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
489
- auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
488
+ auto in_feats = reinterpret_cast<half*>(
489
+ _in_feats.mutable_data_ptr<torch::headeronly::Half>());
490
+ auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
491
+ auto out_feats = reinterpret_cast<half*>(
492
+ _out_feats.mutable_data_ptr<torch::headeronly::Half>());
493
+ auto scaling_factors = reinterpret_cast<half*>(
494
+ _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
495
+ auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
496
int group_size = num_in_channels / _scaling_factors.size(0);
497
498
if (num_out_channels % 64 != 0)
@@ -498,7 +504,7 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
504
if (num_out_channels % group_size != 0)
505
throw std::invalid_argument("OC is not multiple of Group size");
506
501
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
507
+ const cudaStream_t stream = get_current_cuda_stream();
508
if (num_out_channels % 128 == 0) {
509
int j_factors1 = num_out_channels / 128 / 1;
510
dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
@@ -522,5 +528,5 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
528
group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
529
num_in_feats, num_in_channels, num_out_channels, out_feats);
530
}
525
- return _out_feats.sum(0);
531
+ return torch::stable::sum(_out_feats, 0);
532
}
csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu
renamed
+56
-43
@@ -1,20 +1,28 @@
1
#include "allspark_utils.cuh"
2
-#include <torch/all.h>
3
-#include "core/registration.h"
2
+
3
+#include <torch/csrc/stable/c/shim.h>
4
+#include <torch/csrc/stable/library.h>
5
+#include <torch/csrc/stable/tensor.h>
6
+#include <torch/headeronly/core/ScalarType.h>
7
+
8
#include <cublas_v2.h>
9
6
-at::Tensor as_g_workspace;
10
+#include "core/registration.h"
11
+#include "libtorch_stable/torch_utils.h"
12
+
13
+torch::stable::Tensor as_g_workspace;
14
15
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
16
10
-torch::Tensor allspark_w8a16_gemm(
11
- torch::Tensor const& a, torch::Tensor const& b_qweight,
12
- torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
13
- int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
17
+torch::stable::Tensor allspark_w8a16_gemm(
18
+ torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
19
+ torch::stable::Tensor const& b_scales,
20
+ std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
21
+ int64_t group_size, int64_t sm_count, int64_t sm_version,
22
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
15
- TORCH_CHECK_NOT_IMPLEMENTED(
23
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
24
false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
17
- return torch::empty({1, 1});
25
+ return torch::stable::empty({1, 1});
26
}
27
28
#else
@@ -848,8 +856,8 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
856
const int N_32align, const int N,
857
const int K, const int GroupSize,
858
cudaStream_t stream) {
851
- TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
852
- "Unsupported shape");
859
+ STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
860
+ "Unsupported shape");
861
if (GroupSize == -1) {
862
const int BLOCK = 128;
863
dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
@@ -859,7 +867,7 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
867
}
868
// TODO: Support SubChannel
869
else {
862
- TORCH_CHECK(false, "Now only support PerChannel");
870
+ STD_TORCH_CHECK(false, "Now only support PerChannel");
871
}
872
}
873
@@ -916,24 +924,27 @@ void allspark_qgemm_w8a16_perc_ampere(
924
925
} // namespace allspark
926
919
-torch::Tensor allspark_w8a16_gemm(
920
- torch::Tensor const& a, torch::Tensor const& b_qweight,
921
- torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
922
- int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
927
+torch::stable::Tensor allspark_w8a16_gemm(
928
+ torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
929
+ torch::stable::Tensor const& b_scales,
930
+ std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
931
+ int64_t group_size, int64_t sm_count, int64_t sm_version,
932
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
933
// Verify device and strides
925
- TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
926
- TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
934
+ STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
935
+ STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
936
928
- TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
929
- TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
937
+ STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
938
+ STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
939
931
- TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
932
- TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
940
+ STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
941
+ STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
942
943
if (has_zp) {
935
- TORCH_CHECK(b_qzeros.value().device().is_cuda(), "b_qzeros is not on GPU");
936
- TORCH_CHECK(b_qzeros.value().is_contiguous(), "b_qzeros is not contiguous");
944
+ STD_TORCH_CHECK(b_qzeros.value().device().is_cuda(),
945
+ "b_qzeros is not on GPU");
946
+ STD_TORCH_CHECK(b_qzeros.value().is_contiguous(),
947
+ "b_qzeros is not contiguous");
948
}
949
950
int m = a.size(0);
@@ -941,16 +952,17 @@ torch::Tensor allspark_w8a16_gemm(
952
int k = a.size(1);
953
954
// Verify shape
944
- TORCH_CHECK(b_qweight.size(0) == n_32align,
945
- "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
946
- ", n_32align = ", n_32align);
947
- TORCH_CHECK(b_qweight.size(1) == k,
948
- "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
949
- ", k = ", k);
955
+ STD_TORCH_CHECK(b_qweight.size(0) == n_32align,
956
+ "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
957
+ ", n_32align = ", n_32align);
958
+ STD_TORCH_CHECK(b_qweight.size(1) == k,
959
+ "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
960
+ ", k = ", k);
961
951
- TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
962
+ STD_TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
963
953
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
964
+ const torch::stable::accelerator::DeviceGuard device_guard(
965
+ a.get_device_index());
966
const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
967
const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
968
const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
@@ -959,12 +971,12 @@ torch::Tensor allspark_w8a16_gemm(
971
b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
972
}
973
962
- auto c_options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
963
- torch::Tensor c = torch::empty({m, n}, c_options);
964
- void* c_ptr = reinterpret_cast<void*>(c.data_ptr());
974
+ auto c =
975
+ torch::stable::empty({m, n}, a.scalar_type(), std::nullopt, a.device());
976
+ void* c_ptr = reinterpret_cast<void*>(c.mutable_data_ptr());
977
966
- cudaStream_t stream = at::cuda::getCurrentCUDAStream();
967
- cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
978
+ cudaStream_t stream = get_current_cuda_stream();
979
+ cublasHandle_t handle = get_current_cuda_blas_handle();
980
981
allspark::BlockTileSplitkParams fused_gemm_params;
982
@@ -976,14 +988,15 @@ torch::Tensor allspark_w8a16_gemm(
988
m, n, k, sm_count, fused_gemm_params);
989
}
990
979
- auto ws_options = torch::TensorOptions().dtype(at::kChar).device(a.device());
991
if (as_g_workspace.numel() <
992
ws_size) { // ws_options: kChar, so numel() is bytes
982
- as_g_workspace = torch::empty({long(ws_size)}, ws_options);
993
+ as_g_workspace = torch::stable::empty({static_cast<int64_t>(ws_size)},
994
+ torch::headeronly::ScalarType::Char,
995
+ std::nullopt, a.device());
996
}
997
void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());
998
986
- if (a.dtype() == at::ScalarType::Half) {
999
+ if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
1000
allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
1001
reinterpret_cast<const __half*>(a_ptr), b_ptr,
1002
reinterpret_cast<const __half*>(b_scale_ptr),
@@ -991,7 +1004,7 @@ torch::Tensor allspark_w8a16_gemm(
1004
reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
1005
fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
1006
handle);
994
- } else if (a.dtype() == at::ScalarType::BFloat16) {
1007
+ } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
1008
allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
1009
reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
1010
reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
@@ -1006,6 +1019,6 @@ torch::Tensor allspark_w8a16_gemm(
1019
1020
#endif
1021
1009
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1010
- m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
1022
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
1023
+ m.impl("allspark_w8a16_gemm", TORCH_BOX(&allspark_w8a16_gemm));
1024
}
csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu
renamed
+44
-33
@@ -1,6 +1,11 @@
1
#include "allspark_utils.cuh"
2
-#include <torch/all.h>
2
+
3
+#include <torch/csrc/stable/library.h>
4
+#include <torch/csrc/stable/tensor.h>
5
+#include <torch/headeronly/core/ScalarType.h>
6
+
7
#include "core/registration.h"
8
+#include "libtorch_stable/torch_utils.h"
9
10
namespace allspark {
11
@@ -99,36 +104,40 @@ void rearrange_kn_weight_as_n32k16_order_ldg16(
104
} // namespace allspark
105
106
void rearrange_kn_weight_as_n32k16_order(
102
- torch::Tensor const& b_qweight, torch::Tensor const& b_scales,
103
- std::optional<torch::Tensor> const& b_zeros, bool has_zp,
104
- torch::Tensor& b_qweight_reorder, torch::Tensor& b_scales_reorder,
105
- std::optional<torch::Tensor> const& b_zeros_reorder, const int64_t K,
106
- const int64_t N, const int64_t N_32align) {
107
+ torch::stable::Tensor const& b_qweight,
108
+ torch::stable::Tensor const& b_scales,
109
+ std::optional<torch::stable::Tensor> const& b_zeros, bool has_zp,
110
+ torch::stable::Tensor& b_qweight_reorder,
111
+ torch::stable::Tensor& b_scales_reorder,
112
+ std::optional<torch::stable::Tensor> const& b_zeros_reorder,
113
+ const int64_t K, const int64_t N, const int64_t N_32align) {
114
// Verify device and strides
108
- TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
109
- TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
115
+ STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
116
+ STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
117
111
- TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
112
- TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
118
+ STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
119
+ STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
120
114
- TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
115
- "b_qweight_reorder is not on GPU");
116
- TORCH_CHECK(b_qweight_reorder.is_contiguous(),
117
- "b_qweight_reorder is not contiguous");
121
+ STD_TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
122
+ "b_qweight_reorder is not on GPU");
123
+ STD_TORCH_CHECK(b_qweight_reorder.is_contiguous(),
124
+ "b_qweight_reorder is not contiguous");
125
119
- TORCH_CHECK(b_scales_reorder.device().is_cuda(),
120
- "b_scales_reorder is not on GPU");
121
- TORCH_CHECK(b_scales_reorder.is_contiguous(),
122
- "b_scales_reorder is not contiguous");
126
+ STD_TORCH_CHECK(b_scales_reorder.device().is_cuda(),
127
+ "b_scales_reorder is not on GPU");
128
+ STD_TORCH_CHECK(b_scales_reorder.is_contiguous(),
129
+ "b_scales_reorder is not contiguous");
130
131
if (has_zp) {
125
- TORCH_CHECK(b_zeros.value().device().is_cuda(), "b_zeros is not on GPU");
126
- TORCH_CHECK(b_zeros.value().is_contiguous(), "b_zeros is not contiguous");
127
-
128
- TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
129
- "b_zeros_reorder is not on GPU");
130
- TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
131
- "b_zeros_reorder is not contiguous");
132
+ STD_TORCH_CHECK(b_zeros.value().device().is_cuda(),
133
+ "b_zeros is not on GPU");
134
+ STD_TORCH_CHECK(b_zeros.value().is_contiguous(),
135
+ "b_zeros is not contiguous");
136
+
137
+ STD_TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
138
+ "b_zeros_reorder is not on GPU");
139
+ STD_TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
140
+ "b_zeros_reorder is not contiguous");
141
}
142
143
const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
@@ -136,18 +145,20 @@ void rearrange_kn_weight_as_n32k16_order(
145
const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;
146
147
uint8_t* matB_reorder =
139
- reinterpret_cast<uint8_t*>(b_qweight_reorder.data_ptr());
140
- void* b_scale_reorder = b_scales_reorder.data_ptr();
141
- void* b_zero_reorder = has_zp ? b_zeros_reorder.value().data_ptr() : nullptr;
148
+ reinterpret_cast<uint8_t*>(b_qweight_reorder.mutable_data_ptr());
149
+ void* b_scale_reorder = b_scales_reorder.mutable_data_ptr();
150
+ void* b_zero_reorder =
151
+ has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;
152
143
- cudaStream_t stream = at::cuda::getCurrentCUDAStream();
144
- if (b_scales.dtype() == at::ScalarType::Half) {
153
+ cudaStream_t stream = get_current_cuda_stream();
154
+ if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
155
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
156
matB, reinterpret_cast<const __half*>(b_scale),
157
reinterpret_cast<const __half*>(b_zero), matB_reorder,
158
reinterpret_cast<__half*>(b_scale_reorder),
159
reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
150
- } else if (b_scales.dtype() == at::ScalarType::BFloat16) {
160
+ } else if (b_scales.scalar_type() ==
161
+ torch::headeronly::ScalarType::BFloat16) {
162
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
163
matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
164
reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
@@ -157,7 +168,7 @@ void rearrange_kn_weight_as_n32k16_order(
168
}
169
}
170
160
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
171
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
172
m.impl("rearrange_kn_weight_as_n32k16_order",
162
- &rearrange_kn_weight_as_n32k16_order);
173
+ TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
174
}
csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh
renamed
+5
-6
@@ -1,13 +1,12 @@
1
#pragma once
2
3
-#include <torch/all.h>
4
-#include <c10/cuda/CUDAGuard.h>
5
-#include <ATen/cuda/CUDAContext.h>
6
-#include <cuda_runtime.h>
7
-#include <cuda_fp16.h>
3
#include <cuda_bf16.h>
4
+#include <cuda_fp16.h>
5
+#include <cuda_runtime.h>
6
+
7
#include <iostream>
10
-#include "../marlin/marlin_dtypes.cuh"
8
+
9
+#include "quantization/marlin/marlin_dtypes.cuh"
10
using marlin::MarlinScalarType2;
11
12
namespace allspark {
csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
renamed
+41
-43
@@ -11,18 +11,16 @@ Redistribution and use in source and binary forms, with or without modification,
11
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
12
***********/
13
14
-#include <torch/all.h>
14
+#include "libtorch_stable/torch_utils.h"
15
+#include "libtorch_stable/dispatch_utils.h"
16
+
17
+#include <torch/csrc/stable/library.h>
18
+#include <torch/csrc/stable/macros.h>
19
+
20
#include <stdint.h>
21
#include <cuda_runtime.h>
22
#include <mma.h>
23
#include <cuda/annotated_ptr>
19
-#include <c10/cuda/CUDAException.h>
20
-
21
-#include <ATen/cuda/CUDAContext.h>
22
-#include <c10/cuda/CUDAGuard.h>
23
-
24
-#include "core/registration.h"
25
-#include "dispatch_utils.h"
24
25
namespace hadacore {
26
@@ -65,12 +63,12 @@ constexpr int launch_configs_big[7][3] = {
63
};
64
65
// a 4x2, b 2x2, c 2x2
68
-template <torch::ScalarType dtype>
66
+template <torch::headeronly::ScalarType dtype>
67
__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
70
- static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
68
+ static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16);
69
// d, a, b, c
70
b32 zero = 0;
73
- if constexpr(dtype == torch::ScalarType::Half) {
71
+ if constexpr(dtype == torch::headeronly::ScalarType::Half) {
72
asm (
73
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
74
"{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
@@ -89,7 +87,7 @@ __device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1,
87
}
88
89
// a 4x2, b 4x2, c 4x2
92
-template <torch::ScalarType dtype>
90
+template <torch::headeronly::ScalarType dtype>
91
__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
92
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
93
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
@@ -108,11 +106,11 @@ __device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
106
#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
107
#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)
108
111
-template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
109
+template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::headeronly::ScalarType dtype>
110
__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
111
// a is column major, b is row major
112
hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
115
- static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
113
+ static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
114
115
b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)
116
@@ -162,8 +160,8 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
160
constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
161
constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};
162
165
- #define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
166
- #define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
163
+ #define val_type_1p(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
164
+ #define val_type_1n(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
165
constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
166
constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};
167
@@ -684,14 +682,14 @@ constexpr int64_t ceil_div(int64_t a, int64_t b) {
682
return (a + b - 1) / b;
683
}
684
687
-template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
685
+template <torch::headeronly::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
686
void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
687
int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
688
dim3 block_size = 32 * warps_per_block;
689
690
#define CHECK_SHARED_LIM() { \
691
if (shared_size > 48 * 1024) { \
694
- C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
692
+ STD_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
693
} \
694
} \
695
@@ -714,10 +712,10 @@ void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaSt
712
kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
713
}
714
717
- C10_CUDA_KERNEL_LAUNCH_CHECK();
715
+ STD_CUDA_KERNEL_LAUNCH_CHECK();
716
}
717
720
-template <torch::ScalarType dtype>
718
+template <torch::headeronly::ScalarType dtype>
719
void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
720
int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
721
// for size 256, use (2, 1)
@@ -764,54 +762,54 @@ void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cu
762
}
763
}
764
767
-template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
768
-template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
765
+template void run_fht<torch::headeronly::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
766
+template void run_fht<torch::headeronly::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
767
768
} // namespace hadacore
769
770
constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }
771
774
-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
772
+torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace) {
773
auto dtype = x.scalar_type();
776
- TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
777
- TORCH_CHECK(x.is_cuda());
778
-
774
+ STD_TORCH_CHECK(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
775
+ STD_TORCH_CHECK(x.is_cuda());
776
+
777
const int had_size = x.size(-1);
780
- TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
778
+ STD_TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
779
"Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
782
-
780
+
781
const auto res_shape = x.sizes();
784
- x = x.reshape({-1, had_size});
785
-
782
+ x = torch::stable::reshape(x, {-1, had_size});
783
+
784
auto numel = x.numel();
785
if (numel % 256 != 0) {
788
- x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
786
+ x = torch::stable::pad(x, {0, 0, 0, (256 - numel % 256) / had_size});
787
}
790
-
788
+
789
if (x.stride(-1) != 1) {
792
- x = x.contiguous();
790
+ x = torch::stable::contiguous(x);
791
}
794
- torch::Tensor out = inplace ? x : torch::empty_like(x);
792
+ torch::stable::Tensor out = inplace ? x : torch::stable::empty_like(x);
793
796
- at::cuda::CUDAGuard device_guard{(char)x.get_device()};
797
- auto stream = at::cuda::getCurrentCUDAStream().stream();
794
+ torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
795
+ auto stream = get_current_cuda_stream();
796
799
- VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
800
- auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
797
+ VLLM_STABLE_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
798
+ auto constexpr SCALAR_TYPE = torch::headeronly::CppTypeToScalarType<scalar_t>::value;
799
hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
800
});
801
802
if (numel % 256 != 0) {
805
- out = out.narrow(0, 0, numel / had_size);
803
+ out = torch::stable::narrow(out, 0, 0, numel / had_size);
804
}
805
806
if (inplace && out.data_ptr() != x.data_ptr()) {
809
- x.copy_(out.view(res_shape));
807
+ torch::stable::copy_(x, torch::stable::view(out, res_shape));
808
return x;
809
}
812
- return out.reshape(res_shape);
810
+ return torch::stable::reshape(out, res_shape);
811
}
812
815
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
816
- m.impl("hadacore_transform", &hadacore_transform);
813
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
814
+ m.impl("hadacore_transform", TORCH_BOX(&hadacore_transform));
815
}
csrc/libtorch_stable/torch_bindings.cpp
+57
@@ -218,7 +218,54 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
218
ops.def(
219
"cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
220
"Tensor)");
221
+
222
+ // SM100 CUTLASS MLA decode
223
+ // conditionally compiled so impl registrations are in source file
224
+ ops.def(
225
+ "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
226
+ " Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
227
+ " Tensor seq_lens, Tensor page_table,"
228
+ " Tensor workspace, float scale,"
229
+ " int num_kv_splits) -> ()");
230
+
231
+ ops.def(
232
+ "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
233
+ " int sm_count, int num_kv_splits) "
234
+ "-> int");
235
+ // Quantized GEMM for AWQ.
236
+ ops.def(
237
+ "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
238
+ "Tensor _zeros, SymInt split_k_iters) -> Tensor");
239
+
240
+ // Dequantization for AWQ.
241
+ ops.def(
242
+ "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
243
+ "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
244
+
245
+ // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
246
+ // conditionally compiled so impl registration is in source file
247
+ ops.def(
248
+ "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
249
+
250
+ // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
251
+ ops.def(
252
+ "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
253
+ "Tensor? b_zeros, "
254
+ "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
255
+ "Tensor!? b_zeros_reorder, "
256
+ "int K, int N, int N_32align) -> ()");
257
+
258
+ // AllSpark quantization ops
259
+ ops.def(
260
+ "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
261
+ "Tensor? b_qzeros, "
262
+ "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
263
+ "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
264
#endif
265
+
266
+ // Hadamard transforms
267
+ // conditionally compiled so impl registration is in source file
268
+ ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
269
}
270
271
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -254,6 +301,16 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
301
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
302
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
303
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
304
+
305
+ // AWQ ops
306
+ ops.impl("awq_gemm", TORCH_BOX(&awq_gemm));
307
+ ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize));
308
+
309
+ // DSV3 fused A GEMM: conditionally compiled so impl registration is in
310
+ // source file (dsv3_fused_a_gemm.cu)
311
+
312
+ // AllSpark ops: conditionally compiled so impl registrations are in source
313
+ // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
314
#endif
315
}
316
csrc/libtorch_stable/torch_utils.h
+66
@@ -6,12 +6,71 @@
6
#include <torch/csrc/stable/tensor.h>
7
#include <torch/headeronly/util/shim_utils.h>
8
9
+#include <cublas_v2.h>
10
#include <cuda_runtime.h>
11
12
+#include <deque>
13
+#include <mutex>
14
+#include <string>
15
+#include <vector>
16
+
17
// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED.
18
#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \
19
STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__)
20
21
+// Device properties cache for stable ABI compatibility.
22
+// Uses raw CUDA/HIP APIs instead of ATen functions.
23
+// Using inline ensures a single instance across all translation units.
24
+inline std::deque<std::once_flag> device_flags;
25
+inline std::vector<cudaDeviceProp> device_properties;
26
+inline std::once_flag vectors_init_flag;
27
+
28
+inline void do_init_device_vectors() {
29
+ int device_count;
30
+ cudaError_t err = cudaGetDeviceCount(&device_count);
31
+ if (err != cudaSuccess) {
32
+ STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " +
33
+ std::string(cudaGetErrorString(err)));
34
+ }
35
+ device_flags.resize(device_count);
36
+ device_properties.resize(device_count);
37
+}
38
+
39
+inline void initDeviceVectors() {
40
+ std::call_once(vectors_init_flag, do_init_device_vectors);
41
+}
42
+
43
+inline void initDeviceProperty(int device_index) {
44
+ cudaDeviceProp device_prop{};
45
+ cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index);
46
+ if (err != cudaSuccess) {
47
+ STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " +
48
+ std::string(cudaGetErrorString(err)));
49
+ }
50
+ device_properties[device_index] = device_prop;
51
+}
52
+
53
+// Get device properties using raw CUDA/HIP APIs (stable ABI compatible).
54
+// Caches results per device so cudaGetDeviceProperties is called at most once
55
+// per device.
56
+inline cudaDeviceProp* get_device_prop() {
57
+ initDeviceVectors();
58
+ int device_index;
59
+ cudaError_t err = cudaGetDevice(&device_index);
60
+ if (err != cudaSuccess) {
61
+ STD_TORCH_CHECK(
62
+ false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err)));
63
+ }
64
+ STD_TORCH_CHECK(device_index >= 0 && static_cast<size_t>(device_index) <
65
+ device_properties.size(),
66
+ "CUDA device index " + std::to_string(device_index) +
67
+ " out of range [0, " +
68
+ std::to_string(device_properties.size()) + ")");
69
+
70
+ std::call_once(device_flags[device_index], initDeviceProperty, device_index);
71
+ return &device_properties[device_index];
72
+}
73
+
74
// Utility to get the current CUDA stream for a given device using stable APIs.
75
// Returns a cudaStream_t for use in kernel launches.
76
inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
@@ -20,3 +79,10 @@ inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
79
aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
80
return reinterpret_cast<cudaStream_t>(stream_ptr);
81
}
82
+
83
+// Utility to get the current cuBLAS handle using stable APIs.
84
+inline cublasHandle_t get_current_cuda_blas_handle() {
85
+ void* blas_handle_ptr = nullptr;
86
+ TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr));
87
+ return reinterpret_cast<cublasHandle_t>(blas_handle_ptr);
88
+}
csrc/ops.h
-20
@@ -200,19 +200,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
200
201
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
202
203
-#ifndef USE_ROCM
204
-
205
-torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
206
- torch::Tensor _scaling_factors, torch::Tensor _zeros,
207
- int64_t split_k_iters);
208
-
209
-torch::Tensor awq_dequantize(torch::Tensor _kernel,
210
- torch::Tensor _scaling_factors,
211
- torch::Tensor _zeros, int64_t split_k_iters,
212
- int64_t thx, int64_t thy);
213
-
214
-#endif
215
-
203
torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
204
int64_t n,
205
std::optional<at::ScalarType> const& dtype);
@@ -302,8 +289,6 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
289
int64_t open_mem_handle(torch::Tensor& mem_handle);
290
void free_shared_buffer(int64_t buffer);
291
305
-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
306
-
292
#ifdef USE_ROCM
293
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
294
std::optional<int64_t> qr_max_size = std::nullopt);
@@ -315,11 +300,6 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
300
int64_t qr_max_size();
301
#endif
302
318
-#ifndef USE_ROCM
319
-void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
320
- torch::Tensor const& mat_b);
321
-#endif
322
-
303
#ifndef USE_ROCM
304
torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
305
torch::Tensor const& norm_weight,
csrc/quantization/marlin/marlin.cuh
+8
-4
@@ -2,10 +2,14 @@
2
3
#ifndef _marlin_cuh
4
#define _marlin_cuh
5
- #include <torch/all.h>
6
-
7
- #include <ATen/cuda/CUDAContext.h>
8
- #include <c10/cuda/CUDAGuard.h>
5
+ // These torch headers are only needed by non-stable callers (e.g. ops.cu).
6
+ // Guard them so that stable ABI targets can still include marlin.cuh
7
+ // for Vec, constants, and cp_async helpers without pulling in torch/all.h.
8
+ #ifndef TORCH_TARGET_VERSION
9
+ #include <torch/all.h>
10
+ #include <ATen/cuda/CUDAContext.h>
11
+ #include <c10/cuda/CUDAGuard.h>
12
+ #endif
13
#include <cuda.h>
14
#include <cuda_fp16.h>
15
#include <cuda_runtime.h>
csrc/torch_bindings.cpp
-51
@@ -263,22 +263,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
263
264
// Quantization ops
265
#ifndef USE_ROCM
266
- // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
267
- ops.def(
268
- "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
269
- // conditionally compiled so impl registration is in source file
270
-
271
- // Quantized GEMM for AWQ.
272
- ops.def(
273
- "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
274
- "Tensor _zeros, SymInt split_k_iters) -> Tensor");
275
- ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
276
-
277
- // Dequantization for AWQ.
278
- ops.def(
279
- "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
280
- "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
281
- ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
266
267
// Note about marlin kernel 'workspace' arguments:
268
// Technically these should be mutable since they are modified by the kernel.
@@ -408,22 +392,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
392
" -> ()");
393
// conditionally compiled so impl registration is in source file
394
411
- // SM100 CUTLASS MLA decode
412
- ops.def(
413
- "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
414
- " Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
415
- " Tensor seq_lens, Tensor page_table,"
416
- " Tensor workspace, float scale,"
417
- " int num_kv_splits) -> ()");
418
- // conditionally compiled so impl in source file
419
-
420
- // SM100 CUTLASS MLA workspace
421
- ops.def(
422
- "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
423
- " int sm_count, int num_kv_splits) "
424
- "-> int");
425
- // conditionally compiled so impl in source file
426
-
395
#endif
396
397
// Quantized GEMM for GPTQ.
@@ -496,26 +464,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
464
"Tensor? last_chunk_indices) -> ()");
465
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
466
499
- // Hadamard transforms
500
- ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
501
-
467
#ifndef USE_ROCM
503
- // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
504
- ops.def(
505
- "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
506
- "Tensor? b_zeros, "
507
- "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
508
- "Tensor!? b_zeros_reorder, "
509
- "int K, int N, int N_32align) -> ()");
510
- // conditionally compiled so impl in source file
511
-
512
- // AllSpark quantization ops
513
- ops.def(
514
- "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
515
- "Tensor? b_qzeros, "
516
- "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
517
- "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
518
-
468
ops.def(
469
"minimax_allreduce_rms("
470
"Tensor input,"