A high-throughput and memory-efficient inference and serving engine for LLMs
[11b/n] Migrate Machete kernels to torch stable ABI (#45304)
Signed-off-by: Chris Leonard <chleonar@redhat.com> Signed-off-by: Shengqi Chen <harry-chen@outlook.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
Jun 12, 2026 at 06:36 UTC
a014dddbaa67661236a8a7d0dc3d5773d4e0f60a
17 files changed
+341
-323
CMakeLists.txt
+69
-70
@@ -385,76 +385,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
385
endif()
386
endif()
387
388
- #
389
- # Machete kernels
390
-
391
- # The machete kernels only work on hopper and require CUDA 12.0 or later.
392
- # Only build Machete kernels if we are building for something compatible with sm90a
393
- cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}")
394
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS)
395
- #
396
- # For the Machete kernels we automatically generate sources for various
397
- # preselected input type pairs and schedules.
398
- # Generate sources:
399
- set(MACHETE_GEN_SCRIPT
400
- ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py)
401
- file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH)
402
-
403
- message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}")
404
- message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}")
405
-
406
- if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH}
407
- OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH})
408
- execute_process(
409
- COMMAND ${CMAKE_COMMAND} -E env
410
- PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH}
411
- ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT}
412
- RESULT_VARIABLE machete_generation_result
413
- OUTPUT_VARIABLE machete_generation_output
414
- OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
415
- ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
416
- )
417
-
418
- if (NOT machete_generation_result EQUAL 0)
419
- message(FATAL_ERROR "Machete generation failed."
420
- " Result: \"${machete_generation_result}\""
421
- "\nCheck the log for details: "
422
- "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log")
423
- else()
424
- set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH}
425
- CACHE STRING "Last run machete generate script hash" FORCE)
426
- message(STATUS "Machete generation completed successfully.")
427
- endif()
428
- else()
429
- message(STATUS "Machete generation script has not changed, skipping generation.")
430
- endif()
431
-
432
- # Add machete generated sources
433
- file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu")
434
- list(APPEND VLLM_EXT_SRC ${MACHETE_GEN_SOURCES})
435
-
436
- # forward compatible
437
- set_gencode_flags_for_srcs(
438
- SRCS "${MACHETE_GEN_SOURCES}"
439
- CUDA_ARCHS "${MACHETE_ARCHS}")
440
-
441
- list(APPEND VLLM_EXT_SRC
442
- csrc/quantization/machete/machete_pytorch.cu)
443
-
444
- message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}")
445
- else()
446
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
447
- AND MACHETE_ARCHS)
448
- message(STATUS "Not building Machete kernels as CUDA Compiler version is "
449
- "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
450
- "later if you intend on running w4a16 quantized models on "
451
- "Hopper.")
452
- else()
453
- message(STATUS "Not building Machete kernels as no compatible archs "
454
- "found in CUDA target architectures")
455
- endif()
456
- endif()
457
-
388
389
390
# if CUDA endif
@@ -533,6 +463,75 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
463
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
464
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
465
466
+ #
467
+ # Machete kernels
468
+ #
469
+ # The machete kernels only work on hopper and require CUDA 12.0 or later.
470
+ # Only build Machete kernels if we are building for something compatible with sm90a
471
+ cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}")
472
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS)
473
+ #
474
+ # For the Machete kernels we automatically generate sources for various
475
+ # preselected input type pairs and schedules.
476
+ # Generate sources:
477
+ set(MACHETE_GEN_SCRIPT
478
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py)
479
+ file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH)
480
+
481
+ message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}")
482
+ message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}")
483
+
484
+ if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH}
485
+ OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH})
486
+ execute_process(
487
+ COMMAND ${CMAKE_COMMAND} -E env
488
+ PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH}
489
+ ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT}
490
+ RESULT_VARIABLE machete_generation_result
491
+ OUTPUT_VARIABLE machete_generation_output
492
+ OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
493
+ ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log
494
+ )
495
+
496
+ if (NOT machete_generation_result EQUAL 0)
497
+ message(FATAL_ERROR "Machete generation failed."
498
+ " Result: \"${machete_generation_result}\""
499
+ "\nCheck the log for details: "
500
+ "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log")
501
+ else()
502
+ set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH}
503
+ CACHE STRING "Last run machete generate script hash" FORCE)
504
+ message(STATUS "Machete generation completed successfully.")
505
+ endif()
506
+ else()
507
+ message(STATUS "Machete generation script has not changed, skipping generation.")
508
+ endif()
509
+
510
+ # Add machete generated sources
511
+ file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu")
512
+ list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES})
513
+
514
+ # forward compatible
515
+ set_gencode_flags_for_srcs(
516
+ SRCS "${MACHETE_GEN_SOURCES}"
517
+ CUDA_ARCHS "${MACHETE_ARCHS}")
518
+
519
+ list(APPEND VLLM_STABLE_EXT_SRC
520
+ csrc/libtorch_stable/quantization/machete/machete_pytorch.cu)
521
+ message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}")
522
+ else()
523
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
524
+ AND MACHETE_ARCHS)
525
+ message(STATUS "Not building Machete kernels as CUDA Compiler version is "
526
+ "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
527
+ "later if you intend on running w4a16 quantized models on "
528
+ "Hopper.")
529
+ else()
530
+ message(STATUS "Not building Machete kernels as no compatible archs "
531
+ "found in CUDA target architectures")
532
+ endif()
533
+ endif()
534
+
535
set_gencode_flags_for_srcs(
536
SRCS "${VLLM_STABLE_EXT_SRC}"
537
CUDA_ARCHS "${CUDA_ARCHS}")
csrc/cutlass_extensions/vllm_cutlass_library_extension.py
+7
-7
@@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {
57
}
58
59
VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {
60
- DataType.u8: "at::ScalarType::Byte",
61
- DataType.s8: "at::ScalarType::Char",
62
- DataType.e4m3: "at::ScalarType::Float8_e4m3fn",
63
- DataType.s32: "at::ScalarType::Int",
64
- DataType.f16: "at::ScalarType::Half",
65
- DataType.bf16: "at::ScalarType::BFloat16",
66
- DataType.f32: "at::ScalarType::Float",
60
+ DataType.u8: "torch::headeronly::ScalarType::Byte",
61
+ DataType.s8: "torch::headeronly::ScalarType::Char",
62
+ DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn",
63
+ DataType.s32: "torch::headeronly::ScalarType::Int",
64
+ DataType.f16: "torch::headeronly::ScalarType::Half",
65
+ DataType.bf16: "torch::headeronly::ScalarType::BFloat16",
66
+ DataType.f32: "torch::headeronly::ScalarType::Float",
67
}
68
69
VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {
csrc/libtorch_stable/quantization/machete/Readme.md
renamed
csrc/libtorch_stable/quantization/machete/generate.py
renamed
+18
-18
@@ -39,10 +39,10 @@ namespace machete {
39
{% for impl_config in impl_configs %}
40
{% set type_sig = gen_type_sig(impl_config.types) -%}
41
{% for s in impl_config.schedules %}
42
-extern torch::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs);
42
+extern torch::stable::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs);
43
{%- endfor %}
44
45
-torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
45
+torch::stable::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
46
[[maybe_unused]] auto M = args.A.size(0);
47
[[maybe_unused]] auto N = args.B.size(1);
48
[[maybe_unused]] auto K = args.A.size(1);
@@ -59,14 +59,14 @@ torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) {
59
if (*args.maybe_schedule == "{{ gen_sch_sig(s) }}")
60
return impl_{{type_sig}}_sch_{{ gen_sch_sig(s) }}(args);
61
{%- endfor %}
62
- TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for "
62
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for "
63
"schedule = ", *args.maybe_schedule);
64
}
65
{%- endfor %}
66
67
68
-static inline std::optional<at::ScalarType> maybe_scalartype(
69
- std::optional<at::Tensor> const& t) {
68
+static inline std::optional<torch::headeronly::ScalarType> maybe_scalartype(
69
+ std::optional<torch::stable::Tensor> const& t) {
70
if (!t) {
71
return std::nullopt;
72
} else {
@@ -74,7 +74,7 @@ static inline std::optional<at::ScalarType> maybe_scalartype(
74
};
75
}
76
77
-torch::Tensor mm_dispatch(MMArgs args) {
77
+torch::stable::Tensor mm_dispatch(MMArgs args) {
78
auto out_type = args.maybe_out_type.value_or(args.A.scalar_type());
79
auto a_type = args.A.scalar_type();
80
auto maybe_g_scales_type = maybe_scalartype(args.maybe_group_scales);
@@ -105,19 +105,19 @@ torch::Tensor mm_dispatch(MMArgs args) {
105
}
106
{%- endfor %}
107
108
- TORCH_CHECK_NOT_IMPLEMENTED(
108
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
109
false, "machete_mm(..) is not implemented for "
110
- "a_type=", args.A.scalar_type(),
110
+ "a_type=", torch::headeronly::toString(args.A.scalar_type()),
111
", b_type=", args.b_type.str(),
112
- ", out_type=", out_type,
112
+ ", out_type=", torch::headeronly::toString(out_type),
113
", with_group_scale_type=", maybe_g_scales_type
114
- ? toString(*maybe_g_scales_type) : "None",
114
+ ? torch::headeronly::toString(*maybe_g_scales_type) : "None",
115
", with_group_zeropoint_type=", maybe_g_zeros_type
116
- ? toString(*maybe_g_zeros_type) : "None",
116
+ ? torch::headeronly::toString(*maybe_g_zeros_type) : "None",
117
", with_channel_scale_type=", maybe_ch_scales_type
118
- ? toString(*maybe_ch_scales_type) : "None",
118
+ ? torch::headeronly::toString(*maybe_ch_scales_type) : "None",
119
", with_token_scale_type=", maybe_tok_scales_type
120
- ? toString(*maybe_tok_scales_type) : "None",
120
+ ? torch::headeronly::toString(*maybe_tok_scales_type) : "None",
121
"; implemented types are: \\n",
122
{%- for impl_config in impl_configs %}
123
{% set t = impl_config.types -%}
@@ -197,7 +197,7 @@ using Kernel_{{type_sig}} = MacheteKernelTemplate<
197
198
{% for sch in schs %}
199
{% set sch_sig = gen_sch_sig(sch) -%}
200
-torch::Tensor
200
+torch::stable::Tensor
201
impl_{{type_sig}}_sch_{{sch_sig}}(MMArgs args) {
202
return run_impl<Kernel_{{type_sig}}<sch_{{sch_sig}}>>(args);
203
}
@@ -212,7 +212,7 @@ PREPACK_TEMPLATE = """
212
213
namespace machete {
214
215
-torch::Tensor prepack_B_dispatch(PrepackBArgs args) {
215
+torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args) {
216
auto convert_type = args.maybe_group_scales_type.value_or(args.a_type);
217
{%- for t in types %}
218
{% set b_type = unsigned_type_with_bitwidth(t.b_num_bits) %}
@@ -231,12 +231,12 @@ torch::Tensor prepack_B_dispatch(PrepackBArgs args) {
231
}
232
{%- endfor %}
233
234
- TORCH_CHECK_NOT_IMPLEMENTED(false,
234
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
235
"prepack_B_dispatch(..) is not implemented for "
236
- "atype = ", args.a_type,
236
+ "atype = ", torch::headeronly::toString(args.a_type),
237
", b_type = ", args.b_type.str(),
238
", with_group_scales_type= ", args.maybe_group_scales_type ?
239
- toString(*args.maybe_group_scales_type) : "None");
239
+ torch::headeronly::toString(*args.maybe_group_scales_type) : "None");
240
}
241
242
}; // namespace machete
csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh
renamed
csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh
renamed
csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh
renamed
csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh
renamed
+30
-27
@@ -1,8 +1,6 @@
1
#pragma once
2
3
-#include <ATen/cuda/CUDAContext.h>
4
-#include <c10/cuda/CUDAGuard.h>
5
-#include <torch/all.h>
3
+#include <torch/headeronly/util/Exception.h>
4
5
// clang-format off
6
// The cutlass include order matters (annoyingly)
@@ -175,19 +173,23 @@ struct MacheteKernelTemplate {
173
174
static Arguments create_arguments(
175
cudaStream_t stream,
178
- torch::Tensor const& A, // MxK matrix
179
- torch::Tensor const& B, // KxN prepacked matrix
180
- torch::Tensor& D, // MxN matrix
181
- std::optional<torch::Tensor> const& maybe_g_scales, // scale_KxN matrix
182
- std::optional<torch::Tensor> const& maybe_g_zeros, // scale_KxN matrix
176
+ torch::stable::Tensor const& A, // MxK matrix
177
+ torch::stable::Tensor const& B, // KxN prepacked matrix
178
+ torch::stable::Tensor& D, // MxN matrix
179
+ std::optional<torch::stable::Tensor> const&
180
+ maybe_g_scales, // scale_KxN matrix
181
+ std::optional<torch::stable::Tensor> const&
182
+ maybe_g_zeros, // scale_KxN matrix
183
std::optional<int64_t> maybe_group_size,
184
- std::optional<torch::Tensor> const& maybe_ch_scales, // len N vector
185
- std::optional<torch::Tensor> const& maybe_tok_scales) // len M vector
184
+ std::optional<torch::stable::Tensor> const&
185
+ maybe_ch_scales, // len N vector
186
+ std::optional<torch::stable::Tensor> const&
187
+ maybe_tok_scales) // len M vector
188
{
189
static_assert(!with_group_zeropoints || with_group_scales);
190
191
int M = A.size(0), N = B.size(1), K = A.size(1);
190
- TORCH_CHECK(D.size(0) == M && D.size(1) == N);
192
+ STD_TORCH_CHECK(D.size(0) == M && D.size(1) == N);
193
194
auto layout_A = make_cute_layout<StrideA>(A, "A");
195
auto layout_D = make_cute_layout<StrideD>(D, "D");
@@ -216,29 +218,29 @@ struct MacheteKernelTemplate {
218
maybe_group_size == -1 ? K : maybe_group_size.value_or(K);
219
int const scale_k = (K + group_size - 1) / group_size;
220
219
- TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K);
220
- TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N);
221
+ STD_TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K);
222
+ STD_TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N);
223
224
if constexpr (with_group_scales) {
223
- TORCH_CHECK(S_group_ptr && layout_S_group);
224
- TORCH_CHECK((size<0>(*layout_S_group) == scale_k &&
225
- size<1>(*layout_S_group) == N));
225
+ STD_TORCH_CHECK(S_group_ptr && layout_S_group);
226
+ STD_TORCH_CHECK((size<0>(*layout_S_group) == scale_k &&
227
+ size<1>(*layout_S_group) == N));
228
} else {
227
- TORCH_CHECK(!S_group_ptr, "Scales not supported");
229
+ STD_TORCH_CHECK(!S_group_ptr, "Scales not supported");
230
}
231
232
if constexpr (with_group_zeropoints) {
231
- TORCH_CHECK(Z_group_ptr && layout_Z_group);
232
- TORCH_CHECK((size<0>(*layout_Z_group) == scale_k &&
233
- size<1>(*layout_Z_group) == N));
234
- TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group,
235
- "Scales and zeros must have the same layout");
233
+ STD_TORCH_CHECK(Z_group_ptr && layout_Z_group);
234
+ STD_TORCH_CHECK((size<0>(*layout_Z_group) == scale_k &&
235
+ size<1>(*layout_Z_group) == N));
236
+ STD_TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group,
237
+ "Scales and zeros must have the same layout");
238
} else {
237
- TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported");
239
+ STD_TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported");
240
}
241
242
if constexpr (with_channel_scales || with_token_scales) {
241
- TORCH_CHECK(
243
+ STD_TORCH_CHECK(
244
(maybe_ch_scales->numel() == N || maybe_ch_scales->numel() == 1) &&
245
(maybe_tok_scales->numel() == M || maybe_tok_scales->numel() == 1));
246
}
@@ -298,11 +300,12 @@ struct MacheteKernelTemplate {
300
Gemm gemm_op;
301
302
cutlass::Status status = gemm_op.initialize(args, workspace, stream);
301
- TORCH_CHECK(status == cutlass::Status::kSuccess,
302
- "Machete kernel failed to initialize workspace");
303
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
304
+ "Machete kernel failed to initialize workspace");
305
306
status = gemm_op.run(stream);
305
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Machete kernel failed");
307
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
308
+ "Machete kernel failed");
309
}
310
};
311
csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh
new
+80
@@ -0,0 +1,80 @@
1
+#pragma once
2
+
3
+#include "machete_mm_kernel.cuh"
4
+#include "cutlass_extensions/torch_utils.hpp"
5
+#include "core/scalar_type.hpp"
6
+#include "libtorch_stable/torch_utils.h"
7
+
8
+#include <torch/csrc/stable/tensor.h>
9
+#include <torch/headeronly/core/ScalarType.h>
10
+
11
+#include <optional>
12
+#include <string>
13
+#include <vector>
14
+
15
+namespace machete {
16
+
17
+struct MMArgs {
18
+ torch::stable::Tensor const& A;
19
+ torch::stable::Tensor const& B;
20
+ vllm::ScalarType const& b_type;
21
+ std::optional<torch::headeronly::ScalarType> const& maybe_out_type;
22
+ std::optional<torch::stable::Tensor> const& maybe_group_scales;
23
+ std::optional<torch::stable::Tensor> const& maybe_group_zeros;
24
+ std::optional<int64_t> maybe_group_size;
25
+ std::optional<torch::stable::Tensor> const& maybe_channel_scales;
26
+ std::optional<torch::stable::Tensor> const& maybe_token_scales;
27
+ std::optional<std::string> maybe_schedule;
28
+};
29
+
30
+struct SupportedSchedulesArgs {
31
+ torch::headeronly::ScalarType a_type;
32
+ vllm::ScalarType b_type;
33
+ std::optional<torch::headeronly::ScalarType> maybe_group_scales_type;
34
+ std::optional<torch::headeronly::ScalarType> maybe_group_zeros_type;
35
+ std::optional<torch::headeronly::ScalarType> maybe_channel_scales_type;
36
+ std::optional<torch::headeronly::ScalarType> maybe_token_scales_type;
37
+ std::optional<torch::headeronly::ScalarType> maybe_out_type;
38
+};
39
+
40
+torch::stable::Tensor mm_dispatch(MMArgs args);
41
+
42
+std::vector<std::string> supported_schedules_dispatch(
43
+ SupportedSchedulesArgs args);
44
+
45
+template <typename MacheteKernel>
46
+torch::stable::Tensor run_impl(MMArgs args) {
47
+ const torch::stable::accelerator::DeviceGuard device_guard(
48
+ args.A.get_device_index());
49
+
50
+ auto device = args.A.device();
51
+ auto stream = get_current_cuda_stream(device.index());
52
+
53
+ int M = args.A.size(0);
54
+ int N = args.B.size(1);
55
+ int K = args.A.size(1);
56
+
57
+ // Allocate output
58
+ torch::stable::Tensor D = torch::stable::empty(
59
+ {M, N}, equivalent_scalar_type_v<typename MacheteKernel::ElementD>,
60
+ std::nullopt, device);
61
+
62
+ auto arguments = MacheteKernel::create_arguments(
63
+ stream, //
64
+ args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros,
65
+ args.maybe_group_size, args.maybe_channel_scales,
66
+ args.maybe_token_scales);
67
+ STD_TORCH_CHECK(MacheteKernel::can_implement(arguments),
68
+ "Machete kernel cannot be run with these arguments");
69
+
70
+ size_t workspace_size = MacheteKernel::get_workspace_size(arguments);
71
+ torch::stable::Tensor workspace =
72
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
73
+ std::nullopt, device);
74
+
75
+ MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream);
76
+
77
+ return D;
78
+};
79
+
80
+}; // namespace machete
\ No newline at end of file
csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh
renamed
+3
-2
@@ -3,6 +3,7 @@
3
#include "machete_mm_kernel.cuh"
4
#include "cutlass_extensions/cute_utils.cuh"
5
#include "cutlass_extensions/torch_utils.hpp"
6
+#include <torch/headeronly/util/Exception.h>
7
8
namespace machete {
9
@@ -60,8 +61,8 @@ static void prepack_B_template(
61
auto ilvd_NKbNbKL_to_offset =
62
PrepackedLayoutB::ilvd_NKbNbKL_to_offset(shape(B_layout));
63
63
- TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0);
64
- TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0);
64
+ STD_TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0);
65
+ STD_TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0);
66
67
auto N_tiles = size<0>(B_layout) / size<0>(TileShapeNKL{});
68
auto K_tiles = size<1>(B_layout) / size<1>(TileShapeNKL{});
csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh
renamed
+24
-14
@@ -3,39 +3,47 @@
3
#include "machete_prepack_kernel.cuh"
4
#include "cutlass_extensions/torch_utils.hpp"
5
#include "core/scalar_type.hpp"
6
+#include "libtorch_stable/torch_utils.h"
7
+
8
+#include <torch/csrc/stable/tensor.h>
9
+#include <torch/headeronly/core/ScalarType.h>
10
+
11
+#include <optional>
12
13
namespace machete {
14
15
struct PrepackBArgs {
10
- torch::Tensor const& B;
11
- at::ScalarType a_type;
16
+ torch::stable::Tensor const& B;
17
+ torch::headeronly::ScalarType a_type;
18
vllm::ScalarType b_type;
13
- std::optional<at::ScalarType> maybe_group_scales_type;
19
+ std::optional<torch::headeronly::ScalarType> maybe_group_scales_type;
20
};
21
22
template <typename PrepackedLayoutB>
17
-torch::Tensor prepack_impl(torch::Tensor const B) {
18
- const at::cuda::OptionalCUDAGuard device_guard(device_of(B));
23
+torch::stable::Tensor prepack_impl(torch::stable::Tensor const& B) {
24
+ const torch::stable::accelerator::DeviceGuard device_guard(
25
+ B.get_device_index());
26
using ElementB = typename PrepackedLayoutB::ElementB;
27
using PPBlockShape_NK = typename PrepackedLayoutB::PPBlockShape_NK;
28
29
auto device = B.device();
23
- auto stream = at::cuda::getCurrentCUDAStream(device.index());
30
+ auto stream = get_current_cuda_stream(device.index());
31
auto B_ptr = static_cast<ElementB const*>(B.const_data_ptr());
32
// elements per storage item for B
33
auto eles_per_storage =
27
- (B.dtype().itemsize() * 8) / cute::sizeof_bits_v<ElementB>;
34
+ (B.element_size() * 8) / cute::sizeof_bits_v<ElementB>;
35
36
// torch B passed in is/should be (packed_K,N), the kernel expects (N,K,L) (to
37
// match cutlass using (N,K,L) for B), so we transpose B to (N,packed_K,L)
31
- auto Bt_packed = B.t();
38
+ auto Bt_packed = torch::stable::transpose(B, 0, 1);
39
33
- TORCH_CHECK(
40
+ STD_TORCH_CHECK(
41
(B.size(0) * eles_per_storage) % size<1>(PPBlockShape_NK{}) == 0,
42
"B.shape[0] (in terms of unpacked elements) must be a multiple of ",
43
size<1>(PPBlockShape_NK{}));
37
- TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0,
38
- "B.shape[1] must be a multiple of ", size<0>(PPBlockShape_NK{}));
44
+ STD_TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0,
45
+ "B.shape[1] must be a multiple of ",
46
+ size<0>(PPBlockShape_NK{}));
47
48
using StrideB = cutlass::detail::TagToStrideB_t<cutlass::layout::ColumnMajor>;
49
auto const l_Bt_packed = make_cute_layout<StrideB>(Bt_packed, "B");
@@ -49,7 +57,7 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
57
// new_shape = (N, packed_K, L) * (1, eles_per_storage, 1) -> (N, K, L)
58
// new_stride = (s0, s1, s2) * (eles_per_storage, 1, eles_per_storage)
59
// when s1 == 1
52
- TORCH_CHECK(stride<1>(l_Bt_packed) == 1);
60
+ STD_TORCH_CHECK(stride<1>(l_Bt_packed) == 1);
61
// clang-format off
62
auto const layout_Bt = make_layout(
63
transform_with_idx(l_Bt_packed.shape(), [&](auto ele, auto idx) {
@@ -61,7 +69,9 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
69
// clang-format on
70
71
// Allocate output
64
- torch::Tensor D = torch::empty_like(B, {}, at::MemoryFormat::Contiguous);
72
+ torch::stable::Tensor D = torch::stable::empty(
73
+ B.sizes(), B.scalar_type(), std::nullopt, B.device(), std::nullopt,
74
+ torch::headeronly::MemoryFormat::Contiguous);
75
76
prepack_B_template<PrepackedLayoutB>(
77
stream, B_ptr, layout_Bt, static_cast<ElementB*>(D.mutable_data_ptr()));
@@ -69,6 +79,6 @@ torch::Tensor prepack_impl(torch::Tensor const B) {
79
return D;
80
};
81
72
-torch::Tensor prepack_B_dispatch(PrepackBArgs args);
82
+torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args);
83
84
}; // namespace machete
\ No newline at end of file
csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh
renamed
-4
@@ -1,9 +1,5 @@
1
#pragma once
2
3
-#include <ATen/cuda/CUDAContext.h>
4
-#include <c10/cuda/CUDAGuard.h>
5
-#include <torch/all.h>
6
-
3
// clang-format off
4
// The cutlass include order matters (annoyingly)
5
csrc/libtorch_stable/quantization/machete/machete_pytorch.cu
new
+77
@@ -0,0 +1,77 @@
1
+#include "machete_mm_launcher.cuh"
2
+#include "machete_prepack_launcher.cuh"
3
+#include "core/scalar_type.hpp"
4
+
5
+#include <torch/csrc/stable/library.h>
6
+#include <torch/csrc/stable/tensor.h>
7
+#include <torch/headeronly/core/ScalarType.h>
8
+
9
+namespace machete {
10
+
11
+using namespace vllm;
12
+
13
+std::vector<std::string> supported_schedules(
14
+ torch::headeronly::ScalarType a_type, int64_t b_type_id,
15
+ std::optional<torch::headeronly::ScalarType> maybe_group_scales_type,
16
+ std::optional<torch::headeronly::ScalarType> maybe_group_zeros_type,
17
+ std::optional<torch::headeronly::ScalarType> maybe_channel_scales_type,
18
+ std::optional<torch::headeronly::ScalarType> maybe_token_scales_type,
19
+ std::optional<torch::headeronly::ScalarType> maybe_out_type) {
20
+ ScalarType const b_type = ScalarType::from_id(b_type_id);
21
+ return supported_schedules_dispatch({
22
+ .a_type = a_type,
23
+ .b_type = b_type,
24
+ .maybe_group_scales_type = maybe_group_scales_type,
25
+ .maybe_group_zeros_type = maybe_group_zeros_type,
26
+ .maybe_channel_scales_type = maybe_channel_scales_type,
27
+ .maybe_token_scales_type = maybe_token_scales_type,
28
+ .maybe_out_type = maybe_out_type,
29
+ });
30
+}
31
+
32
+torch::stable::Tensor mm(
33
+ torch::stable::Tensor const& A, torch::stable::Tensor const& B,
34
+ int64_t b_type_id,
35
+ std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
36
+ std::optional<torch::stable::Tensor> const& maybe_group_scales,
37
+ std::optional<torch::stable::Tensor> const& maybe_group_zeros,
38
+ std::optional<int64_t> maybe_group_size,
39
+ std::optional<torch::stable::Tensor> const& maybe_channel_scales,
40
+ std::optional<torch::stable::Tensor> const& maybe_token_scales,
41
+ std::optional<std::string> maybe_schedule) {
42
+ ScalarType const b_type = ScalarType::from_id(b_type_id);
43
+ return mm_dispatch({.A = A,
44
+ .B = B,
45
+ .b_type = b_type,
46
+ .maybe_out_type = maybe_out_type,
47
+ .maybe_group_scales = maybe_group_scales,
48
+ .maybe_group_zeros = maybe_group_zeros,
49
+ .maybe_group_size = maybe_group_size,
50
+ .maybe_channel_scales = maybe_channel_scales,
51
+ .maybe_token_scales = maybe_token_scales,
52
+ .maybe_schedule = maybe_schedule});
53
+}
54
+
55
+torch::stable::Tensor prepack_B(
56
+ torch::stable::Tensor const& B, torch::headeronly::ScalarType const& a_type,
57
+ int64_t b_type_id,
58
+ std::optional<torch::headeronly::ScalarType> const&
59
+ maybe_group_scales_type) {
60
+ ScalarType const b_type = ScalarType::from_id(b_type_id);
61
+ return prepack_B_dispatch(
62
+ {.B = B,
63
+ .a_type = a_type,
64
+ .b_type = b_type,
65
+ .maybe_group_scales_type = maybe_group_scales_type});
66
+}
67
+
68
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
69
+ m.impl("machete_prepack_B", TORCH_BOX(&prepack_B));
70
+ m.impl("machete_mm", TORCH_BOX(&mm));
71
+}
72
+
73
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
74
+ m.impl("machete_supported_schedules", TORCH_BOX(&supported_schedules));
75
+}
76
+
77
+}; // namespace machete
csrc/libtorch_stable/torch_bindings.cpp
+33
@@ -34,6 +34,39 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
34
// TODO: Remove this once ROCm upgrade to torch 2.11.
35
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
36
37
+ // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
38
+ ops.def(
39
+ "machete_supported_schedules("
40
+ " ScalarType a_type,"
41
+ " int b_type,"
42
+ " ScalarType? maybe_group_scales_type,"
43
+ " ScalarType? maybe_group_zeros_type,"
44
+ " ScalarType? maybe_channel_scales_type,"
45
+ " ScalarType? maybe_token_scales_type,"
46
+ " ScalarType? maybe_out_type"
47
+ ") -> str[]");
48
+ ops.def(
49
+ "machete_mm("
50
+ " Tensor A,"
51
+ " Tensor B,"
52
+ " int b_type,"
53
+ " ScalarType? out_type,"
54
+ " Tensor? group_scales,"
55
+ " Tensor? group_zeros,"
56
+ " int? group_size,"
57
+ " Tensor? channel_scales,"
58
+ " Tensor? token_scales,"
59
+ " str? schedule"
60
+ ") -> Tensor");
61
+ ops.def(
62
+ "machete_prepack_B("
63
+ " Tensor B,"
64
+ " ScalarType a_type,"
65
+ " int b_type,"
66
+ " ScalarType? group_scales_type"
67
+ ") -> Tensor");
68
+ // conditionally compiled so impl registration is in source file
69
+
70
// Marlin GEMM
71
ops.def(
72
"marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
csrc/quantization/machete/machete_mm_launcher.cuh
deleted
-75
@@ -1,75 +0,0 @@
1
-#pragma once
2
-
3
-#include <torch/all.h>
4
-#include <Python.h>
5
-
6
-#include "machete_mm_kernel.cuh"
7
-#include "cutlass_extensions/torch_utils.hpp"
8
-#include "core/scalar_type.hpp"
9
-
10
-namespace machete {
11
-
12
-struct MMArgs {
13
- torch::Tensor const& A;
14
- torch::Tensor const& B;
15
- vllm::ScalarType const& b_type;
16
- std::optional<at::ScalarType> const& maybe_out_type;
17
- std::optional<torch::Tensor> const& maybe_group_scales;
18
- std::optional<torch::Tensor> const& maybe_group_zeros;
19
- std::optional<int64_t> maybe_group_size;
20
- std::optional<torch::Tensor> const& maybe_channel_scales;
21
- std::optional<torch::Tensor> const& maybe_token_scales;
22
- std::optional<std::string> maybe_schedule;
23
-};
24
-
25
-struct SupportedSchedulesArgs {
26
- at::ScalarType a_type;
27
- vllm::ScalarType b_type;
28
- std::optional<at::ScalarType> maybe_group_scales_type;
29
- std::optional<at::ScalarType> maybe_group_zeros_type;
30
- std::optional<at::ScalarType> maybe_channel_scales_type;
31
- std::optional<at::ScalarType> maybe_token_scales_type;
32
- std::optional<at::ScalarType> maybe_out_type;
33
-};
34
-
35
-torch::Tensor mm_dispatch(MMArgs args);
36
-
37
-std::vector<std::string> supported_schedules_dispatch(
38
- SupportedSchedulesArgs args);
39
-
40
-template <typename MacheteKernel>
41
-torch::Tensor run_impl(MMArgs args) {
42
- const at::cuda::OptionalCUDAGuard device_guard(device_of(args.A));
43
-
44
- auto device = args.A.device();
45
- auto stream = at::cuda::getCurrentCUDAStream(device.index());
46
-
47
- int M = args.A.size(0);
48
- int N = args.B.size(1);
49
- int K = args.A.size(1);
50
-
51
- // Allocate output
52
- torch::Tensor D = torch::empty(
53
- {M, N},
54
- torch::TensorOptions()
55
- .dtype(equivalent_scalar_type_v<typename MacheteKernel::ElementD>)
56
- .device(device));
57
-
58
- auto arguments = MacheteKernel::create_arguments(
59
- stream, //
60
- args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros,
61
- args.maybe_group_size, args.maybe_channel_scales,
62
- args.maybe_token_scales);
63
- TORCH_CHECK(MacheteKernel::can_implement(arguments),
64
- "Machete kernel cannot be run with these arguments");
65
-
66
- size_t workspace_size = MacheteKernel::get_workspace_size(arguments);
67
- torch::Tensor workspace = torch::empty(
68
- workspace_size, torch::TensorOptions().dtype(torch::kU8).device(device));
69
-
70
- MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream);
71
-
72
- return D;
73
-};
74
-
75
-}; // namespace machete
\ No newline at end of file
csrc/quantization/machete/machete_pytorch.cu
deleted
-73
@@ -1,73 +0,0 @@
1
-#include "machete_mm_launcher.cuh"
2
-#include "machete_prepack_launcher.cuh"
3
-#include "core/scalar_type.hpp"
4
-
5
-#include "core/registration.h"
6
-
7
-namespace machete {
8
-
9
-using namespace vllm;
10
-
11
-std::vector<std::string> supported_schedules(
12
- at::ScalarType a_type, int64_t b_type_id,
13
- std::optional<at::ScalarType> maybe_group_scales_type,
14
- std::optional<at::ScalarType> maybe_group_zeros_type,
15
- std::optional<at::ScalarType> maybe_channel_scales_type,
16
- std::optional<at::ScalarType> maybe_token_scales_type,
17
- std::optional<at::ScalarType> maybe_out_type) {
18
- ScalarType const b_type = ScalarType::from_id(b_type_id);
19
- return supported_schedules_dispatch({
20
- .a_type = a_type,
21
- .b_type = b_type,
22
- .maybe_group_scales_type = maybe_group_scales_type,
23
- .maybe_group_zeros_type = maybe_group_zeros_type,
24
- .maybe_channel_scales_type = maybe_channel_scales_type,
25
- .maybe_token_scales_type = maybe_token_scales_type,
26
- .maybe_out_type = maybe_out_type,
27
- });
28
-}
29
-
30
-torch::Tensor mm(torch::Tensor const& A, torch::Tensor const& B,
31
- int64_t b_type_id,
32
- std::optional<at::ScalarType> const& maybe_out_type,
33
- std::optional<torch::Tensor> const& maybe_group_scales,
34
- std::optional<torch::Tensor> const& maybe_group_zeros,
35
- std::optional<int64_t> maybe_group_size,
36
- std::optional<torch::Tensor> const& maybe_channel_scales,
37
- std::optional<torch::Tensor> const& maybe_token_scales,
38
- std::optional<std::string> maybe_schedule) {
39
- ScalarType const b_type = ScalarType::from_id(b_type_id);
40
- return mm_dispatch({.A = A,
41
- .B = B,
42
- .b_type = b_type,
43
- .maybe_out_type = maybe_out_type,
44
- .maybe_group_scales = maybe_group_scales,
45
- .maybe_group_zeros = maybe_group_zeros,
46
- .maybe_group_size = maybe_group_size,
47
- .maybe_channel_scales = maybe_channel_scales,
48
- .maybe_token_scales = maybe_token_scales,
49
- .maybe_schedule = maybe_schedule});
50
-}
51
-
52
-torch::Tensor prepack_B(
53
- torch::Tensor const& B, at::ScalarType const& a_type, int64_t b_type_id,
54
- std::optional<at::ScalarType> const& maybe_group_scales_type) {
55
- ScalarType const b_type = ScalarType::from_id(b_type_id);
56
- return prepack_B_dispatch(
57
- {.B = B,
58
- .a_type = a_type,
59
- .b_type = b_type,
60
- .maybe_group_scales_type = maybe_group_scales_type});
61
-}
62
-
63
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
64
- m.impl("machete_prepack_B", &prepack_B);
65
- m.impl("machete_mm", &mm);
66
-}
67
-
68
-// use CatchAll since supported_schedules has no tensor arguments
69
-TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) {
70
- m.impl("machete_supported_schedules", &supported_schedules);
71
-}
72
-
73
-}; // namespace machete
csrc/torch_bindings.cpp
-33
@@ -68,39 +68,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
68
// custom types:
69
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
70
71
- // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
72
- ops.def(
73
- "machete_supported_schedules("
74
- " ScalarType a_type,"
75
- " int b_type,"
76
- " ScalarType? maybe_group_scales_type,"
77
- " ScalarType? maybe_group_zeros_type,"
78
- " ScalarType? maybe_channel_scales_type,"
79
- " ScalarType? maybe_token_scales_type,"
80
- " ScalarType? maybe_out_type"
81
- ") -> str[]");
82
- ops.def(
83
- "machete_mm("
84
- " Tensor A,"
85
- " Tensor B,"
86
- " int b_type,"
87
- " ScalarType? out_type,"
88
- " Tensor? group_scales,"
89
- " Tensor? group_zeros,"
90
- " int? group_size,"
91
- " Tensor? channel_scales,"
92
- " Tensor? token_scales,"
93
- " str? schedule"
94
- ") -> Tensor");
95
- ops.def(
96
- "machete_prepack_B("
97
- " Tensor B,"
98
- " ScalarType a_type,"
99
- " int b_type,"
100
- " ScalarType? group_scales_type"
101
- ") -> Tensor");
102
- // conditionally compiled so impl registration is in source file
103
-
71
#endif
72
}
73