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 2, 23:15 UTC
[7/n] Migrate pos_encoding and norm kernels to libtorch stable ABI (continued) (#43209)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Signed-off-by: Chris Leonard <chleonar@redhat.com> Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
May 23, 2026 at 01:20 UTC
a7be0f342dd376bf8234eb1b7397be7a79be8f42
15 files changed
+577
-419
CMakeLists.txt
+6
-6
@@ -311,14 +311,9 @@ set(VLLM_EXT_SRC
311
"csrc/attention/paged_attention_v1.cu"
312
"csrc/attention/paged_attention_v2.cu"
313
"csrc/attention/merge_attn_states.cu"
314
- "csrc/pos_encoding_kernels.cu"
315
- "csrc/layernorm_kernels.cu"
316
- "csrc/fused_qknorm_rope_kernel.cu"
317
- "csrc/layernorm_quant_kernels.cu"
314
"csrc/sampler.cu"
315
"csrc/topk.cu"
316
"csrc/cuda_view.cu"
321
- "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
317
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
318
"csrc/quantization/activation_kernels.cu"
319
"csrc/cuda_utils_kernels.cu"
@@ -633,7 +628,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
628
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
629
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
630
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
636
- "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu")
631
+ "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
632
+ "csrc/libtorch_stable/pos_encoding_kernels.cu"
633
+ "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
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
638
if(VLLM_GPU_LANG STREQUAL "CUDA")
639
list(APPEND VLLM_STABLE_EXT_SRC
csrc/libtorch_stable/dispatch_utils.h
+82
@@ -58,6 +58,35 @@
58
THO_DISPATCH_SWITCH(TYPE, NAME, \
59
VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
60
61
+// Quant type dispatch (FP8 + INT8)
62
+#ifdef USE_ROCM
63
+ #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \
64
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \
65
+ __VA_ARGS__) \
66
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \
67
+ __VA_ARGS__) \
68
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
69
+#else
70
+ #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \
71
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \
72
+ __VA_ARGS__) \
73
+ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)
74
+#endif
75
+
76
+#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
77
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
78
+ VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
79
+
80
+// Group size dispatch (pure C++ if/else, no ATen dependency)
81
+#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
82
+ if (group_size == 128) { \
83
+ constexpr int const_group_size = 128; \
84
+ __VA_ARGS__(); \
85
+ } else if (group_size == 64) { \
86
+ constexpr int const_group_size = 64; \
87
+ __VA_ARGS__(); \
88
+ }
89
+
90
// Boolean dispatch
91
#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
92
if (expr) { \
@@ -67,3 +96,56 @@
96
constexpr bool const_expr = false; \
97
__VA_ARGS__(); \
98
}
99
+
100
+// Vec size dispatch (pure C++ switch, no ATen dependency)
101
+#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \
102
+ switch (VEC_SIZE) { \
103
+ case 16: { \
104
+ constexpr int vec_size = 16; \
105
+ __VA_ARGS__(); \
106
+ break; \
107
+ } \
108
+ case 8: { \
109
+ constexpr int vec_size = 8; \
110
+ __VA_ARGS__(); \
111
+ break; \
112
+ } \
113
+ case 4: { \
114
+ constexpr int vec_size = 4; \
115
+ __VA_ARGS__(); \
116
+ break; \
117
+ } \
118
+ case 2: { \
119
+ constexpr int vec_size = 2; \
120
+ __VA_ARGS__(); \
121
+ break; \
122
+ } \
123
+ default: { \
124
+ constexpr int vec_size = 1; \
125
+ __VA_ARGS__(); \
126
+ break; \
127
+ } \
128
+ }
129
+
130
+// Tensor rank dispatch (2D, 3D, 4D)
131
+#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \
132
+ switch (NUM_DIMS) { \
133
+ case 2: { \
134
+ constexpr int tensor_rank = 2; \
135
+ __VA_ARGS__(); \
136
+ break; \
137
+ } \
138
+ case 3: { \
139
+ constexpr int tensor_rank = 3; \
140
+ __VA_ARGS__(); \
141
+ break; \
142
+ } \
143
+ case 4: { \
144
+ constexpr int tensor_rank = 4; \
145
+ __VA_ARGS__(); \
146
+ break; \
147
+ } \
148
+ default: \
149
+ STD_TORCH_CHECK( \
150
+ false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \
151
+ }
csrc/libtorch_stable/fused_qknorm_rope_kernel.cu
renamed
+80
-77
@@ -18,21 +18,20 @@
18
#include <cuda_runtime.h>
19
#include <type_traits>
20
21
-#include <torch/cuda.h>
22
-#include <ATen/cuda/CUDAContext.h>
23
-#include <c10/cuda/CUDAGuard.h>
21
+#include "torch_utils.h"
22
25
-#include "async_util.cuh"
26
-#include "cuda_compat.h"
23
+#include "../async_util.cuh"
24
+#include "../cuda_compat.h"
25
+#include "../type_convert.cuh"
26
#include "dispatch_utils.h"
28
-#include "type_convert.cuh"
27
30
-#define CHECK_TYPE(x, st) \
31
- TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \
32
- ", while ", st, " is expected")
33
-#define CHECK_TH_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
28
+#define CHECK_TYPE(x, st) \
29
+ STD_TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \
30
+ ", while ", st, " is expected")
31
+#define CHECK_TH_CUDA(x) \
32
+ STD_TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
33
#define CHECK_CONTIGUOUS(x) \
35
- TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
34
+ STD_TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
35
#define CHECK_INPUT(x) \
36
CHECK_TH_CUDA(x); \
37
CHECK_CONTIGUOUS(x)
@@ -589,8 +588,8 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
588
});
589
break;
590
default:
592
- TORCH_CHECK(false,
593
- "Unsupported head dimension for fusedQKNormRope: ", head_dim);
591
+ STD_TORCH_CHECK(
592
+ false, "Unsupported head dimension for fusedQKNormRope: ", head_dim);
593
}
594
}
595
@@ -604,10 +603,10 @@ void launchFusedQKNormRopeNTokenHeads(
603
void const* k_weight, void const* cos_sin_cache, bool const interleave,
604
int64_t const* position_ids, int const token_heads_per_warp,
605
cudaStream_t stream) {
607
- TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
608
- token_heads_per_warp == 4 || token_heads_per_warp == 8,
609
- "token_heads_per_warp must be 1, 2, 4, or 8, got ",
610
- token_heads_per_warp);
606
+ STD_TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
607
+ token_heads_per_warp == 4 || token_heads_per_warp == 8,
608
+ "token_heads_per_warp must be 1, 2, 4, or 8, got ",
609
+ token_heads_per_warp);
610
611
// token_heads_per_warp == 1: delegate to the 1-head baseline kernel.
612
if (token_heads_per_warp == 1) {
@@ -691,7 +690,7 @@ void launchFusedQKNormRopeNTokenHeads(
690
}); \
691
break; \
692
default: \
694
- TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
693
+ STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
694
} \
695
} while (0)
696
@@ -708,19 +707,21 @@ void launchFusedQKNormRopeNTokenHeads(
707
} // namespace tensorrt_llm::kernels
708
709
void fused_qk_norm_rope(
711
- torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
712
- // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
713
- int64_t num_heads_q, // Number of query heads
714
- int64_t num_heads_k, // Number of key heads
715
- int64_t num_heads_v, // Number of value heads
716
- int64_t head_dim, // Dimension per head
717
- double eps, // Epsilon for RMS normalization
718
- torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
719
- torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
720
- torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim]
721
- bool is_neox, // Whether RoPE is applied in Neox style
722
- torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens]
723
- int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value
710
+ torch::stable::Tensor&
711
+ qkv, // Combined QKV tensor [num_tokens,
712
+ // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
713
+ int64_t num_heads_q, // Number of query heads
714
+ int64_t num_heads_k, // Number of key heads
715
+ int64_t num_heads_v, // Number of value heads
716
+ int64_t head_dim, // Dimension per head
717
+ double eps, // Epsilon for RMS normalization
718
+ torch::stable::Tensor& q_weight, // RMSNorm weights for query [head_dim]
719
+ torch::stable::Tensor& k_weight, // RMSNorm weights for key [head_dim]
720
+ torch::stable::Tensor& cos_sin_cache, // Cos/sin cache [max_position,
721
+ // head_dim]
722
+ bool is_neox, // Whether RoPE is applied in Neox style
723
+ torch::stable::Tensor& position_ids, // Position IDs for RoPE [num_tokens]
724
+ int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value
725
) {
726
// Input validation
727
CHECK_INPUT(qkv);
@@ -728,40 +729,42 @@ void fused_qk_norm_rope(
729
CHECK_INPUT(q_weight);
730
CHECK_INPUT(k_weight);
731
CHECK_INPUT(cos_sin_cache);
731
- CHECK_TYPE(position_ids, torch::kInt64);
732
-
733
- TORCH_CHECK(qkv.dim() == 2,
734
- "QKV tensor must be 2D: [num_tokens, "
735
- "(num_heads_q+num_heads_k+num_heads_v)*head_dim]");
736
- TORCH_CHECK(position_ids.dim() == 1, "Position IDs must be 1D: [num_tokens]");
737
- TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]");
738
- TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]");
739
- TORCH_CHECK(cos_sin_cache.dim() == 2,
740
- "Cos/sin cache must be 2D: [max_position, head_dim]");
741
- TORCH_CHECK(q_weight.size(0) == head_dim,
742
- "Query weights size must match head dimension");
743
- TORCH_CHECK(k_weight.size(0) == head_dim,
744
- "Key weights size must match head dimension");
745
-
746
- TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even");
747
- TORCH_CHECK(cos_sin_cache.size(1) <= head_dim,
748
- "rotary_dim must be less than or equal to head_dim");
749
-
750
- TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
751
- qkv.scalar_type() == k_weight.scalar_type(),
752
- "qkv, q_weight and k_weight must have the same dtype");
732
+ CHECK_TYPE(position_ids, torch::headeronly::ScalarType::Long);
733
+
734
+ STD_TORCH_CHECK(qkv.dim() == 2,
735
+ "QKV tensor must be 2D: [num_tokens, "
736
+ "(num_heads_q+num_heads_k+num_heads_v)*head_dim]");
737
+ STD_TORCH_CHECK(position_ids.dim() == 1,
738
+ "Position IDs must be 1D: [num_tokens]");
739
+ STD_TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]");
740
+ STD_TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]");
741
+ STD_TORCH_CHECK(cos_sin_cache.dim() == 2,
742
+ "Cos/sin cache must be 2D: [max_position, head_dim]");
743
+ STD_TORCH_CHECK(q_weight.size(0) == head_dim,
744
+ "Query weights size must match head dimension");
745
+ STD_TORCH_CHECK(k_weight.size(0) == head_dim,
746
+ "Key weights size must match head dimension");
747
+
748
+ STD_TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even");
749
+ STD_TORCH_CHECK(cos_sin_cache.size(1) <= head_dim,
750
+ "rotary_dim must be less than or equal to head_dim");
751
+
752
+ STD_TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
753
+ qkv.scalar_type() == k_weight.scalar_type(),
754
+ "qkv, q_weight and k_weight must have the same dtype");
755
756
int64_t num_tokens = qkv.size(0);
755
- TORCH_CHECK(position_ids.size(0) == num_tokens,
756
- "Number of tokens in position_ids must match QKV");
757
+ STD_TORCH_CHECK(position_ids.size(0) == num_tokens,
758
+ "Number of tokens in position_ids must match QKV");
759
760
int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
759
- TORCH_CHECK(
761
+ STD_TORCH_CHECK(
762
qkv.size(1) == total_heads * head_dim,
763
"QKV tensor size must match total number of heads and head dimension");
764
763
- auto device_id = qkv.get_device();
764
- auto stream = at::cuda::getCurrentCUDAStream(device_id);
765
+ const torch::stable::accelerator::DeviceGuard device_guard(
766
+ qkv.get_device_index());
767
+ auto stream = get_current_cuda_stream(qkv.get_device_index());
768
769
// Select token_heads_per_warp: forced value if >0, else auto-select.
770
// Auto thresholds are calibrated on SM 9.0 (H100). On other architectures,
@@ -771,8 +774,7 @@ void fused_qk_norm_rope(
774
token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp);
775
} else {
776
token_heads_per_warp = 1;
774
- auto* dev_prop = at::cuda::getDeviceProperties(device_id);
775
- int sm_version = dev_prop->major * 10 + dev_prop->minor;
777
+ int sm_version = get_device_prop()->major * 10 + get_device_prop()->minor;
778
int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k);
779
if (sm_version == 90) {
780
if (head_dim >= 256) {
@@ -795,21 +797,22 @@ void fused_qk_norm_rope(
797
}
798
}
799
798
- VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
799
- using qkv_scalar_t = scalar_t;
800
- VLLM_DISPATCH_FLOATING_TYPES(
801
- cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
802
- using cache_scalar_t = scalar_t;
803
- tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
804
- qkv_scalar_t, cache_scalar_t>(
805
- qkv.data_ptr(), static_cast<int>(num_tokens),
806
- static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
807
- static_cast<int>(num_heads_v), static_cast<int>(head_dim),
808
- static_cast<int>(cos_sin_cache.size(1)), static_cast<float>(eps),
809
- q_weight.data_ptr(), k_weight.data_ptr(),
810
- cos_sin_cache.data_ptr(), !is_neox,
811
- reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
812
- token_heads_per_warp, stream);
813
- });
814
- });
800
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
801
+ qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
802
+ using qkv_scalar_t = scalar_t;
803
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
804
+ cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
805
+ using cache_scalar_t = scalar_t;
806
+ tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
807
+ qkv_scalar_t, cache_scalar_t>(
808
+ qkv.data_ptr(), static_cast<int>(num_tokens),
809
+ static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
810
+ static_cast<int>(num_heads_v), static_cast<int>(head_dim),
811
+ static_cast<int>(cos_sin_cache.size(1)),
812
+ static_cast<float>(eps), q_weight.data_ptr(),
813
+ k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox,
814
+ reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
815
+ token_heads_per_warp, stream);
816
+ });
817
+ });
818
}
csrc/libtorch_stable/layernorm_kernels.cu
renamed
+55
-49
@@ -1,11 +1,12 @@
1
-#include "type_convert.cuh"
2
-#include "dispatch_utils.h"
3
-#include "cub_helpers.h"
4
-#include "core/batch_invariant.hpp"
5
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
1
+#include <numeric>
2
+
3
+#include "torch_utils.h"
4
7
-#include <torch/cuda.h>
8
-#include <c10/cuda/CUDAGuard.h>
5
+#include "../cub_helpers.h"
6
+#include "../core/batch_invariant.hpp"
7
+#include "../type_convert.cuh"
8
+#include "dispatch_utils.h"
9
+#include "quantization/vectorization_utils.cuh"
10
11
namespace vllm {
12
@@ -189,16 +190,16 @@ fused_add_rms_norm_kernel(
190
191
} // namespace vllm
192
192
-void rms_norm(torch::Tensor& out, // [..., hidden_size]
193
- torch::Tensor& input, // [..., hidden_size]
194
- torch::Tensor& weight, // [hidden_size]
193
+void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]
194
+ torch::stable::Tensor& input, // [..., hidden_size]
195
+ torch::stable::Tensor& weight, // [hidden_size]
196
double epsilon) {
196
- TORCH_CHECK(out.is_contiguous());
197
+ STD_TORCH_CHECK(out.is_contiguous());
198
if (input.stride(-1) != 1) {
198
- input = input.contiguous();
199
+ input = torch::stable::contiguous(input);
200
}
200
- TORCH_CHECK(input.stride(-1) == 1);
201
- TORCH_CHECK(weight.is_contiguous());
201
+ STD_TORCH_CHECK(input.stride(-1) == 1);
202
+ STD_TORCH_CHECK(weight.is_contiguous());
203
204
int hidden_size = input.size(-1);
205
@@ -213,45 +214,49 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]
214
// For large num_tokens, use smaller blocks to increase SM concurrency.
215
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
216
dim3 grid(num_tokens);
216
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
217
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
218
- VLLM_DISPATCH_RANK234(num_dims, [&] {
219
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
220
- const int calculated_vec_size =
221
- std::gcd(16 / sizeof(scalar_t), hidden_size);
222
- const int block_size =
223
- std::min(hidden_size / calculated_vec_size, max_block_size);
224
- dim3 block(block_size);
225
- VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
226
- vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
227
- <<<grid, block, 0, stream>>>(
228
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
229
- input_stride_d2, input_stride_d3, input_stride_d4,
230
- input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
231
- epsilon, num_tokens, hidden_size);
232
- });
233
- });
217
+ const torch::stable::accelerator::DeviceGuard device_guard(
218
+ input.get_device_index());
219
+ const cudaStream_t stream = get_current_cuda_stream();
220
+ VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] {
221
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
222
+ input.scalar_type(), "rms_norm_kernel", [&] {
223
+ const int calculated_vec_size =
224
+ std::gcd(16 / sizeof(scalar_t), hidden_size);
225
+ const int block_size =
226
+ std::min(hidden_size / calculated_vec_size, max_block_size);
227
+ dim3 block(block_size);
228
+ VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
229
+ vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
230
+ <<<grid, block, 0, stream>>>(
231
+ out.mutable_data_ptr<scalar_t>(),
232
+ input.const_data_ptr<scalar_t>(), input_stride_d2,
233
+ input_stride_d3, input_stride_d4, input_shape_d2,
234
+ input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
235
+ num_tokens, hidden_size);
236
+ });
237
+ });
238
});
239
}
240
237
-#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
238
- VLLM_DISPATCH_FLOATING_TYPES( \
239
- input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
240
- vllm::fused_add_rms_norm_kernel<scalar_t, width> \
241
- <<<grid, block, 0, stream>>>( \
242
- input.data_ptr<scalar_t>(), input_stride, \
243
- residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
244
- epsilon, num_tokens, hidden_size); \
241
+#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
242
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
243
+ input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
244
+ vllm::fused_add_rms_norm_kernel<scalar_t, width> \
245
+ <<<grid, block, 0, stream>>>( \
246
+ input.mutable_data_ptr<scalar_t>(), input_stride, \
247
+ residual.mutable_data_ptr<scalar_t>(), \
248
+ weight.const_data_ptr<scalar_t>(), epsilon, num_tokens, \
249
+ hidden_size); \
250
});
251
247
-void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
248
- torch::Tensor& residual, // [..., hidden_size]
249
- torch::Tensor& weight, // [hidden_size]
252
+void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]
253
+ torch::stable::Tensor& residual, // [..., hidden_size]
254
+ torch::stable::Tensor& weight, // [hidden_size]
255
double epsilon) {
251
- TORCH_CHECK(weight.scalar_type() == input.scalar_type());
252
- TORCH_CHECK(input.scalar_type() == residual.scalar_type());
253
- TORCH_CHECK(residual.is_contiguous());
254
- TORCH_CHECK(weight.is_contiguous());
256
+ STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
257
+ STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());
258
+ STD_TORCH_CHECK(residual.is_contiguous());
259
+ STD_TORCH_CHECK(weight.is_contiguous());
260
int hidden_size = input.size(-1);
261
int64_t input_stride = input.stride(-2);
262
int num_tokens = input.numel() / hidden_size;
@@ -263,8 +268,9 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
268
hiding on global mem ops. */
269
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
270
dim3 block(std::min(hidden_size, max_block_size));
266
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
267
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
271
+ const torch::stable::accelerator::DeviceGuard device_guard(
272
+ input.get_device_index());
273
+ const cudaStream_t stream = get_current_cuda_stream();
274
/*If the tensor types are FP16/BF16, try to use the optimized kernel
275
with packed + vectorized ops.
276
Max optimization is achieved with a width-8 vector of FP16/BF16s
csrc/libtorch_stable/layernorm_quant_kernels.cu
renamed
+46
-39
@@ -5,15 +5,16 @@
5
* Currently, only static fp8 quantization is supported.
6
*/
7
8
-#include "type_convert.cuh"
9
-#include "quantization/w8a8/fp8/common.cuh"
10
-#include "dispatch_utils.h"
11
-#include "cub_helpers.h"
12
-#include "core/batch_invariant.hpp"
13
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
8
+#include <numeric>
9
+
10
+#include "torch_utils.h"
11
15
-#include <torch/cuda.h>
16
-#include <c10/cuda/CUDAGuard.h>
12
+#include "../cub_helpers.h"
13
+#include "../core/batch_invariant.hpp"
14
+#include "../quantization/w8a8/fp8/common.cuh"
15
+#include "../type_convert.cuh"
16
+#include "dispatch_utils.h"
17
+#include "quantization/vectorization_utils.cuh"
18
19
namespace vllm {
20
@@ -202,12 +203,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
203
204
} // namespace vllm
205
205
-void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
206
- torch::Tensor& input, // [..., hidden_size]
207
- torch::Tensor& weight, // [hidden_size]
208
- torch::Tensor& scale, // [1]
209
- double epsilon) {
210
- TORCH_CHECK(out.is_contiguous());
206
+void rms_norm_static_fp8_quant(
207
+ torch::stable::Tensor& out, // [..., hidden_size]
208
+ torch::stable::Tensor& input, // [..., hidden_size]
209
+ torch::stable::Tensor& weight, // [hidden_size]
210
+ torch::stable::Tensor& scale, // [1]
211
+ double epsilon) {
212
+ STD_TORCH_CHECK(out.is_contiguous());
213
int hidden_size = input.size(-1);
214
int input_stride = input.stride(-2);
215
int num_tokens = input.numel() / hidden_size;
@@ -215,24 +217,26 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
217
// For large num_tokens, use smaller blocks to increase SM concurrency.
218
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
219
dim3 grid(num_tokens);
218
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
219
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
220
- VLLM_DISPATCH_FLOATING_TYPES(
220
+ const torch::stable::accelerator::DeviceGuard device_guard(
221
+ input.get_device_index());
222
+ const cudaStream_t stream = get_current_cuda_stream();
223
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
224
input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
222
- VLLM_DISPATCH_FP8_TYPES(
225
+ VLLM_STABLE_DISPATCH_FP8_TYPES(
226
out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
227
const int calculated_vec_size =
228
std::gcd(16 / sizeof(scalar_t), hidden_size);
229
const int block_size =
230
std::min(hidden_size / calculated_vec_size, max_block_size);
231
dim3 block(block_size);
229
- VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
232
+ VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
233
vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t,
234
vec_size>
235
<<<grid, block, 0, stream>>>(
233
- out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
234
- input_stride, weight.data_ptr<scalar_t>(),
235
- scale.data_ptr<float>(), epsilon, num_tokens,
236
+ out.mutable_data_ptr<fp8_t>(),
237
+ input.const_data_ptr<scalar_t>(), input_stride,
238
+ weight.const_data_ptr<scalar_t>(),
239
+ scale.const_data_ptr<float>(), epsilon, num_tokens,
240
hidden_size);
241
});
242
});
@@ -240,30 +244,32 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
244
}
245
246
#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
243
- VLLM_DISPATCH_FLOATING_TYPES( \
247
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
248
input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] { \
245
- VLLM_DISPATCH_FP8_TYPES( \
249
+ VLLM_STABLE_DISPATCH_FP8_TYPES( \
250
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
251
vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t, \
252
width, fp8_t> \
253
<<<grid, block, 0, stream>>>( \
250
- out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(), \
251
- input_stride, residual.data_ptr<scalar_t>(), \
252
- weight.data_ptr<scalar_t>(), scale.data_ptr<float>(), \
253
- epsilon, num_tokens, hidden_size); \
254
+ out.mutable_data_ptr<fp8_t>(), \
255
+ input.mutable_data_ptr<scalar_t>(), input_stride, \
256
+ residual.mutable_data_ptr<scalar_t>(), \
257
+ weight.const_data_ptr<scalar_t>(), \
258
+ scale.const_data_ptr<float>(), epsilon, num_tokens, \
259
+ hidden_size); \
260
}); \
261
});
262
void fused_add_rms_norm_static_fp8_quant(
257
- torch::Tensor& out, // [..., hidden_size],
258
- torch::Tensor& input, // [..., hidden_size]
259
- torch::Tensor& residual, // [..., hidden_size]
260
- torch::Tensor& weight, // [hidden_size]
261
- torch::Tensor& scale, // [1]
263
+ torch::stable::Tensor& out, // [..., hidden_size],
264
+ torch::stable::Tensor& input, // [..., hidden_size]
265
+ torch::stable::Tensor& residual, // [..., hidden_size]
266
+ torch::stable::Tensor& weight, // [hidden_size]
267
+ torch::stable::Tensor& scale, // [1]
268
double epsilon) {
263
- TORCH_CHECK(out.is_contiguous());
264
- TORCH_CHECK(residual.is_contiguous());
265
- TORCH_CHECK(residual.scalar_type() == input.scalar_type());
266
- TORCH_CHECK(weight.scalar_type() == input.scalar_type());
269
+ STD_TORCH_CHECK(out.is_contiguous());
270
+ STD_TORCH_CHECK(residual.is_contiguous());
271
+ STD_TORCH_CHECK(residual.scalar_type() == input.scalar_type());
272
+ STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
273
int hidden_size = input.size(-1);
274
int input_stride = input.stride(-2);
275
int num_tokens = input.numel() / hidden_size;
@@ -275,8 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(
281
hiding on global mem ops. */
282
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
283
dim3 block(std::min(hidden_size, max_block_size));
278
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
279
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
284
+ const torch::stable::accelerator::DeviceGuard device_guard(
285
+ input.get_device_index());
286
+ const cudaStream_t stream = get_current_cuda_stream();
287
/*If the tensor types are FP16/BF16, try to use the optimized kernel
288
with packed + vectorized ops.
289
Max optimization is achieved with a width-8 vector of FP16/BF16s
csrc/libtorch_stable/ops.h
+53
@@ -167,6 +167,59 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
167
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
168
bool inplace);
169
170
+// Layernorm kernels (shared CUDA/ROCm)
171
+void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
172
+ torch::stable::Tensor& weight, double epsilon);
173
+
174
+void fused_add_rms_norm(torch::stable::Tensor& input,
175
+ torch::stable::Tensor& residual,
176
+ torch::stable::Tensor& weight, double epsilon);
177
+
178
+// Layernorm-quant kernels (shared CUDA/ROCm)
179
+void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
180
+ torch::stable::Tensor& input,
181
+ torch::stable::Tensor& weight,
182
+ torch::stable::Tensor& scale, double epsilon);
183
+
184
+void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out,
185
+ torch::stable::Tensor& input,
186
+ torch::stable::Tensor& residual,
187
+ torch::stable::Tensor& weight,
188
+ torch::stable::Tensor& scale,
189
+ double epsilon);
190
+
191
+// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
192
+void rms_norm_dynamic_per_token_quant(
193
+ torch::stable::Tensor& out, torch::stable::Tensor const& input,
194
+ torch::stable::Tensor const& weight, torch::stable::Tensor& scales,
195
+ double const var_epsilon, std::optional<torch::stable::Tensor> scale_ub,
196
+ std::optional<torch::stable::Tensor> residual);
197
+
198
+void rms_norm_per_block_quant(torch::stable::Tensor& out,
199
+ torch::stable::Tensor const& input,
200
+ torch::stable::Tensor const& weight,
201
+ torch::stable::Tensor& scales,
202
+ double const var_epsilon,
203
+ std::optional<torch::stable::Tensor> scale_ub,
204
+ std::optional<torch::stable::Tensor> residual,
205
+ int64_t group_size, bool is_scale_transposed);
206
+
207
+// Positional encoding kernels (shared CUDA/ROCm)
208
+void rotary_embedding(torch::stable::Tensor& positions,
209
+ torch::stable::Tensor& query,
210
+ std::optional<torch::stable::Tensor> key,
211
+ int64_t head_size, torch::stable::Tensor& cos_sin_cache,
212
+ bool is_neox, int64_t rope_dim_offset, bool inverse);
213
+
214
+void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
215
+ int64_t num_heads_k, int64_t num_heads_v,
216
+ int64_t head_dim, double eps,
217
+ torch::stable::Tensor& q_weight,
218
+ torch::stable::Tensor& k_weight,
219
+ torch::stable::Tensor& cos_sin_cache, bool is_neox,
220
+ torch::stable::Tensor& position_ids,
221
+ int64_t forced_token_heads_per_warp);
222
+
223
// Activation kernels (shared CUDA/ROCm)
224
void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
225
void silu_and_mul_clamp(torch::stable::Tensor& out,
csrc/libtorch_stable/pos_encoding_kernels.cu
renamed
+52
-46
@@ -1,8 +1,6 @@
1
-#include <torch/all.h>
2
-#include <ATen/cuda/CUDAContext.h>
3
-#include <c10/cuda/CUDAGuard.h>
1
+#include "torch_utils.h"
2
5
-#include "cuda_compat.h"
3
+#include "../cuda_compat.h"
4
#include "dispatch_utils.h"
5
6
namespace vllm {
@@ -103,35 +101,37 @@ __global__ void rotary_embedding_kernel(
101
} // namespace vllm
102
103
void rotary_embedding(
106
- torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
107
- torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or
108
- // [num_tokens, num_heads * head_size] or
109
- // [batch_size, seq_len, num_heads, head_size] or
110
- // [num_tokens, num_heads, head_size]
111
- std::optional<torch::Tensor> key,
104
+ torch::stable::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
105
+ torch::stable::Tensor&
106
+ query, // [batch_size, seq_len, num_heads * head_size] or
107
+ // [num_tokens, num_heads * head_size] or
108
+ // [batch_size, seq_len, num_heads, head_size] or
109
+ // [num_tokens, num_heads, head_size]
110
+ std::optional<torch::stable::Tensor> key,
111
// null or
112
// [batch_size, seq_len, num_kv_heads * head_size] or
113
// [num_tokens, num_kv_heads * head_size] or
114
// [batch_size, seq_len, num_heads, head_size] or
115
// [num_tokens, num_heads, head_size]
116
int64_t head_size,
118
- torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
117
+ torch::stable::Tensor& cos_sin_cache, // [max_position, rot_dim]
118
bool is_neox, int64_t rope_dim_offset, bool inverse) {
119
// num_tokens = batch_size * seq_len
120
int64_t num_tokens = positions.numel();
121
int positions_ndim = positions.dim();
122
123
// Make sure num_tokens dim is consistent across positions, query, and key
125
- TORCH_CHECK(
124
+ STD_TORCH_CHECK(
125
positions_ndim == 1 || positions_ndim == 2,
126
"positions must have shape [num_tokens] or [batch_size, seq_len]");
127
if (positions_ndim == 1) {
129
- TORCH_CHECK(query.size(0) == positions.size(0) &&
130
- (!key.has_value() || key->size(0) == positions.size(0)),
131
- "query, key and positions must have the same number of tokens");
128
+ STD_TORCH_CHECK(
129
+ query.size(0) == positions.size(0) &&
130
+ (!key.has_value() || key->size(0) == positions.size(0)),
131
+ "query, key and positions must have the same number of tokens");
132
}
133
if (positions_ndim == 2) {
134
- TORCH_CHECK(
134
+ STD_TORCH_CHECK(
135
query.size(0) == positions.size(0) &&
136
(!key.has_value() || key->size(0) == positions.size(0)) &&
137
query.size(1) == positions.size(1) &&
@@ -143,20 +143,20 @@ void rotary_embedding(
143
// hidden_size = num_heads * head_size
144
int query_hidden_size = query.numel() / num_tokens;
145
int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
146
- TORCH_CHECK(query_hidden_size % head_size == 0);
147
- TORCH_CHECK(key_hidden_size % head_size == 0);
146
+ STD_TORCH_CHECK(query_hidden_size % head_size == 0);
147
+ STD_TORCH_CHECK(key_hidden_size % head_size == 0);
148
149
// Make sure query and key have consistent number of heads
150
int num_heads = query_hidden_size / head_size;
151
int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
152
- TORCH_CHECK(num_heads % num_kv_heads == 0);
152
+ STD_TORCH_CHECK(num_heads % num_kv_heads == 0);
153
154
int rot_dim = cos_sin_cache.size(1);
155
int seq_dim_idx = positions_ndim - 1;
156
int64_t query_stride = query.stride(seq_dim_idx);
157
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
158
159
- TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
159
+ STD_TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
160
// Determine head stride: for [*, heads, head_size] use stride of last dim;
161
// for flat [*, heads*head_size], heads blocks are contiguous of size
162
// head_size
@@ -166,30 +166,36 @@ void rotary_embedding(
166
167
dim3 grid(num_tokens);
168
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
169
- const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
170
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
171
- VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
172
- using query_t = scalar_t;
173
- VLLM_DISPATCH_FLOATING_TYPES(
174
- cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
175
- using cache_t = scalar_t;
176
- if (is_neox) {
177
- vllm::rotary_embedding_kernel<query_t, cache_t, true>
178
- <<<grid, block, 0, stream>>>(
179
- positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
180
- key.has_value() ? key->data_ptr<query_t>() : nullptr,
181
- cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
182
- key_stride, head_stride, num_heads, num_kv_heads, head_size,
183
- rope_dim_offset, inverse);
184
- } else {
185
- vllm::rotary_embedding_kernel<query_t, cache_t, false>
186
- <<<grid, block, 0, stream>>>(
187
- positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
188
- key.has_value() ? key->data_ptr<query_t>() : nullptr,
189
- cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
190
- key_stride, head_stride, num_heads, num_kv_heads, head_size,
191
- rope_dim_offset, inverse);
192
- }
193
- });
194
- });
169
+ const torch::stable::accelerator::DeviceGuard device_guard(
170
+ query.get_device_index());
171
+ const cudaStream_t stream = get_current_cuda_stream();
172
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
173
+ query.scalar_type(), "rotary_embedding", [&] {
174
+ using query_t = scalar_t;
175
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
176
+ cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
177
+ using cache_t = scalar_t;
178
+ if (is_neox) {
179
+ vllm::rotary_embedding_kernel<query_t, cache_t, true>
180
+ <<<grid, block, 0, stream>>>(
181
+ positions.const_data_ptr<int64_t>(),
182
+ query.mutable_data_ptr<query_t>(),
183
+ key.has_value() ? key->mutable_data_ptr<query_t>()
184
+ : nullptr,
185
+ cos_sin_cache.const_data_ptr<cache_t>(), rot_dim,
186
+ query_stride, key_stride, head_stride, num_heads,
187
+ num_kv_heads, head_size, rope_dim_offset, inverse);
188
+ } else {
189
+ vllm::rotary_embedding_kernel<query_t, cache_t, false>
190
+ <<<grid, block, 0, stream>>>(
191
+ positions.const_data_ptr<int64_t>(),
192
+ query.mutable_data_ptr<query_t>(),
193
+ key.has_value() ? key->mutable_data_ptr<query_t>()
194
+ : nullptr,
195
+ cos_sin_cache.const_data_ptr<cache_t>(), rot_dim,
196
+ query_stride, key_stride, head_stride, num_heads,
197
+ num_kv_heads, head_size, rope_dim_offset, inverse);
198
+ }
199
+ });
200
+ });
201
}
csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
renamed
+118
-99
@@ -1,6 +1,5 @@
1
2
-#include <ATen/cuda/CUDAContext.h>
3
-#include <c10/cuda/CUDAGuard.h>
2
+#include "../../torch_utils.h"
3
4
#include "../../dispatch_utils.h"
5
#include "layernorm_utils.cuh"
@@ -134,63 +133,71 @@ __global__ void rms_norm_per_block_quant_kernel(
133
// Residual add + RMS norm + dynamic per token
134
template <typename scalar_in_t>
135
void rms_norm_dynamic_per_token_quant_dispatch(
137
- torch::Tensor& out, // [..., hidden_size]
138
- torch::Tensor const& input, // [..., hidden_size]
139
- torch::Tensor const& weight, // [hidden_size]
140
- torch::Tensor& scales, // [num_tokens]
141
- double const var_epsilon, // Variance epsilon used in norm calculation
142
- std::optional<at::Tensor> const& scale_ub,
143
- std::optional<at::Tensor>& residual) {
136
+ torch::stable::Tensor& out, // [..., hidden_size]
137
+ torch::stable::Tensor const& input, // [..., hidden_size]
138
+ torch::stable::Tensor const& weight, // [hidden_size]
139
+ torch::stable::Tensor& scales, // [num_tokens]
140
+ double const var_epsilon, // Variance epsilon used in norm calculation
141
+ std::optional<torch::stable::Tensor> const& scale_ub,
142
+ std::optional<torch::stable::Tensor>& residual) {
143
int32_t hidden_size = input.size(-1);
145
- int32_t input_stride = input.view({-1, hidden_size}).stride(0);
144
+ int32_t input_stride =
145
+ torch::stable::view(input, {-1, hidden_size}).stride(0);
146
auto num_tokens = input.numel() / hidden_size;
147
148
dim3 grid(num_tokens);
149
dim3 block(std::min(hidden_size, 1024));
150
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
151
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
150
+ const torch::stable::accelerator::DeviceGuard device_guard(
151
+ input.get_device_index());
152
+ const cudaStream_t stream = get_current_cuda_stream();
153
153
- VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
154
- VLLM_DISPATCH_QUANT_TYPES(
154
+ VLLM_STABLE_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
155
+ VLLM_STABLE_DISPATCH_QUANT_TYPES(
156
out.scalar_type(), "rms_norm_dynamic_per_token_quant_kernel", [&] {
157
vllm::rms_norm_dynamic_per_token_quant_kernel<scalar_in_t, scalar_t,
158
has_residual>
159
<<<grid, block, 0, stream>>>(
159
- out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
160
- input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
161
- scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
160
+ out.mutable_data_ptr<scalar_t>(),
161
+ scales.mutable_data_ptr<float>(),
162
+ input.const_data_ptr<scalar_in_t>(),
163
+ weight.const_data_ptr<scalar_in_t>(),
164
+ scale_ub.has_value() ? scale_ub->const_data_ptr<float>()
165
+ : nullptr,
166
var_epsilon, hidden_size, input_stride,
163
- has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
167
+ has_residual ? residual->mutable_data_ptr<scalar_in_t>()
168
+ : nullptr);
169
});
170
});
171
}
172
173
void rms_norm_dynamic_per_token_quant(
169
- torch::Tensor& out, // [..., hidden_size]
170
- torch::Tensor const& input, // [..., hidden_size]
171
- torch::Tensor const& weight, // [hidden_size]
172
- torch::Tensor& scales, // [num_tokens]
173
- double const var_epsilon, // Variance epsilon used in norm calculation
174
- std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
175
- static c10::ScalarType kFp8Type = is_fp8_ocp()
176
- ? c10::ScalarType::Float8_e4m3fn
177
- : c10::ScalarType::Float8_e4m3fnuz;
178
- TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
179
- TORCH_CHECK(out.is_contiguous());
180
- TORCH_CHECK(input.stride(-1) == 1,
181
- "Input must be contiguous in the last dimension");
174
+ torch::stable::Tensor& out, // [..., hidden_size]
175
+ torch::stable::Tensor const& input, // [..., hidden_size]
176
+ torch::stable::Tensor const& weight, // [hidden_size]
177
+ torch::stable::Tensor& scales, // [num_tokens]
178
+ double const var_epsilon, // Variance epsilon used in norm calculation
179
+ std::optional<torch::stable::Tensor> scale_ub,
180
+ std::optional<torch::stable::Tensor> residual) {
181
+ static torch::headeronly::ScalarType kFp8Type =
182
+ is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
183
+ : torch::headeronly::ScalarType::Float8_e4m3fnuz;
184
+ STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
185
+ out.scalar_type() == torch::headeronly::ScalarType::Char);
186
+ STD_TORCH_CHECK(out.is_contiguous());
187
+ STD_TORCH_CHECK(input.stride(-1) == 1,
188
+ "Input must be contiguous in the last dimension");
189
190
if (scale_ub.has_value()) {
184
- TORCH_CHECK(out.dtype() == kFp8Type);
191
+ STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
192
}
186
- TORCH_CHECK(weight.dtype() == input.dtype());
187
- TORCH_CHECK(scales.dtype() == torch::kFloat32);
193
+ STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
194
+ STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
195
if (residual) {
189
- TORCH_CHECK(residual->scalar_type() == input.scalar_type());
190
- TORCH_CHECK(residual->is_contiguous());
196
+ STD_TORCH_CHECK(residual->scalar_type() == input.scalar_type());
197
+ STD_TORCH_CHECK(residual->is_contiguous());
198
}
199
193
- VLLM_DISPATCH_FLOATING_TYPES(
200
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
201
input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {
202
rms_norm_dynamic_per_token_quant_dispatch<scalar_t>(
203
out, input, weight, scales, var_epsilon, scale_ub, residual);
@@ -199,103 +206,115 @@ void rms_norm_dynamic_per_token_quant(
206
207
// Residual add + RMS norm + dynamic per token
208
void rms_norm_per_block_quant_dispatch(
202
- torch::Tensor& out, // [..., hidden_size]
203
- torch::Tensor const& input, // [..., hidden_size]
204
- torch::Tensor const& weight, // [hidden_size]
205
- torch::Tensor& scales, // [num_tokens, hidden_size / group_size] or
206
- // [hidden_size / group_size, num_tokens]
209
+ torch::stable::Tensor& out, // [..., hidden_size]
210
+ torch::stable::Tensor const& input, // [..., hidden_size]
211
+ torch::stable::Tensor const& weight, // [hidden_size]
212
+ torch::stable::Tensor& scales, // [num_tokens, hidden_size /
213
+ // group_size] or
214
+ // [hidden_size / group_size,
215
+ // num_tokens]
216
int32_t group_size,
217
double const var_epsilon, // Variance epsilon used in norm calculation
209
- std::optional<at::Tensor> const& scale_ub,
210
- std::optional<at::Tensor>& residual, bool is_scale_transposed) {
218
+ std::optional<torch::stable::Tensor> const& scale_ub,
219
+ std::optional<torch::stable::Tensor>& residual, bool is_scale_transposed) {
220
int32_t hidden_size = input.size(-1);
212
- int32_t input_stride = input.view({-1, hidden_size}).stride(0);
221
+ int32_t input_stride =
222
+ torch::stable::view(input, {-1, hidden_size}).stride(0);
223
214
- TORCH_CHECK(hidden_size % 4 == 0,
215
- "Hidden size must be divisible by 4 for vectorized access");
216
- TORCH_CHECK(input_stride % 4 == 0,
217
- "Input stride must be divisible by 4 for vectorized access");
218
- TORCH_CHECK(group_size % 4 == 0,
219
- "Group size must be divisible by 4 for vectorized access");
224
+ STD_TORCH_CHECK(hidden_size % 4 == 0,
225
+ "Hidden size must be divisible by 4 for vectorized access");
226
+ STD_TORCH_CHECK(input_stride % 4 == 0,
227
+ "Input stride must be divisible by 4 for vectorized access");
228
+ STD_TORCH_CHECK(group_size % 4 == 0,
229
+ "Group size must be divisible by 4 for vectorized access");
230
231
auto num_tokens = input.numel() / hidden_size;
232
233
dim3 grid(num_tokens);
234
const int max_block_size = (num_tokens <= 256) ? 512 : 256;
235
dim3 block(std::min(hidden_size, max_block_size));
226
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
227
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
236
+ const torch::stable::accelerator::DeviceGuard device_guard(
237
+ input.get_device_index());
238
+ const cudaStream_t stream = get_current_cuda_stream();
239
229
- VLLM_DISPATCH_FLOATING_TYPES(
240
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
241
input.scalar_type(), "rms_norm_per_block_quant_fp_dispatch", [&] {
242
using scalar_in_t = scalar_t;
232
- VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
233
- VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
234
- VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
235
- VLLM_DISPATCH_QUANT_TYPES(
236
- out.scalar_type(), "rms_norm_per_block_quant_kernel", [&] {
237
- vllm::rms_norm_per_block_quant_kernel<scalar_in_t, scalar_t,
238
- has_residual,
239
- transpose_scale, gs>
240
- <<<grid, block, 0, stream>>>(
241
- out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
242
- input.data_ptr<scalar_in_t>(),
243
- weight.data_ptr<scalar_in_t>(),
244
- scale_ub.has_value() ? scale_ub->data_ptr<float>()
245
- : nullptr,
243
+ VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
244
+ VLLM_STABLE_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
245
+ VLLM_STABLE_DISPATCH_BOOL(
246
+ is_scale_transposed, transpose_scale, [&] {
247
+ VLLM_STABLE_DISPATCH_QUANT_TYPES(
248
+ out.scalar_type(), "rms_norm_per_block_quant_kernel",
249
+ [&] {
250
+ vllm::rms_norm_per_block_quant_kernel<
251
+ scalar_in_t, scalar_t, has_residual,
252
+ transpose_scale, gs><<<grid, block, 0, stream>>>(
253
+ out.mutable_data_ptr<scalar_t>(),
254
+ scales.mutable_data_ptr<float>(),
255
+ input.const_data_ptr<scalar_in_t>(),
256
+ weight.const_data_ptr<scalar_in_t>(),
257
+ scale_ub.has_value()
258
+ ? scale_ub->const_data_ptr<float>()
259
+ : nullptr,
260
var_epsilon, hidden_size, input_stride,
247
- has_residual ? residual->data_ptr<scalar_in_t>()
248
- : nullptr,
261
+ has_residual
262
+ ? residual->mutable_data_ptr<scalar_in_t>()
263
+ : nullptr,
264
scales.stride(1));
250
- });
251
- });
265
+ });
266
+ });
267
});
268
});
269
});
270
}
271
257
-void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
258
- torch::Tensor const& weight,
259
- torch::Tensor& scales, double const var_epsilon,
260
- std::optional<torch::Tensor> scale_ub,
261
- std::optional<torch::Tensor> residual,
272
+void rms_norm_per_block_quant(torch::stable::Tensor& out,
273
+ torch::stable::Tensor const& input,
274
+ torch::stable::Tensor const& weight,
275
+ torch::stable::Tensor& scales,
276
+ double const var_epsilon,
277
+ std::optional<torch::stable::Tensor> scale_ub,
278
+ std::optional<torch::stable::Tensor> residual,
279
int64_t group_size, bool is_scale_transposed) {
263
- static c10::ScalarType kFp8Type = is_fp8_ocp()
264
- ? c10::ScalarType::Float8_e4m3fn
265
- : c10::ScalarType::Float8_e4m3fnuz;
266
- TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
267
- TORCH_CHECK(out.is_contiguous());
268
- TORCH_CHECK(input.stride(-1) == 1,
269
- "Input must be contiguous in the last dimension");
280
+ static torch::headeronly::ScalarType kFp8Type =
281
+ is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
282
+ : torch::headeronly::ScalarType::Float8_e4m3fnuz;
283
+ STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
284
+ out.scalar_type() == torch::headeronly::ScalarType::Char);
285
+ STD_TORCH_CHECK(out.is_contiguous());
286
+ STD_TORCH_CHECK(input.stride(-1) == 1,
287
+ "Input must be contiguous in the last dimension");
288
289
if (scale_ub.has_value()) {
272
- TORCH_CHECK(out.dtype() == kFp8Type);
290
+ STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
291
}
274
- TORCH_CHECK(weight.dtype() == input.dtype());
275
- TORCH_CHECK(scales.dtype() == torch::kFloat32);
292
+ STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
293
+ STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
294
if (residual) {
277
- TORCH_CHECK(residual->scalar_type() == input.scalar_type());
278
- TORCH_CHECK(residual->is_contiguous());
295
+ STD_TORCH_CHECK(residual->scalar_type() == input.scalar_type());
296
+ STD_TORCH_CHECK(residual->is_contiguous());
297
}
298
281
- TORCH_CHECK(group_size == 128 || group_size == 64,
282
- "Unsupported group size: ", group_size);
299
+ STD_TORCH_CHECK(group_size == 128 || group_size == 64,
300
+ "Unsupported group size: ", group_size);
301
302
if (scales.stride(1) > 1) {
285
- TORCH_CHECK(is_scale_transposed,
286
- "Outer scale stride must be 1 when scales are not transposed");
303
+ STD_TORCH_CHECK(
304
+ is_scale_transposed,
305
+ "Outer scale stride must be 1 when scales are not transposed");
306
}
307
308
int64_t hidden_size = input.size(-1);
290
- TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
291
- "hidden_size must be a positive multiple of group_size");
309
+ STD_TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
310
+ "hidden_size must be a positive multiple of group_size");
311
int64_t num_tokens = input.numel() / hidden_size;
312
int64_t num_groups = hidden_size / group_size;
294
- TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
295
- "scales buffer too small: need ", num_tokens * num_groups,
296
- " elements, got ", scales.numel());
313
+ STD_TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
314
+ "scales buffer too small: need ", num_tokens * num_groups,
315
+ " elements, got ", scales.numel());
316
317
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
318
var_epsilon, scale_ub, residual,
319
is_scale_transposed);
301
-}
\ No newline at end of file
320
+}
csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh
renamed
+2
-2
@@ -8,8 +8,8 @@
8
#include "quantization/utils.cuh"
9
#include "quant_conversions.cuh"
10
11
-#include "../../cub_helpers.h"
12
-#include "../../cuda_compat.h"
11
+#include "../../../cub_helpers.h"
12
+#include "../../../cuda_compat.h"
13
14
namespace vllm {
15
csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh
renamed
+1
-1
@@ -6,7 +6,7 @@
6
7
#include "libtorch_stable/quantization/vectorization.cuh"
8
// TODO(luka/varun):refactor common.cuh to use this file instead
9
-#include "../w8a8/fp8/common.cuh"
9
+#include "../../../quantization/w8a8/fp8/common.cuh"
10
11
namespace vllm {
12
csrc/libtorch_stable/torch_bindings.cpp
+71
@@ -267,9 +267,62 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
267
// conditionally compiled so impl registration is in source file
268
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
269
270
+ // Apply Root Mean Square (RMS) Normalization to the input tensor.
271
+ ops.def(
272
+ "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
273
+ "()");
274
+
275
+ // In-place fused Add and RMS Normalization.
276
+ ops.def(
277
+ "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
278
+ "float epsilon) -> ()");
279
+
280
+ // Layernorm-quant
281
+ // Apply Root Mean Square (RMS) Normalization to the input tensor.
282
+ ops.def(
283
+ "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
284
+ "Tensor scale, float epsilon) -> "
285
+ "()");
286
+
287
+ // In-place fused Add and RMS Normalization.
288
+ ops.def(
289
+ "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
290
+ "Tensor! residual, Tensor weight, "
291
+ "Tensor scale, float epsilon) -> ()");
292
+
293
+ // Fused Layernorm + Quant kernels
294
+ ops.def(
295
+ "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
296
+ "Tensor weight, Tensor! scale, float epsilon, "
297
+ "Tensor? scale_ub, Tensor!? residual) -> ()");
298
+
299
+ // Fused Layernorm + Block quant kernels
300
+ ops.def(
301
+ "rms_norm_per_block_quant(Tensor! result, Tensor input, "
302
+ "Tensor weight, Tensor! scale, float epsilon, "
303
+ "Tensor? scale_ub, Tensor!? residual, int group_size, "
304
+ "bool is_scale_transposed) -> ()");
305
+
306
+ // Rotary embedding
307
+ // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
308
+ ops.def(
309
+ "rotary_embedding(Tensor positions, Tensor! query,"
310
+ " Tensor!? key, int head_size,"
311
+ " Tensor cos_sin_cache, bool is_neox, int "
312
+ "rope_dim_offset=0, bool inverse=False) -> ()");
313
+
314
+ // Function for fused QK Norm and RoPE
315
+ ops.def(
316
+ "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
317
+ "int num_heads_k, int num_heads_v, int head_dim, float eps, "
318
+ "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
319
+ "bool is_neox, Tensor position_ids, "
320
+ "int forced_token_heads_per_warp=-1) -> ()");
321
+
322
// Activation ops
323
// Activation function used in SwiGLU.
324
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
325
+
326
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
327
328
// SwiGLU activation with input clamping.
@@ -416,6 +469,24 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
469
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
470
#endif
471
472
+ // Layernorm kernels (shared CUDA/ROCm)
473
+ ops.impl("rms_norm", TORCH_BOX(&rms_norm));
474
+ ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
475
+
476
+ // Layernorm-quant kernels (shared CUDA/ROCm)
477
+ ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant));
478
+ ops.impl("fused_add_rms_norm_static_fp8_quant",
479
+ TORCH_BOX(&fused_add_rms_norm_static_fp8_quant));
480
+
481
+ // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
482
+ ops.impl("rms_norm_dynamic_per_token_quant",
483
+ TORCH_BOX(&rms_norm_dynamic_per_token_quant));
484
+ ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
485
+
486
+ // Positional encoding kernels (shared CUDA/ROCm)
487
+ ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
488
+ ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
489
+
490
// Activation kernels (shared CUDA/ROCm)
491
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
492
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
csrc/ops.h
+6
-32
@@ -61,19 +61,15 @@ void merge_attn_states(
61
const std::optional<int64_t> prefill_tokens_with_context,
62
const std::optional<torch::Tensor>& output_scale = std::nullopt);
63
64
+// rms_norm and fused_add_rms_norm declarations also exist in
65
+// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
66
+// because the CPU build still uses these torch::Tensor declarations.
67
void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
68
double epsilon);
69
70
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
71
torch::Tensor& weight, double epsilon);
72
70
-void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
71
- int64_t num_heads_k, int64_t num_heads_v,
72
- int64_t head_dim, double eps, torch::Tensor& q_weight,
73
- torch::Tensor& k_weight, torch::Tensor& cos_sin_cache,
74
- bool is_neox, torch::Tensor& position_ids,
75
- int64_t forced_token_heads_per_warp);
76
-
73
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
74
torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
75
torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
@@ -99,37 +95,15 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
95
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
96
int64_t max_seq_len);
97
102
-void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
103
- torch::Tensor& weight, torch::Tensor& scale,
104
- double epsilon);
105
-
106
-void fused_add_rms_norm_static_fp8_quant(torch::Tensor& out,
107
- torch::Tensor& input,
108
- torch::Tensor& residual,
109
- torch::Tensor& weight,
110
- torch::Tensor& scale, double epsilon);
111
-
112
-void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
113
- torch::Tensor const& input,
114
- torch::Tensor const& weight,
115
- torch::Tensor& scales,
116
- double const epsilon,
117
- std::optional<torch::Tensor> scale_ub,
118
- std::optional<torch::Tensor> residual);
119
-
120
-void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
121
- torch::Tensor const& weight,
122
- torch::Tensor& scales, double const epsilon,
123
- std::optional<torch::Tensor> scale_ub,
124
- std::optional<torch::Tensor> residual,
125
- int64_t group_size, bool is_scale_transposed);
126
-
98
void silu_and_mul_per_block_quant(torch::Tensor& out,
99
torch::Tensor const& input,
100
torch::Tensor& scales, int64_t group_size,
101
std::optional<torch::Tensor> scale_ub,
102
bool is_scale_transposed);
103
104
+// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
105
+// ABI for CUDA). It remains here because the CPU build still uses these
106
+// torch::Tensor declarations.
107
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
108
std::optional<torch::Tensor> key, int64_t head_size,
109
torch::Tensor& cos_sin_cache, bool is_neox,
csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu
+1
-2
@@ -5,8 +5,7 @@
5
#include <c10/cuda/CUDAGuard.h>
6
7
#include "../../dispatch_utils.h"
8
-#include "quant_conversions.cuh"
9
-#include "../w8a8/fp8/common.cuh"
8
+#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
9
10
namespace vllm {
11
csrc/torch_bindings.cpp
-64
@@ -94,28 +94,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
94
ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
95
&silu_and_mul_per_block_quant);
96
97
- // Layernorm
98
- // Apply Root Mean Square (RMS) Normalization to the input tensor.
99
- ops.def(
100
- "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
101
- "()");
102
- ops.impl("rms_norm", torch::kCUDA, &rms_norm);
103
-
104
- // In-place fused Add and RMS Normalization.
105
- ops.def(
106
- "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
107
- "float epsilon) -> ()");
108
- ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
109
-
110
- // Function for fused QK Norm and RoPE
111
- ops.def(
112
- "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
113
- "int num_heads_k, int num_heads_v, int head_dim, float eps, "
114
- "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
115
- "bool is_neox, Tensor position_ids, "
116
- "int forced_token_heads_per_warp=-1) -> ()");
117
- ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
118
-
97
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
98
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
99
// kernel launch.
@@ -152,48 +130,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
130
"Tensor workspace, int k, int max_seq_len) -> ()");
131
ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
132
155
- // Layernorm-quant
156
- // Apply Root Mean Square (RMS) Normalization to the input tensor.
157
- ops.def(
158
- "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
159
- "Tensor scale, float epsilon) -> "
160
- "()");
161
- ops.impl("rms_norm_static_fp8_quant", torch::kCUDA,
162
- &rms_norm_static_fp8_quant);
163
-
164
- // In-place fused Add and RMS Normalization.
165
- ops.def(
166
- "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
167
- "Tensor! residual, Tensor weight, "
168
- "Tensor scale, float epsilon) -> ()");
169
- ops.impl("fused_add_rms_norm_static_fp8_quant", torch::kCUDA,
170
- &fused_add_rms_norm_static_fp8_quant);
171
-
172
- // Fused Layernorm + Quant kernels
173
- ops.def(
174
- "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
175
- "Tensor weight, Tensor! scale, float epsilon, "
176
- "Tensor? scale_ub, Tensor!? residual) -> ()");
177
- ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
178
- &rms_norm_dynamic_per_token_quant);
179
-
180
- // Fused Layernorm + Block quant kernels
181
- ops.def(
182
- "rms_norm_per_block_quant(Tensor! result, Tensor input, "
183
- "Tensor weight, Tensor! scale, float epsilon, "
184
- "Tensor? scale_ub, Tensor!? residual, int group_size, "
185
- "bool is_scale_transposed) -> ()");
186
- ops.impl("rms_norm_per_block_quant", torch::kCUDA, &rms_norm_per_block_quant);
187
-
188
- // Rotary embedding
189
- // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
190
- ops.def(
191
- "rotary_embedding(Tensor positions, Tensor! query,"
192
- " Tensor!? key, int head_size,"
193
- " Tensor cos_sin_cache, bool is_neox, int "
194
- "rope_dim_offset=0, bool inverse=False) -> ()");
195
- ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
196
-
133
// Quantization ops
134
#ifndef USE_ROCM
135
csrc/type_convert.cuh
+4
-2
@@ -1,8 +1,10 @@
1
#pragma once
2
3
-#include <torch/all.h>
3
+#include <torch/headeronly/util/BFloat16.h>
4
+#include <torch/headeronly/util/Half.h>
5
6
#ifndef USE_ROCM
7
+ #include <cuda.h>
8
#include <cuda_bf16.h>
9
#include <cuda_fp16.h>
10
#else
@@ -191,4 +193,4 @@ struct alignas(16) _f16Vec {
193
return result;
194
}
195
};
194
-} // namespace vllm
\ No newline at end of file
196
+} // namespace vllm