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 6, 00:45 UTC
[2/n] Migrate per_token_group_quant to torch stable ABI (#36058)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
mikaylagawarecki committed
Mar 25, 2026 at 13:15 UTC
bf4cc9ed2d1f2075a1dbdeed233a9a44a7d48247
22 files changed
+207
-133
CMakeLists.txt
+5
-4
@@ -343,9 +343,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
343
"csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
344
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
345
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
346
- "csrc/cutlass_extensions/common.cpp"
347
- "csrc/quantization/w8a8/fp8/per_token_group_quant.cu"
348
- "csrc/quantization/w8a8/int8/per_token_group_quant.cu")
346
+ "csrc/cutlass_extensions/common.cpp")
347
348
set_gencode_flags_for_srcs(
349
SRCS "${VLLM_EXT_SRC}"
@@ -969,7 +967,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
967
"csrc/libtorch_stable/torch_bindings.cpp")
968
969
if(VLLM_GPU_LANG STREQUAL "CUDA")
972
- list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/permute_cols.cu")
970
+ list(APPEND VLLM_STABLE_EXT_SRC
971
+ "csrc/libtorch_stable/permute_cols.cu"
972
+ "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
973
+ "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
974
endif()
975
976
if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/cache_kernels.cu
+2
-1
@@ -7,7 +7,8 @@
7
#include "cuda_utils.h"
8
#include "cuda_compat.h"
9
#include "dispatch_utils.h"
10
-#include "quantization/vectorization_utils.cuh"
10
+
11
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
12
#include "concat_mla_q.cuh"
13
14
#ifdef USE_ROCM
csrc/layernorm_kernels.cu
+1
-1
@@ -2,7 +2,7 @@
2
#include "dispatch_utils.h"
3
#include "cub_helpers.h"
4
#include "core/batch_invariant.hpp"
5
-#include "quantization/vectorization_utils.cuh"
5
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
6
7
#include <torch/cuda.h>
8
#include <c10/cuda/CUDAGuard.h>
csrc/layernorm_quant_kernels.cu
+1
-1
@@ -10,7 +10,7 @@
10
#include "dispatch_utils.h"
11
#include "cub_helpers.h"
12
#include "core/batch_invariant.hpp"
13
-#include "quantization/vectorization_utils.cuh"
13
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
14
15
#include <torch/cuda.h>
16
#include <c10/cuda/CUDAGuard.h>
csrc/libtorch_stable/dispatch_utils.h
new
+60
@@ -0,0 +1,60 @@
1
+/*
2
+ * Stable ABI compatible dispatch utilities for vLLM.
3
+ * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros
4
+ * instead of the ATen (AT_*) macros.
5
+ *
6
+ * These macros use:
7
+ * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH
8
+ * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE
9
+ * - torch::headeronly::ScalarType instead of at::ScalarType
10
+ *
11
+ * Add more macros here as needed when migrating additional kernels.
12
+ */
13
+#pragma once
14
+
15
+#include <torch/headeronly/core/Dispatch.h>
16
+#include <torch/headeronly/core/ScalarType.h>
17
+#include <torch/headeronly/util/Exception.h>
18
+
19
+// Need a special dispatch case macro since we will nest the FP8 dispatch.
20
+// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
21
+#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \
22
+ THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
23
+
24
+#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \
25
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
26
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
27
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__)
28
+
29
+#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
30
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
31
+ VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
32
+
33
+// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format
34
+#ifdef USE_ROCM
35
+ #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \
36
+ VLLM_STABLE_DISPATCH_FP8_CASE( \
37
+ torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
38
+ VLLM_STABLE_DISPATCH_FP8_CASE( \
39
+ torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
40
+#else
41
+ #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \
42
+ VLLM_STABLE_DISPATCH_FP8_CASE( \
43
+ torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__)
44
+#endif
45
+
46
+// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
47
+// See VLLM_STABLE_DISPATCH_FP8_CASE above.
48
+#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
49
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
50
+ VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
51
+
52
+// Boolean dispatch
53
+#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
54
+ if (expr) { \
55
+ constexpr bool const_expr = true; \
56
+ __VA_ARGS__(); \
57
+ } else { \
58
+ constexpr bool const_expr = false; \
59
+ __VA_ARGS__(); \
60
+ }
csrc/libtorch_stable/ops.h
+21
@@ -6,4 +6,25 @@
6
#ifndef USE_ROCM
7
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
8
torch::stable::Tensor const& perm);
9
+
10
+void per_token_group_quant_fp8(const torch::stable::Tensor& input,
11
+ torch::stable::Tensor& output_q,
12
+ torch::stable::Tensor& output_s,
13
+ int64_t group_size, double eps, double fp8_min,
14
+ double fp8_max, bool scale_ue8m0,
15
+ bool dummy_is_scale_transposed,
16
+ bool dummy_is_tma_aligned);
17
+
18
+// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales.
19
+void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
20
+ torch::stable::Tensor& output_q,
21
+ torch::stable::Tensor& output_s_packed,
22
+ int64_t group_size, double eps,
23
+ double min_8bit, double max_8bit);
24
+
25
+void per_token_group_quant_int8(const torch::stable::Tensor& input,
26
+ torch::stable::Tensor& output_q,
27
+ torch::stable::Tensor& output_s,
28
+ int64_t group_size, double eps, double int8_min,
29
+ double int8_max);
30
#endif
csrc/libtorch_stable/quantization/vectorization.cuh
renamed
+2
-2
@@ -4,8 +4,8 @@
4
*/
5
6
// Include both AMD and NVIDIA fp8 types to avoid circular import
7
-#include <c10/util/Float8_e4m3fnuz.h>
8
-#include <c10/util/Float8_e4m3fn.h>
7
+#include <torch/headeronly/util/Float8_e4m3fnuz.h>
8
+#include <torch/headeronly/util/Float8_e4m3fn.h>
9
10
namespace vllm {
11
csrc/libtorch_stable/quantization/vectorization_utils.cuh
renamed
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
renamed
+50
-46
@@ -1,16 +1,18 @@
1
-#include <ATen/cuda/CUDAContext.h>
1
+#include <torch/csrc/stable/tensor.h>
2
+#include <torch/csrc/stable/ops.h>
3
+#include <torch/headeronly/util/Exception.h>
4
+#include <torch/headeronly/core/ScalarType.h>
5
3
-#include "quantization/w8a8/per_token_group_quant_8bit.h"
6
+#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h"
7
8
#include <cmath>
9
10
#include <cuda_fp8.h>
11
9
-#include <torch/all.h>
10
-
11
-#include "quantization/vectorization.cuh"
12
-#include "quantization/vectorization_utils.cuh"
13
-#include "dispatch_utils.h"
12
+#include "libtorch_stable/quantization/vectorization.cuh"
13
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
14
+#include "libtorch_stable/dispatch_utils.h"
15
+#include "libtorch_stable/torch_utils.h"
16
17
__device__ __forceinline__ float GroupReduceMax(float val) {
18
unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
@@ -154,20 +156,20 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
156
return 1;
157
}
158
157
-void per_token_group_quant_8bit(const torch::Tensor& input,
158
- torch::Tensor& output_q,
159
- torch::Tensor& output_s, int64_t group_size,
160
- double eps, double min_8bit, double max_8bit,
161
- bool scale_ue8m0) {
162
- TORCH_CHECK(input.is_contiguous());
163
- TORCH_CHECK(output_q.is_contiguous());
159
+void per_token_group_quant_8bit(const torch::stable::Tensor& input,
160
+ torch::stable::Tensor& output_q,
161
+ torch::stable::Tensor& output_s,
162
+ int64_t group_size, double eps, double min_8bit,
163
+ double max_8bit, bool scale_ue8m0) {
164
+ STD_TORCH_CHECK(input.is_contiguous());
165
+ STD_TORCH_CHECK(output_q.is_contiguous());
166
167
const int num_groups = input.numel() / group_size;
168
167
- TORCH_CHECK(input.numel() % group_size == 0);
168
- TORCH_CHECK(output_s.dim() == 2);
169
+ STD_TORCH_CHECK(input.numel() % group_size == 0);
170
+ STD_TORCH_CHECK(output_s.dim() == 2);
171
170
- cudaStream_t stream = at::cuda::getCurrentCUDAStream();
172
+ cudaStream_t stream = get_current_cuda_stream();
173
174
constexpr int THREADS_PER_GROUP = 16;
175
@@ -222,11 +224,11 @@ void per_token_group_quant_8bit(const torch::Tensor& input,
224
} \
225
} while (0)
226
225
- VLLM_DISPATCH_FLOATING_TYPES(
227
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
228
input.scalar_type(), "per_token_group_quant_8bit", ([&] {
227
- if (dst_type == at::ScalarType::Float8_e4m3fn) {
229
+ if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
230
LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
229
- } else if (dst_type == at::ScalarType::Char) {
231
+ } else if (dst_type == torch::headeronly::ScalarType::Char) {
232
LAUNCH_KERNEL(scalar_t, int8_t);
233
}
234
}));
@@ -294,41 +296,42 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
296
threads_per_group, y_s, min_8bit, max_8bit);
297
}
298
297
-void per_token_group_quant_8bit_packed(const torch::Tensor& input,
298
- torch::Tensor& output_q,
299
- torch::Tensor& output_s_packed,
299
+void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
300
+ torch::stable::Tensor& output_q,
301
+ torch::stable::Tensor& output_s_packed,
302
int64_t group_size, double eps,
303
double min_8bit, double max_8bit) {
302
- TORCH_CHECK(input.is_contiguous());
303
- TORCH_CHECK(output_q.is_contiguous());
304
+ STD_TORCH_CHECK(input.is_contiguous());
305
+ STD_TORCH_CHECK(output_q.is_contiguous());
306
307
const int64_t k = input.size(-1);
306
- TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
307
- ") must be divisible by group_size (", group_size, ").");
308
+ STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
309
+ ") must be divisible by group_size (", group_size, ").");
310
311
const int64_t mn = input.numel() / k;
312
const int64_t groups_per_row = k / group_size;
313
const int64_t num_groups = mn * groups_per_row;
314
313
- TORCH_CHECK(output_s_packed.dim() == 2,
314
- "output_s_packed must be 2D, got dim=", output_s_packed.dim(),
315
- ".");
315
+ STD_TORCH_CHECK(output_s_packed.dim() == 2,
316
+ "output_s_packed must be 2D, got dim=", output_s_packed.dim(),
317
+ ".");
318
319
const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4;
320
const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4;
321
320
- TORCH_CHECK(output_s_packed.scalar_type() == at::ScalarType::Int,
321
- "output_s_packed must have dtype int32 for UE8M0-packed scales.");
322
+ STD_TORCH_CHECK(
323
+ output_s_packed.scalar_type() == torch::headeronly::ScalarType::Int,
324
+ "output_s_packed must have dtype int32 for UE8M0-packed scales.");
325
// DeepGEMM expects SFA scales in MN-major form with shape
326
// [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last
327
// dimension.
325
- TORCH_CHECK(output_s_packed.size(0) == mn &&
326
- output_s_packed.size(1) == k_num_packed_sfk,
327
- "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
328
- "], but got [", output_s_packed.size(0), ", ",
329
- output_s_packed.size(1), "].");
328
+ STD_TORCH_CHECK(output_s_packed.size(0) == mn &&
329
+ output_s_packed.size(1) == k_num_packed_sfk,
330
+ "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
331
+ "], but got [", output_s_packed.size(0), ", ",
332
+ output_s_packed.size(1), "].");
333
331
- cudaStream_t stream = at::cuda::getCurrentCUDAStream();
334
+ cudaStream_t stream = get_current_cuda_stream();
335
336
constexpr int THREADS_PER_GROUP = 16;
337
@@ -340,7 +343,7 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
343
344
// zero-initialize packed scales, since we use atomicOr to accumulate
345
// exponents from different groups.
343
- output_s_packed.zero_();
346
+ torch::stable::zero_(output_s_packed);
347
348
#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \
349
do { \
@@ -359,14 +362,14 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
362
static_cast<float>(max_8bit)); \
363
} while (0)
364
362
- VLLM_DISPATCH_FLOATING_TYPES(
365
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
366
input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] {
364
- if (dst_type == at::ScalarType::Float8_e4m3fn) {
367
+ if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
368
LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3);
366
- } else if (dst_type == at::ScalarType::Char) {
369
+ } else if (dst_type == torch::headeronly::ScalarType::Char) {
370
LAUNCH_PACKED_KERNEL(scalar_t, int8_t);
371
} else {
369
- TORCH_CHECK(
372
+ STD_TORCH_CHECK(
373
false,
374
"per_token_group_quant_8bit_packed only supports FP8/INT8 "
375
"outputs.");
@@ -376,12 +379,13 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
379
#undef LAUNCH_PACKED_KERNEL
380
}
381
379
-void per_token_group_quant_fp8(const torch::Tensor& input,
380
- torch::Tensor& output_q, torch::Tensor& output_s,
382
+void per_token_group_quant_fp8(const torch::stable::Tensor& input,
383
+ torch::stable::Tensor& output_q,
384
+ torch::stable::Tensor& output_s,
385
int64_t group_size, double eps, double fp8_min,
386
double fp8_max, bool scale_ue8m0,
387
bool dummy_is_scale_transposed = false,
388
bool dummy_is_tma_aligned = false) {
389
per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
390
fp8_min, fp8_max, scale_ue8m0);
387
-}
\ No newline at end of file
391
+}
csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu
new
+12
@@ -0,0 +1,12 @@
1
+#include <torch/csrc/stable/tensor.h>
2
+
3
+#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h"
4
+
5
+void per_token_group_quant_int8(const torch::stable::Tensor& input,
6
+ torch::stable::Tensor& output_q,
7
+ torch::stable::Tensor& output_s,
8
+ int64_t group_size, double eps, double int8_min,
9
+ double int8_max) {
10
+ per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
11
+ int8_min, int8_max);
12
+}
csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h
new
+10
@@ -0,0 +1,10 @@
1
+#pragma once
2
+
3
+#include <torch/csrc/stable/tensor.h>
4
+
5
+// 8-bit per-token-group quantization helper used by both FP8 and INT8
6
+void per_token_group_quant_8bit(const torch::stable::Tensor& input,
7
+ torch::stable::Tensor& output_q,
8
+ torch::stable::Tensor& output_s,
9
+ int64_t group_size, double eps, double min_8bit,
10
+ double max_8bit, bool scale_ue8m0 = false);
csrc/libtorch_stable/torch_bindings.cpp
+35
-4
@@ -6,15 +6,46 @@
6
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
7
// Note: We register under namespace "_C" so ops are accessible as
8
// torch.ops._C.<op_name> for compatibility with existing code.
9
-STABLE_TORCH_LIBRARY_FRAGMENT(_C, m) {
9
+STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
10
#ifndef USE_ROCM
11
- m.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
11
+ ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
12
+#endif
13
+
14
+#ifndef USE_ROCM
15
+ // Compute per-token-group FP8 quantized tensor and scaling factor.
16
+ // The dummy arguments are here so we can correctly fuse with RMSNorm.
17
+ ops.def(
18
+ "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
19
+ "output_s, "
20
+ "int group_size, float eps, float fp8_min, float fp8_max, bool "
21
+ "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
22
+ ") -> ()");
23
+ // Compute per-token-group 8-bit quantized tensor and UE8M0-packed,
24
+ // TMA-aligned scales for DeepGEMM.
25
+ ops.def(
26
+ "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, "
27
+ "Tensor! output_s_packed, int group_size, float eps, float fp8_min, "
28
+ "float fp8_max) -> ()");
29
+ // Compute per-token-group INT8 quantized tensor and scaling factor.
30
+ ops.def(
31
+ "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
32
+ "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
33
+ "()");
34
#endif
35
}
36
15
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
37
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
38
+#ifndef USE_ROCM
39
+ ops.impl("permute_cols", TORCH_BOX(&permute_cols));
40
+#endif
41
+
42
#ifndef USE_ROCM
17
- m.impl("permute_cols", TORCH_BOX(&permute_cols));
43
+ // Per-token group quantization
44
+ ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
45
+ ops.impl("per_token_group_fp8_quant_packed",
46
+ TORCH_BOX(&per_token_group_quant_8bit_packed));
47
+ ops.impl("per_token_group_quant_int8",
48
+ TORCH_BOX(&per_token_group_quant_int8));
49
#endif
50
}
51
csrc/libtorch_stable/torch_utils.h
+3
-1
@@ -1,11 +1,13 @@
1
#pragma once
2
3
#include <torch/csrc/inductor/aoti_torch/c/shim.h>
4
+#include <torch/headeronly/util/shim_utils.h>
5
+
6
#include <cuda_runtime.h>
7
8
// Utility to get the current CUDA stream for a given device using stable APIs.
9
// Returns a cudaStream_t for use in kernel launches.
8
-inline cudaStream_t get_current_cuda_stream(int32_t device_index) {
10
+inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
11
void* stream_ptr = nullptr;
12
TORCH_ERROR_CODE_CHECK(
13
aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
csrc/ops.h
-19
@@ -306,25 +306,6 @@ void silu_and_mul_scaled_fp4_experts_quant(
306
torch::Tensor const& input_offset_by_experts,
307
torch::Tensor const& output_scale_offset_by_experts);
308
309
-void per_token_group_quant_fp8(const torch::Tensor& input,
310
- torch::Tensor& output_q, torch::Tensor& output_s,
311
- int64_t group_size, double eps, double fp8_min,
312
- double fp8_max, bool scale_ue8m0,
313
- bool dummy_is_scale_transposed,
314
- bool dummy_is_tma_aligned);
315
-
316
-void per_token_group_quant_int8(const torch::Tensor& input,
317
- torch::Tensor& output_q,
318
- torch::Tensor& output_s, int64_t group_size,
319
- double eps, double int8_min, double int8_max);
320
-
321
-// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales.
322
-void per_token_group_quant_8bit_packed(const torch::Tensor& input,
323
- torch::Tensor& output_q,
324
- torch::Tensor& output_s_packed,
325
- int64_t group_size, double eps,
326
- double min_8bit, double max_8bit);
327
-
309
#endif
310
311
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
csrc/quantization/fused_kernels/layernorm_utils.cuh
+1
-1
@@ -4,7 +4,7 @@
4
* __device__ layernorm utilities.
5
*/
6
7
-#include "quantization/vectorization.cuh"
7
+#include "libtorch_stable/quantization/vectorization.cuh"
8
#include "quantization/utils.cuh"
9
#include "quant_conversions.cuh"
10
csrc/quantization/fused_kernels/quant_conversions.cuh
+1
-1
@@ -4,7 +4,7 @@
4
* __device__ helper functions to deal with float -> quant datatype conversion
5
*/
6
7
-#include "quantization/vectorization.cuh"
7
+#include "libtorch_stable/quantization/vectorization.cuh"
8
// TODO(luka/varun):refactor common.cuh to use this file instead
9
#include "quantization/w8a8/fp8/common.cuh"
10
csrc/quantization/w8a8/fp8/common.cu
+1
-1
@@ -1,7 +1,7 @@
1
#include "common.cuh"
2
#include "dispatch_utils.h"
3
#include "cub_helpers.h"
4
-#include "quantization/vectorization_utils.cuh"
4
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
5
#include <c10/cuda/CUDAGuard.h>
6
#include <ATen/cuda/Exceptions.h>
7
#include <tuple>
csrc/quantization/w8a8/fp8/common.cuh
+1
-1
@@ -1,6 +1,6 @@
1
#pragma once
2
3
-#include "quantization/vectorization.cuh"
3
+#include "libtorch_stable/quantization/vectorization.cuh"
4
#include "quantization/utils.cuh"
5
6
#include <cmath>
csrc/quantization/w8a8/int8/per_token_group_quant.cu
deleted
-12
@@ -1,12 +0,0 @@
1
-#include <ATen/cuda/CUDAContext.h>
2
-#include <torch/all.h>
3
-
4
-#include "quantization/w8a8/per_token_group_quant_8bit.h"
5
-
6
-void per_token_group_quant_int8(const torch::Tensor& input,
7
- torch::Tensor& output_q,
8
- torch::Tensor& output_s, int64_t group_size,
9
- double eps, double int8_min, double int8_max) {
10
- per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
11
- int8_min, int8_max);
12
-}
\ No newline at end of file
csrc/quantization/w8a8/int8/scaled_quant.cu
+1
-1
@@ -5,7 +5,7 @@
5
#include <cmath>
6
7
#include "dispatch_utils.h"
8
-#include "quantization/vectorization_utils.cuh"
8
+#include "libtorch_stable/quantization/vectorization_utils.cuh"
9
#include "cub_helpers.h"
10
11
static inline __device__ int8_t float_to_int8_rn(float x) {
csrc/quantization/w8a8/per_token_group_quant_8bit.h
deleted
-9
@@ -1,9 +0,0 @@
1
-#pragma once
2
-#include <torch/all.h>
3
-
4
-// 8-bit per-token-group quantization helper used by both FP8 and INT8
5
-void per_token_group_quant_8bit(const torch::Tensor& input,
6
- torch::Tensor& output_q,
7
- torch::Tensor& output_s, int64_t group_size,
8
- double eps, double min_8bit, double max_8bit,
9
- bool scale_ue8m0 = false);
\ No newline at end of file
csrc/torch_bindings.cpp
-28
@@ -653,34 +653,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
653
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
654
655
#ifndef USE_ROCM
656
- // Compute per-token-group FP8 quantized tensor and scaling factor.
657
- // The dummy arguments are here so we can correctly fuse with RMSNorm.
658
- ops.def(
659
- "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
660
- "output_s, "
661
- "int group_size, float eps, float fp8_min, float fp8_max, bool "
662
- "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
663
- ") -> ()");
664
- ops.impl("per_token_group_fp8_quant", torch::kCUDA,
665
- &per_token_group_quant_fp8);
666
-
667
- // Compute per-token-group 8-bit quantized tensor and UE8M0-packed,
668
- // TMA-aligned scales for DeepGEMM.
669
- ops.def(
670
- "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, "
671
- "Tensor! output_s_packed, int group_size, float eps, float fp8_min, "
672
- "float fp8_max) -> ()");
673
- ops.impl("per_token_group_fp8_quant_packed", torch::kCUDA,
674
- &per_token_group_quant_8bit_packed);
675
-
676
- // Compute per-token-group INT8 quantized tensor and scaling factor.
677
- ops.def(
678
- "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
679
- "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
680
- "()");
681
- ops.impl("per_token_group_quant_int8", torch::kCUDA,
682
- &per_token_group_quant_int8);
683
-
656
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
657
ops.def(
658
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "