@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 56aff0dd15
Fork 0

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 1, 22:30 UTC

[10/n] Migrate cuda_view and silu_and_mul_per_block_quant kernels to torch stale ABI. (#44334)

Chris Leonard committed Jun 4, 2026 at 23:14 UTC 56aff0dd15c0d01abc3f1de2b773aff2ed61d7a0
25 files changed +187 -164
CMakeLists.txt
+7 -11
@@ -307,10 +307,7 @@ endif()
307 #
308
309 set(VLLM_EXT_SRC
310 - "csrc/cuda_view.cu"
311 - "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
310 "csrc/quantization/activation_kernels.cu"
313 - "csrc/cuda_utils_kernels.cu"
311 "csrc/torch_bindings.cpp")
312
313 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -346,9 +343,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
343 endif()
344 FetchContent_MakeAvailable(cutlass)
345
349 - list(APPEND VLLM_EXT_SRC
350 - "csrc/cutlass_extensions/common.cpp")
351 -
346 set_gencode_flags_for_srcs(
347 SRCS "${VLLM_EXT_SRC}"
348 CUDA_ARCHS "${CUDA_ARCHS}")
@@ -627,6 +621,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
621 #
622 set(VLLM_STABLE_EXT_SRC
623 "csrc/libtorch_stable/torch_bindings.cpp"
624 + "csrc/libtorch_stable/cuda_view.cu"
625 "csrc/libtorch_stable/activation_kernels.cu"
626 "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
627 "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
@@ -639,6 +634,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
634 "csrc/libtorch_stable/layernorm_kernels.cu"
635 "csrc/libtorch_stable/layernorm_quant_kernels.cu"
636 "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
637 + "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
638 "csrc/libtorch_stable/attention/merge_attn_states.cu"
639 "csrc/libtorch_stable/sampler.cu"
640 "csrc/libtorch_stable/topk.cu"
@@ -653,8 +649,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
649
650 if(VLLM_GPU_LANG STREQUAL "CUDA")
651 list(APPEND VLLM_STABLE_EXT_SRC
656 - "csrc/cuda_utils_kernels.cu"
657 - "csrc/cutlass_extensions/common.cpp"
652 + "csrc/libtorch_stable/cuda_utils_kernels.cu"
653 + "csrc/libtorch_stable/cutlass_extensions/common.cpp"
654 "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
655 "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
656 "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
@@ -1076,11 +1072,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1072 WITH_SOABI)
1073
1074 # Set TORCH_TARGET_VERSION for stable ABI compatibility.
1079 - # This ensures we only use C-shim APIs available in PyTorch 2.10.
1075 + # This ensures we only use C-shim APIs available in PyTorch 2.11.
1076 # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
1081 - # which is currently set to 2.10.
1077 + # which is currently set to 2.11.
1078 target_compile_definitions(_C_stable_libtorch PRIVATE
1083 - TORCH_TARGET_VERSION=0x020A000000000000ULL)
1079 + TORCH_TARGET_VERSION=0x020B000000000000ULL)
1080
1081 # Needed to use cuda/hip APIs from C-shim
1082 if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/cuda_view.cu deleted
-59
@@ -1,59 +0,0 @@
1 -#include <torch/all.h>
2 -#include <torch/cuda.h>
3 -#include <cuda_runtime.h>
4 -
5 -// This function assumes that `cpu_tensor` is a CPU tensor,
6 -// and that UVA (Unified Virtual Addressing) is enabled.
7 -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
8 - TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
9 -
10 - // handle empty tensor
11 - if (cpu_tensor.numel() == 0) {
12 - return torch::empty(cpu_tensor.sizes(),
13 - cpu_tensor.options().device(torch::kCUDA));
14 - }
15 -
16 - if (cpu_tensor.is_pinned()) {
17 - // If CPU tensor is pinned, directly get the device pointer.
18 - void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
19 - void* device_ptr = nullptr;
20 - cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
21 - TORCH_CHECK(err == cudaSuccess,
22 - "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
23 -
24 - return torch::from_blob(
25 - device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
26 - [base = cpu_tensor](void*) {}, // keep cpu tensor alive
27 - cpu_tensor.options().device(torch::kCUDA));
28 - }
29 -
30 - // If CPU tensor is not pinned, allocate a new pinned memory buffer.
31 - torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
32 - size_t nbytes = contiguous_cpu.nbytes();
33 -
34 - void* host_ptr = nullptr;
35 - cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
36 - if (err != cudaSuccess) {
37 - AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
38 - }
39 -
40 - err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
41 - cudaMemcpyDefault);
42 - if (err != cudaSuccess) {
43 - cudaFreeHost(host_ptr);
44 - AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
45 - }
46 -
47 - void* device_ptr = nullptr;
48 - err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
49 - if (err != cudaSuccess) {
50 - cudaFreeHost(host_ptr);
51 - AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
52 - }
53 -
54 - auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
55 -
56 - return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
57 - contiguous_cpu.strides(), deleter,
58 - contiguous_cpu.options().device(torch::kCUDA));
59 -}
\ No newline at end of file
csrc/libtorch_stable/cuda_utils_kernels.cu renamed
csrc/libtorch_stable/cuda_view.cu new
+76
@@ -0,0 +1,76 @@
1 +#include <torch/csrc/stable/tensor.h>
2 +#include <torch/csrc/stable/ops.h>
3 +#include <torch/csrc/stable/accelerator.h>
4 +#include <torch/headeronly/core/ScalarType.h>
5 +#include <torch/csrc/stable/device.h>
6 +#include <torch/csrc/stable/c/shim.h>
7 +#include <torch/headeronly/version.h>
8 +#include <cuda_runtime.h>
9 +
10 +#include <array>
11 +#include <optional>
12 +
13 +// This function assumes that `cpu_tensor` is a CPU tensor,
14 +// and that UVA (Unified Virtual Addressing) is enabled.
15 +torch::stable::Tensor get_cuda_view_from_cpu_tensor(
16 + torch::stable::Tensor& cpu_tensor) {
17 + STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
18 +
19 + const auto dtype = cpu_tensor.scalar_type();
20 + const auto layout = cpu_tensor.layout();
21 + const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA);
22 +
23 + // handle empty tensor
24 + if (cpu_tensor.numel() == 0) {
25 + return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev);
26 + }
27 +
28 + std::array<StableIValue, 2> is_pinned_stack{
29 + torch::stable::detail::from(cpu_tensor),
30 + torch::stable::detail::from(std::nullopt)};
31 + TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
32 + "aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
33 + if (torch::stable::detail::to<bool>(is_pinned_stack[0])) {
34 + // If CPU tensor is pinned, directly get the device pointer.
35 + void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr());
36 + void* device_ptr = nullptr;
37 + cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
38 + STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ",
39 + cudaGetErrorString(err));
40 +
41 + return torch::stable::from_blob(
42 + device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype,
43 + [base = cpu_tensor](void*) {}); // keep cpu tensor alive
44 + }
45 +
46 + // If CPU tensor is not pinned, allocate a new pinned memory buffer.
47 + torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor);
48 + size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size();
49 +
50 + void* host_ptr = nullptr;
51 + cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
52 + if (err != cudaSuccess) {
53 + STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err));
54 + }
55 +
56 + err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes,
57 + cudaMemcpyDefault);
58 + if (err != cudaSuccess) {
59 + cudaFreeHost(host_ptr);
60 + STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err));
61 + }
62 +
63 + void* device_ptr = nullptr;
64 + err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
65 + if (err != cudaSuccess) {
66 + cudaFreeHost(host_ptr);
67 + STD_TORCH_CHECK(
68 + false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
69 + }
70 +
71 + auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
72 +
73 + return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(),
74 + contiguous_cpu.strides(), cuda_dev,
75 + contiguous_cpu.scalar_type(), deleter);
76 +}
csrc/libtorch_stable/cutlass_extensions/common.cpp renamed
+1 -1
@@ -1,4 +1,4 @@
1 -#include "cutlass_extensions/common.hpp"
1 +#include "common.hpp"
2
3 int32_t get_sm_version_num() {
4 int32_t major_capability, minor_capability;
csrc/libtorch_stable/cutlass_extensions/common.hpp renamed
csrc/libtorch_stable/ops.h
+11
@@ -164,6 +164,10 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
164
165 #endif
166
167 +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
168 +torch::stable::Tensor get_cuda_view_from_cpu_tensor(
169 + torch::stable::Tensor& cpu_tensor);
170 +
171 // Attention kernels (shared CUDA/ROCm)
172 void merge_attn_states(
173 torch::stable::Tensor& output,
@@ -215,6 +219,13 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
219 std::optional<torch::stable::Tensor> residual,
220 int64_t group_size, bool is_scale_transposed);
221
222 +void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
223 + torch::stable::Tensor const& input,
224 + torch::stable::Tensor& scales,
225 + int64_t group_size,
226 + std::optional<torch::stable::Tensor> scale_ub,
227 + bool is_scale_transposed);
228 +
229 // Positional encoding kernels (shared CUDA/ROCm)
230 void rotary_embedding(torch::stable::Tensor& positions,
231 torch::stable::Tensor& query,
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu
+1 -1
@@ -18,7 +18,7 @@
18 #include <torch/csrc/stable/tensor.h>
19 #include "libtorch_stable/torch_utils.h"
20 #include "cutlass_extensions/torch_utils.hpp"
21 -#include "cutlass_extensions/common.hpp"
21 +#include "libtorch_stable/cutlass_extensions/common.hpp"
22
23 #include "get_group_starts.cuh"
24 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu
+1 -1
@@ -21,7 +21,7 @@
21 #include "cutlass/util/packed_stride.hpp"
22 #include "cutlass/util/mixed_dtype_utils.hpp"
23
24 -#include "cutlass_extensions/common.hpp"
24 +#include "libtorch_stable/cutlass_extensions/common.hpp"
25 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
26
27 #include <cuda_runtime.h>
csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu
+1 -1
@@ -12,7 +12,7 @@
12
13 #include <cutlass/arch/arch.h>
14
15 -#include "cutlass_extensions/common.hpp"
15 +#include "libtorch_stable/cutlass_extensions/common.hpp"
16
17 #include "cute/tensor.hpp"
18 #include "cutlass/tensor_ref.h"
csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+1 -1
@@ -20,7 +20,7 @@
20
21 #include <cutlass/arch/arch.h>
22
23 -#include "cutlass_extensions/common.hpp"
23 +#include "libtorch_stable/cutlass_extensions/common.hpp"
24
25 #include "cute/tensor.hpp"
26 #include "cutlass/tensor_ref.h"
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu
+1 -1
@@ -18,7 +18,7 @@
18
19 #include "libtorch_stable/torch_utils.h"
20
21 -#include "cutlass_extensions/common.hpp"
21 +#include "libtorch_stable/cutlass_extensions/common.hpp"
22 #include "nvfp4_utils.cuh"
23
24 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu
+1 -1
@@ -18,7 +18,7 @@
18
19 #include "libtorch_stable/torch_utils.h"
20
21 -#include "cutlass_extensions/common.hpp"
21 +#include "libtorch_stable/cutlass_extensions/common.hpp"
22
23 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
24 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu
+1 -1
@@ -18,7 +18,7 @@
18
19 #include "libtorch_stable/torch_utils.h"
20
21 -#include "cutlass_extensions/common.hpp"
21 +#include "libtorch_stable/cutlass_extensions/common.hpp"
22
23 #include "cutlass/cutlass.h"
24
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
+1 -1
@@ -18,7 +18,7 @@
18
19 #include "libtorch_stable/torch_utils.h"
20
21 -#include "cutlass_extensions/common.hpp"
21 +#include "libtorch_stable/cutlass_extensions/common.hpp"
22
23 #include "cutlass/cutlass.h"
24
csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu renamed
+45 -40
@@ -1,11 +1,10 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -#include <ATen/cuda/CUDAContext.h>
5 -#include <c10/cuda/CUDAGuard.h>
4 +#include "../../torch_utils.h"
5
6 #include "../../dispatch_utils.h"
8 -#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
7 +#include "quant_conversions.cuh"
8
9 namespace vllm {
10
@@ -105,64 +104,70 @@ __global__ void silu_and_mul_per_block_quant_kernel(
104
105 } // namespace vllm
106
108 -void silu_and_mul_per_block_quant(torch::Tensor& out,
109 - torch::Tensor const& input,
110 - torch::Tensor& scales, int64_t group_size,
111 - std::optional<torch::Tensor> scale_ub,
107 +void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
108 + torch::stable::Tensor const& input,
109 + torch::stable::Tensor& scales,
110 + int64_t group_size,
111 + std::optional<torch::stable::Tensor> scale_ub,
112 bool is_scale_transposed) {
113 - static c10::ScalarType kFp8Type = is_fp8_ocp()
114 - ? c10::ScalarType::Float8_e4m3fn
115 - : c10::ScalarType::Float8_e4m3fnuz;
116 -
117 - TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
118 - TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
119 - TORCH_CHECK(
120 - input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
113 + static torch::headeronly::ScalarType kFp8Type =
114 + is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
115 + : torch::headeronly::ScalarType::Float8_e4m3fnuz;
116 +
117 + STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
118 + out.scalar_type() == torch::headeronly::ScalarType::Char);
119 + STD_TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
120 + STD_TORCH_CHECK(
121 + input.scalar_type() == torch::headeronly::ScalarType::Half ||
122 + input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
123 "Input must be FP16 or BF16");
122 - TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
123 - TORCH_CHECK(group_size == 128 || group_size == 64,
124 - "Unsupported group size: ", group_size);
124 + STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
125 + STD_TORCH_CHECK(group_size == 128 || group_size == 64,
126 + "Unsupported group size: ", group_size);
127
128 if (scale_ub.has_value()) {
127 - TORCH_CHECK(out.dtype() == kFp8Type);
129 + STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
130 }
131
132 int32_t hidden_size = out.size(-1);
133 auto num_tokens = input.size(0);
134 int32_t num_groups = hidden_size / group_size;
135
134 - TORCH_CHECK(input.size(-1) == hidden_size * 2,
135 - "input last dim must be 2x output hidden_size");
136 - TORCH_CHECK(hidden_size % group_size == 0,
137 - "hidden_size must be divisible by group_size");
136 + STD_TORCH_CHECK(input.size(-1) == hidden_size * 2,
137 + "input last dim must be 2x output hidden_size");
138 + STD_TORCH_CHECK(hidden_size % group_size == 0,
139 + "hidden_size must be divisible by group_size");
140
139 - const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
140 - const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
141 + const torch::stable::accelerator::DeviceGuard device_guard(
142 + input.get_device_index());
143 + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
144
145 dim3 grid(num_tokens, num_groups);
146 dim3 block(group_size);
147
145 - VLLM_DISPATCH_FLOATING_TYPES(
148 + VLLM_STABLE_DISPATCH_FLOATING_TYPES(
149 input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
150 using scalar_in_t = scalar_t;
151
149 - VLLM_DISPATCH_QUANT_TYPES(
152 + VLLM_STABLE_DISPATCH_QUANT_TYPES(
153 out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
154 using scalar_out_t = scalar_t;
155
153 - VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
154 - VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
155 - vllm::silu_and_mul_per_block_quant_kernel<
156 - scalar_in_t, scalar_out_t, transpose_scale, gs>
157 - <<<grid, block, 0, stream>>>(
158 - out.data_ptr<scalar_out_t>(),
159 - scales.data_ptr<float>(),
160 - input.data_ptr<scalar_in_t>(),
161 - scale_ub.has_value() ? scale_ub->data_ptr<float>()
162 - : nullptr,
163 - hidden_size);
164 - });
156 + VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
157 + VLLM_STABLE_DISPATCH_BOOL(
158 + is_scale_transposed, transpose_scale, [&] {
159 + vllm::silu_and_mul_per_block_quant_kernel<
160 + scalar_in_t, scalar_out_t, transpose_scale, gs>
161 + <<<grid, block, 0, stream>>>(
162 + out.mutable_data_ptr<scalar_out_t>(),
163 + scales.mutable_data_ptr<float>(),
164 + input.const_data_ptr<scalar_in_t>(),
165 + scale_ub.has_value()
166 + ? scale_ub->const_data_ptr<float>()
167 + : nullptr,
168 + hidden_size);
169 + });
170 });
171 });
172 });
168 -}
\ No newline at end of file
173 +}
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh
+1 -1
@@ -20,7 +20,7 @@
20 #include "cutlass/util/packed_stride.hpp"
21
22 #include "core/math.hpp"
23 -#include "cutlass_extensions/common.hpp"
23 +#include "libtorch_stable/cutlass_extensions/common.hpp"
24 // clang-format on
25
26 namespace vllm::c3x {
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh
+1 -1
@@ -15,7 +15,7 @@
15 #include "cutlass/gemm/collective/collective_builder.hpp"
16
17 #include "core/math.hpp"
18 -#include "cutlass_extensions/common.hpp"
18 +#include "libtorch_stable/cutlass_extensions/common.hpp"
19 // clang-format on
20
21 /*
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp
+1 -1
@@ -1,7 +1,7 @@
1 #include <torch/csrc/stable/tensor.h>
2 #include <torch/headeronly/core/ScalarType.h>
3 #include "cuda_utils.h"
4 -#include "cutlass_extensions/common.hpp"
4 +#include "libtorch_stable/cutlass_extensions/common.hpp"
5
6 template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
7 void dispatch_scaled_mm(torch::stable::Tensor& c,
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh
+1 -1
@@ -8,7 +8,7 @@
8
9 #include <torch/csrc/stable/ops.h>
10 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
11 -#include "cutlass_extensions/common.hpp"
11 +#include "libtorch_stable/cutlass_extensions/common.hpp"
12 #include "get_group_starts.cuh"
13
14 using namespace cute;
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh
+1 -1
@@ -23,7 +23,7 @@
23 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
24
25 #include "core/math.hpp"
26 -#include "cutlass_extensions/common.hpp"
26 +#include "libtorch_stable/cutlass_extensions/common.hpp"
27 // clang-format on
28
29 using namespace cute;
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu
+1 -1
@@ -4,7 +4,7 @@
4
5 #include "libtorch_stable/torch_utils.h"
6
7 -#include "cutlass_extensions/common.hpp"
7 +#include "libtorch_stable/cutlass_extensions/common.hpp"
8
9 void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
10 torch::stable::Tensor const& a,
csrc/libtorch_stable/torch_bindings.cpp
+33
@@ -1,4 +1,5 @@
1 #include "ops.h"
2 +#include "cuda_utils.h"
3 #include "core/registration.h"
4
5 #include <torch/csrc/stable/library.h>
@@ -27,6 +28,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
28 "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
29 "()");
30
31 + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
32 +
33 #ifndef USE_ROCM
34 ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
35 #endif
@@ -321,6 +324,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
324 "Tensor? scale_ub, Tensor!? residual, int group_size, "
325 "bool is_scale_transposed) -> ()");
326
327 + // Fused SiLU+Mul + per-block quantization
328 + ops.def(
329 + "silu_and_mul_per_block_quant("
330 + "Tensor! out, "
331 + "Tensor input, "
332 + "Tensor! scales, "
333 + "int group_size, "
334 + "Tensor? scale_ub=None, "
335 + "bool is_scale_transposed=False) -> ()");
336 +
337 // Rotary embedding
338 // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
339 ops.def(
@@ -599,6 +612,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
612 ops.impl("rms_norm_dynamic_per_token_quant",
613 TORCH_BOX(&rms_norm_dynamic_per_token_quant));
614 ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
615 + ops.impl("silu_and_mul_per_block_quant",
616 + TORCH_BOX(&silu_and_mul_per_block_quant));
617
618 // Positional encoding kernels (shared CUDA/ROCm)
619 ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
@@ -661,6 +676,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
676 ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
677 }
678
679 +STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
680 + ops.impl("get_cuda_view_from_cpu_tensor",
681 + TORCH_BOX(&get_cuda_view_from_cpu_tensor));
682 +}
683 +
684 // These capability-check functions take only primitive args (no tensors), so
685 // there is no device to dispatch on. CompositeExplicitAutograd makes them
686 // available for all backends. This is the stable ABI equivalent of calling
@@ -681,6 +701,19 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
701 ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
702 }
703
704 +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
705 + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
706 + cuda_utils.def(
707 + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
708 +}
709 +
710 +STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
711 + cuda_utils) {
712 + cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
713 + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
714 + TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
715 +}
716 +
717 // Cache ops
718 STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
719 // Swap in (out) the cache blocks from src to dst.
csrc/ops.h
-8
@@ -40,12 +40,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
40 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
41 torch::Tensor& weight, double epsilon);
42
43 -void silu_and_mul_per_block_quant(torch::Tensor& out,
44 - torch::Tensor const& input,
45 - torch::Tensor& scales, int64_t group_size,
46 - std::optional<torch::Tensor> scale_ub,
47 - bool is_scale_transposed);
48 -
43 // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
44 // ABI for CUDA). It remains here because the CPU build still uses these
45 // torch::Tensor declarations.
@@ -84,8 +78,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
78 torch::Tensor const& seq_lens,
79 torch::Tensor const& page_table, double scale);
80
87 -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
88 -
81 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
82 torch::Tensor const& scale,
83 std::optional<torch::Tensor> const& azp);
csrc/torch_bindings.cpp
-31
@@ -2,7 +2,6 @@
2 // cache.h, which is no longer included here after cache ops moved to
3 // _C_stable_libtorch).
4 #include <torch/all.h>
5 -#include "cuda_utils.h"
5 #include "ops.h"
6 #include "core/registration.h"
7 #include <torch/library.h>
@@ -32,27 +31,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
31 ops.def("weak_ref_tensor(Tensor input) -> Tensor");
32 ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
33
35 - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
36 - ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
37 - &get_cuda_view_from_cpu_tensor);
38 -
34 // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
35 ops.def(
36 "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
37 ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
38
44 - // Fused SiLU+Mul + per-block quantization
45 - ops.def(
46 - "silu_and_mul_per_block_quant("
47 - "Tensor! out, "
48 - "Tensor input, "
49 - "Tensor! scales, "
50 - "int group_size, "
51 - "Tensor? scale_ub=None, "
52 - "bool is_scale_transposed=False) -> ()");
53 - ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
54 - &silu_and_mul_per_block_quant);
55 -
39 // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
40 // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
41 // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
@@ -178,18 +161,4 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
161 }
162 #endif
163
181 -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
182 - // Cuda utils
183 -
184 - // Gets the specified device attribute.
185 - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
186 - cuda_utils.impl("get_device_attribute", &get_device_attribute);
187 -
188 - // Gets the maximum shared memory per block device attribute.
189 - cuda_utils.def(
190 - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
191 - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
192 - &get_max_shared_memory_per_block_device_attribute);
193 -}
194 -
164 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)