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, 12:15 UTC
[9/n] Migrate attention and cache kernels to torch stable ABI (continued) (#43717)
Signed-off-by: Chris Leonard <chleonar@redhat.com> Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Chris Leonard committed
May 29, 2026 at 00:44 UTC
22a58640b4563f5945aa2052e9e61d425351588d
23 files changed
+911
-722
CMakeLists.txt
+9
-11
@@ -305,10 +305,6 @@ endif()
305
#
306
307
set(VLLM_EXT_SRC
308
- "csrc/cache_kernels.cu"
309
- "csrc/cache_kernels_fused.cu"
310
- "csrc/attention/paged_attention_v1.cu"
311
- "csrc/attention/paged_attention_v2.cu"
308
"csrc/cuda_view.cu"
309
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
310
"csrc/quantization/activation_kernels.cu"
@@ -647,7 +643,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
643
"csrc/libtorch_stable/attention/merge_attn_states.cu"
644
"csrc/libtorch_stable/sampler.cu"
645
"csrc/libtorch_stable/topk.cu"
650
- "csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
646
+ "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
647
+ "csrc/libtorch_stable/attention/paged_attention_v1.cu"
648
+ "csrc/libtorch_stable/attention/paged_attention_v2.cu"
649
+ "csrc/libtorch_stable/cache_kernels.cu"
650
+ "csrc/libtorch_stable/cache_kernels_fused.cu")
651
652
if(VLLM_GPU_LANG STREQUAL "CUDA")
653
list(APPEND VLLM_STABLE_EXT_SRC
@@ -924,13 +924,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
924
SRCS "${SRCS}"
925
CUDA_ARCHS "${FP4_ARCHS}")
926
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
927
- # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
928
- # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
929
- set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
927
+ set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
928
set_gencode_flags_for_srcs(
929
SRCS "${NVFP4_KV_SRC}"
930
CUDA_ARCHS "${FP4_ARCHS}")
933
- target_sources(_C PRIVATE ${NVFP4_KV_SRC})
931
+ list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
932
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
933
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
934
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
@@ -960,11 +958,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
958
SRCS "${SRCS}"
959
CUDA_ARCHS "${FP4_ARCHS}")
960
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
963
- set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
961
+ set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
962
set_gencode_flags_for_srcs(
963
SRCS "${NVFP4_KV_SRC}"
964
CUDA_ARCHS "${FP4_ARCHS}")
967
- target_sources(_C PRIVATE ${NVFP4_KV_SRC})
965
+ list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
966
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
967
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
968
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
csrc/libtorch_stable/activation_kernels.cu
+1
-1
@@ -4,7 +4,7 @@
4
#include <cmath>
5
6
#include "../cuda_compat.h"
7
-#include "../cuda_vec_utils.cuh"
7
+#include "cuda_vec_utils.cuh"
8
#include "dispatch_utils.h"
9
#include "torch_utils.h"
10
csrc/libtorch_stable/attention/attention_kernels.cuh
renamed
+4
-7
@@ -17,21 +17,18 @@
17
* limitations under the License.
18
*/
19
20
-#include <torch/all.h>
21
-#include <ATen/cuda/CUDAContext.h>
22
-#include <c10/cuda/CUDAGuard.h>
20
#include <algorithm>
21
25
-#include "attention_dtypes.h"
22
+#include "../../attention/attention_dtypes.h"
23
#include "attention_utils.cuh"
27
-#include "../cuda_compat.h"
24
+#include "../../cuda_compat.h"
25
26
#ifdef USE_ROCM
27
#include <hip/hip_bf16.h>
31
- #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
28
+ #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
29
typedef __hip_bfloat16 __nv_bfloat16;
30
#else
34
- #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
31
+ #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
32
#endif
33
34
#define MAX(a, b) ((a) > (b) ? (a) : (b))
csrc/libtorch_stable/attention/attention_utils.cuh
renamed
+2
-2
@@ -18,8 +18,8 @@
18
*/
19
#pragma once
20
21
-#include "../cuda_compat.h"
22
-#include "attention_dtypes.h"
21
+#include "../../cuda_compat.h"
22
+#include "../../attention/attention_dtypes.h"
23
24
#include <float.h>
25
#include <type_traits>
csrc/libtorch_stable/attention/merge_attn_states.cu
+1
-1
@@ -7,7 +7,7 @@
7
#include <torch/headeronly/core/ScalarType.h>
8
9
#include "../../attention/attention_dtypes.h"
10
-#include "../../attention/attention_utils.cuh"
10
+#include "attention_utils.cuh"
11
#include "../../quantization/w8a8/fp8/common.cuh"
12
13
namespace vllm {
csrc/libtorch_stable/attention/paged_attention_v1.cu
renamed
+41
-37
@@ -16,8 +16,9 @@
16
* See the License for the specific language governing permissions and
17
* limitations under the License.
18
*/
19
+#include "../torch_utils.h"
20
#include "attention_kernels.cuh"
20
-#include "../cuda_compat.h"
21
+#include "../../cuda_compat.h"
22
23
#define MAX(a, b) ((a) > (b) ? (a) : (b))
24
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,13 +45,15 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
45
vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
46
int NUM_THREADS = 128>
47
void paged_attention_v1_launcher(
47
- torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
48
- torch::Tensor& value_cache, int num_kv_heads, float scale,
49
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
50
- const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
51
- torch::Tensor& v_scale, const int tp_rank,
52
- const int blocksparse_local_blocks, const int blocksparse_vert_stride,
53
- const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
48
+ torch::stable::Tensor& out, torch::stable::Tensor& query,
49
+ torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
50
+ int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
51
+ torch::stable::Tensor& seq_lens, int max_seq_len,
52
+ const std::optional<torch::stable::Tensor>& alibi_slopes,
53
+ torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
54
+ const int tp_rank, const int blocksparse_local_blocks,
55
+ const int blocksparse_vert_stride, const int blocksparse_block_size,
56
+ const int blocksparse_head_sliding_step) {
57
int num_seqs = query.size(0);
58
int num_heads = query.size(1);
59
int head_size = query.size(2);
@@ -69,8 +72,8 @@ void paged_attention_v1_launcher(
72
T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
73
CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
74
CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
72
- int* block_tables_ptr = block_tables.data_ptr<int>();
73
- int* seq_lens_ptr = seq_lens.data_ptr<int>();
75
+ int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
76
+ int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
77
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
78
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
79
@@ -85,8 +88,9 @@ void paged_attention_v1_launcher(
88
89
dim3 grid(num_heads, num_seqs, 1);
90
dim3 block(NUM_THREADS);
88
- const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
89
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
91
+ const torch::stable::accelerator::DeviceGuard device_guard(
92
+ query.get_device_index());
93
+ const cudaStream_t stream = get_current_cuda_stream();
94
switch (head_size) {
95
// NOTE(woosuk): To reduce the compilation time, we only compile for the
96
// head sizes that we use in the model. However, we can easily extend this
@@ -119,7 +123,7 @@ void paged_attention_v1_launcher(
123
LAUNCH_PAGED_ATTENTION_V1(256);
124
break;
125
default:
122
- TORCH_CHECK(false, "Unsupported head size: ", head_size);
126
+ STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
127
break;
128
}
129
}
@@ -141,43 +145,43 @@ void paged_attention_v1_launcher(
145
146
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
147
// 1, 2, 4, 64, 128, 256.
144
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
145
- switch (block_size) { \
146
- case 8: \
147
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
148
- break; \
149
- case 16: \
150
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
151
- break; \
152
- case 32: \
153
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
154
- break; \
155
- default: \
156
- TORCH_CHECK(false, "Unsupported block size: ", block_size); \
157
- break; \
148
+#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
149
+ switch (block_size) { \
150
+ case 8: \
151
+ CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
152
+ break; \
153
+ case 16: \
154
+ CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
155
+ break; \
156
+ case 32: \
157
+ CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
158
+ break; \
159
+ default: \
160
+ STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
161
+ break; \
162
}
163
164
void paged_attention_v1(
161
- torch::Tensor& out, // [num_seqs, num_heads, head_size]
162
- torch::Tensor& query, // [num_seqs, num_heads, head_size]
163
- torch::Tensor&
165
+ torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
166
+ torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
167
+ torch::stable::Tensor&
168
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
165
- torch::Tensor&
169
+ torch::stable::Tensor&
170
value_cache, // [num_blocks, num_heads, head_size, block_size]
171
int64_t num_kv_heads, // [num_heads]
172
double scale,
169
- torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
170
- torch::Tensor& seq_lens, // [num_seqs]
173
+ torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
174
+ torch::stable::Tensor& seq_lens, // [num_seqs]
175
int64_t block_size, int64_t max_seq_len,
172
- const std::optional<torch::Tensor>& alibi_slopes,
173
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
174
- torch::Tensor& v_scale, const int64_t tp_rank,
176
+ const std::optional<torch::stable::Tensor>& alibi_slopes,
177
+ const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
178
+ torch::stable::Tensor& v_scale, const int64_t tp_rank,
179
const int64_t blocksparse_local_blocks,
180
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
181
const int64_t blocksparse_head_sliding_step) {
182
const bool is_block_sparse = (blocksparse_vert_stride > 1);
183
180
- DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
184
+ DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
185
CALL_V1_LAUNCHER_BLOCK_SIZE)
186
}
187
csrc/libtorch_stable/attention/paged_attention_v2.cu
renamed
+47
-41
@@ -16,8 +16,9 @@
16
* See the License for the specific language governing permissions and
17
* limitations under the License.
18
*/
19
+#include "../torch_utils.h"
20
#include "attention_kernels.cuh"
20
-#include "../cuda_compat.h"
21
+#include "../../cuda_compat.h"
22
23
#define MAX(a, b) ((a) > (b) ? (a) : (b))
24
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -44,14 +45,16 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
45
vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
46
int NUM_THREADS = 128, int PARTITION_SIZE = 512>
47
void paged_attention_v2_launcher(
47
- torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
48
- torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
49
- torch::Tensor& value_cache, int num_kv_heads, float scale,
50
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
51
- const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
52
- torch::Tensor& v_scale, const int tp_rank,
53
- const int blocksparse_local_blocks, const int blocksparse_vert_stride,
54
- const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
48
+ torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
49
+ torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
50
+ torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
51
+ torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
52
+ torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
53
+ int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
54
+ torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
55
+ const int tp_rank, const int blocksparse_local_blocks,
56
+ const int blocksparse_vert_stride, const int blocksparse_block_size,
57
+ const int blocksparse_head_sliding_step) {
58
int num_seqs = query.size(0);
59
int num_heads = query.size(1);
60
int head_size = query.size(2);
@@ -73,8 +76,8 @@ void paged_attention_v2_launcher(
76
T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
77
CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
78
CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
76
- int* block_tables_ptr = block_tables.data_ptr<int>();
77
- int* seq_lens_ptr = seq_lens.data_ptr<int>();
79
+ int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
80
+ int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
81
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
82
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
83
@@ -91,8 +94,9 @@ void paged_attention_v2_launcher(
94
int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
95
96
dim3 block(NUM_THREADS);
94
- const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
95
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
97
+ const torch::stable::accelerator::DeviceGuard device_guard(
98
+ query.get_device_index());
99
+ const cudaStream_t stream = get_current_cuda_stream();
100
switch (head_size) {
101
// NOTE(woosuk): To reduce the compilation time, we only compile for the
102
// head sizes that we use in the model. However, we can easily extend this
@@ -125,7 +129,7 @@ void paged_attention_v2_launcher(
129
LAUNCH_PAGED_ATTENTION_V2(256);
130
break;
131
default:
128
- TORCH_CHECK(false, "Unsupported head size: ", head_size);
132
+ STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
133
break;
134
}
135
}
@@ -148,46 +152,48 @@ void paged_attention_v2_launcher(
152
153
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
154
// 1, 2, 4, 64, 128, 256.
151
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
152
- switch (block_size) { \
153
- case 8: \
154
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
155
- break; \
156
- case 16: \
157
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
158
- break; \
159
- case 32: \
160
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
161
- break; \
162
- default: \
163
- TORCH_CHECK(false, "Unsupported block size: ", block_size); \
164
- break; \
155
+#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
156
+ switch (block_size) { \
157
+ case 8: \
158
+ CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
159
+ break; \
160
+ case 16: \
161
+ CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
162
+ break; \
163
+ case 32: \
164
+ CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
165
+ break; \
166
+ default: \
167
+ STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
168
+ break; \
169
}
170
171
void paged_attention_v2(
168
- torch::Tensor& out, // [num_seqs, num_heads, head_size]
169
- torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions]
170
- torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions]
171
- torch::Tensor&
172
+ torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
173
+ torch::stable::Tensor&
174
+ exp_sums, // [num_seqs, num_heads, max_num_partitions]
175
+ torch::stable::Tensor&
176
+ max_logits, // [num_seqs, num_heads, max_num_partitions]
177
+ torch::stable::Tensor&
178
tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
173
- torch::Tensor& query, // [num_seqs, num_heads, head_size]
174
- torch::Tensor&
179
+ torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
180
+ torch::stable::Tensor&
181
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
176
- torch::Tensor&
182
+ torch::stable::Tensor&
183
value_cache, // [num_blocks, num_heads, head_size, block_size]
184
int64_t num_kv_heads, // [num_heads]
185
double scale,
180
- torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
181
- torch::Tensor& seq_lens, // [num_seqs]
186
+ torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
187
+ torch::stable::Tensor& seq_lens, // [num_seqs]
188
int64_t block_size, int64_t max_seq_len,
183
- const std::optional<torch::Tensor>& alibi_slopes,
184
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
185
- torch::Tensor& v_scale, const int64_t tp_rank,
189
+ const std::optional<torch::stable::Tensor>& alibi_slopes,
190
+ const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
191
+ torch::stable::Tensor& v_scale, const int64_t tp_rank,
192
const int64_t blocksparse_local_blocks,
193
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
194
const int64_t blocksparse_head_sliding_step) {
195
const bool is_block_sparse = (blocksparse_vert_stride > 1);
190
- DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
196
+ DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
197
CALL_V2_LAUNCHER_BLOCK_SIZE)
198
}
199
csrc/libtorch_stable/cache_kernels.cu
renamed
+360
-311
@@ -1,20 +1,16 @@
1
-#include <torch/all.h>
2
-#include <ATen/cuda/CUDAContext.h>
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <c10/cuda/CUDAException.h>
5
-#include <c10/util/Optional.h>
6
-
7
-#include "cuda_utils.h"
8
-#include "cuda_compat.h"
1
+#include "torch_utils.h"
2
#include "dispatch_utils.h"
3
11
-#include "libtorch_stable/quantization/vectorization_utils.cuh"
4
+#include "../cuda_utils.h"
5
+#include "../cuda_compat.h"
6
+
7
+#include "quantization/vectorization_utils.cuh"
8
#include "concat_mla_q.cuh"
9
10
#ifdef USE_ROCM
15
- #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
11
+ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
12
#else
17
- #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
13
+ #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
14
#endif
15
16
#include <algorithm>
@@ -34,40 +30,45 @@ constexpr float kFp8ScaleDivisor = 224.f;
30
constexpr float kFp8ScaleDivisor = 448.f;
31
#endif
32
37
-void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
33
+void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
34
int64_t block_size_in_bytes,
39
- const torch::Tensor& block_mapping) {
40
- torch::Device src_device = src.device();
41
- torch::Device dst_device = dst.device();
35
+ const torch::stable::Tensor& block_mapping) {
36
+ torch::stable::Device src_device = src.device();
37
+ torch::stable::Device dst_device = dst.device();
38
cudaMemcpyKind memcpy_type;
39
if (src_device.is_cuda() && dst_device.is_cuda()) {
44
- TORCH_CHECK(src_device.index() == dst_device.index(),
45
- "src and dst must be on the same GPU");
40
+ STD_TORCH_CHECK(src_device.index() == dst_device.index(),
41
+ "src and dst must be on the same GPU");
42
memcpy_type = cudaMemcpyDeviceToDevice;
43
} else if (src_device.is_cuda() && dst_device.is_cpu()) {
44
memcpy_type = cudaMemcpyDeviceToHost;
45
} else if (src_device.is_cpu() && dst_device.is_cuda()) {
46
memcpy_type = cudaMemcpyHostToDevice;
47
} else {
52
- TORCH_CHECK(false, "Invalid device combination");
48
+ STD_TORCH_CHECK(false, "Invalid device combination");
49
}
50
51
// NOTE(youkaichao): keep in mind that `block_mapping` should be
52
// a cpu tensor, otherwise every `item` call will require a gpu-cpu
53
// synchronization.
58
- TORCH_CHECK(block_mapping.device().is_cpu(), "block_mapping must be on CPU");
54
+ STD_TORCH_CHECK(block_mapping.device().is_cpu(),
55
+ "block_mapping must be on CPU");
56
57
char* src_ptr = static_cast<char*>(src.data_ptr());
58
char* dst_ptr = static_cast<char*>(dst.data_ptr());
59
63
- const at::cuda::OptionalCUDAGuard device_guard(
64
- src_device.is_cuda() ? src_device : dst_device);
65
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
60
+ auto guard_device = src_device.is_cuda() ? src_device : dst_device;
61
+ const torch::stable::accelerator::DeviceGuard device_guard(
62
+ guard_device.index());
63
+ const cudaStream_t stream = get_current_cuda_stream();
64
// NOTE(woosuk): This can be slow if the number of blocks is large.
65
const int64_t num_blocks = block_mapping.size(0);
66
+ const int64_t* bm_ptr = block_mapping.const_data_ptr<int64_t>();
67
+ const int64_t bm_stride0 = block_mapping.stride(0);
68
+ const int64_t bm_stride1 = block_mapping.stride(1);
69
for (size_t i = 0; i < num_blocks; i++) {
69
- int64_t src_block_number = block_mapping[i][0].item<int64_t>();
70
- int64_t dst_block_number = block_mapping[i][1].item<int64_t>();
70
+ int64_t src_block_number = bm_ptr[i * bm_stride0];
71
+ int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1];
72
int64_t src_offset = src_block_number * block_size_in_bytes;
73
int64_t dst_offset = dst_block_number * block_size_in_bytes;
74
cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset,
@@ -75,20 +76,23 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
76
}
77
}
78
78
-void swap_blocks_batch(const torch::Tensor& src_ptrs,
79
- const torch::Tensor& dst_ptrs,
80
- const torch::Tensor& sizes,
79
+void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
80
+ const torch::stable::Tensor& dst_ptrs,
81
+ const torch::stable::Tensor& sizes,
82
bool is_src_access_order_any) {
82
- TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
83
- TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
84
- TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
85
- TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64");
86
- TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64");
87
- TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64");
83
+ STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
84
+ STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
85
+ STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
86
+ STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
87
+ "src_ptrs must be int64");
88
+ STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
89
+ "dst_ptrs must be int64");
90
+ STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long,
91
+ "sizes must be int64");
92
93
const int64_t n = src_ptrs.size(0);
90
- TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs");
91
- TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs");
94
+ STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs");
95
+ STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs");
96
97
if (n == 0) return;
98
@@ -96,7 +100,7 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
100
int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
101
int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
102
99
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
103
+ const cudaStream_t stream = get_current_cuda_stream();
104
105
// Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a
106
// single driver call, amortizing per-copy submission overhead. int64_t
@@ -138,8 +142,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
142
reinterpret_cast<size_t*>(size_data),
143
static_cast<size_t>(n), &attr, &attrs_idx, 1,
144
&fail_idx, static_cast<CUstream>(stream));
141
- TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
142
- fail_idx, " with error ", result);
145
+ STD_TORCH_CHECK(result == CUDA_SUCCESS,
146
+ "cuMemcpyBatchAsync failed at index ", fail_idx,
147
+ " with error ", result);
148
return;
149
}
150
#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
@@ -155,8 +160,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
160
reinterpret_cast<void**>(dst_data), reinterpret_cast<void**>(src_data),
161
reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
162
&attrs_idx, 0, &fail_idx, static_cast<hipStream_t>(stream));
158
- TORCH_CHECK(result == hipSuccess, "hipMemcpyBatchAsync failed at index ",
159
- fail_idx, " with error ", result);
163
+ STD_TORCH_CHECK(result == hipSuccess,
164
+ "hipMemcpyBatchAsync failed at index ", fail_idx,
165
+ " with error ", result);
166
return;
167
}
168
#endif
@@ -675,28 +681,28 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
681
// KV_T is the data type of key and value tensors.
682
// CACHE_T is the stored data type of kv-cache.
683
// KV_DTYPE is the real data type of kv-cache.
678
-#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
679
- vllm::reshape_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
680
- <<<grid, block, 0, stream>>>( \
681
- reinterpret_cast<KV_T*>(key.data_ptr()), \
682
- reinterpret_cast<KV_T*>(value.data_ptr()), \
683
- reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
684
- reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
685
- slot_mapping.data_ptr<int64_t>(), key_stride, value_stride, \
686
- num_heads, head_size, block_size, x, \
687
- reinterpret_cast<const float*>(k_scale.data_ptr()), \
684
+#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
685
+ vllm::reshape_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
686
+ <<<grid, block, 0, stream>>>( \
687
+ reinterpret_cast<KV_T*>(key.data_ptr()), \
688
+ reinterpret_cast<KV_T*>(value.data_ptr()), \
689
+ reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
690
+ reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
691
+ slot_mapping.const_data_ptr<int64_t>(), key_stride, value_stride, \
692
+ num_heads, head_size, block_size, x, \
693
+ reinterpret_cast<const float*>(k_scale.data_ptr()), \
694
reinterpret_cast<const float*>(v_scale.data_ptr()));
695
696
void reshape_and_cache(
691
- torch::Tensor& key, // [num_tokens, num_heads, head_size]
692
- torch::Tensor& value, // [num_tokens, num_heads, head_size]
693
- torch::Tensor&
697
+ torch::stable::Tensor& key, // [num_tokens, num_heads, head_size]
698
+ torch::stable::Tensor& value, // [num_tokens, num_heads, head_size]
699
+ torch::stable::Tensor&
700
key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
695
- torch::Tensor&
701
+ torch::stable::Tensor&
702
value_cache, // [num_blocks, num_heads, head_size, block_size]
697
- torch::Tensor& slot_mapping, // [num_tokens]
698
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
699
- torch::Tensor& v_scale) {
703
+ torch::stable::Tensor& slot_mapping, // [num_tokens]
704
+ const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
705
+ torch::stable::Tensor& v_scale) {
706
int num_tokens = slot_mapping.size(0);
707
int num_heads = key.size(1);
708
int head_size = key.size(2);
@@ -709,39 +715,41 @@ void reshape_and_cache(
715
716
dim3 grid(num_tokens);
717
dim3 block(std::min(num_heads * head_div_x, 512));
712
- const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
713
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
718
+ const torch::stable::accelerator::DeviceGuard device_guard(
719
+ key.get_device_index());
720
+ const cudaStream_t stream = get_current_cuda_stream();
721
715
- DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
722
+ DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
723
CALL_RESHAPE_AND_CACHE);
724
}
725
726
// KV_T is the data type of key and value tensors.
727
// CACHE_T is the stored data type of kv-cache.
728
// KV_DTYPE is the real data type of kv-cache.
722
-#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \
723
- vllm::reshape_and_cache_flash_kernel<KV_T, CACHE_T, KV_DTYPE> \
724
- <<<grid, block, 0, stream>>>( \
725
- reinterpret_cast<KV_T*>(key.data_ptr()), \
726
- reinterpret_cast<KV_T*>(value.data_ptr()), \
727
- reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
728
- reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
729
- slot_mapping.data_ptr<int64_t>(), block_stride, page_stride, \
730
- head_stride, key_stride, value_stride, num_heads, head_size, \
731
- block_size, reinterpret_cast<const float*>(k_scale.data_ptr()), \
732
- reinterpret_cast<const float*>(v_scale.data_ptr()), \
729
+#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \
730
+ vllm::reshape_and_cache_flash_kernel<KV_T, CACHE_T, KV_DTYPE> \
731
+ <<<grid, block, 0, stream>>>( \
732
+ reinterpret_cast<KV_T*>(key.data_ptr()), \
733
+ reinterpret_cast<KV_T*>(value.data_ptr()), \
734
+ reinterpret_cast<CACHE_T*>(key_cache.data_ptr()), \
735
+ reinterpret_cast<CACHE_T*>(value_cache.data_ptr()), \
736
+ slot_mapping.const_data_ptr<int64_t>(), block_stride, page_stride, \
737
+ head_stride, key_stride, value_stride, num_heads, head_size, \
738
+ block_size, reinterpret_cast<const float*>(k_scale.data_ptr()), \
739
+ reinterpret_cast<const float*>(v_scale.data_ptr()), \
740
kv_scale_stride);
741
742
void reshape_and_cache_flash(
736
- torch::Tensor& key, // [num_tokens, num_heads, head_size]
737
- torch::Tensor& value, // [num_tokens, num_heads, head_size]
738
- torch::Tensor& key_cache, // [num_blocks, block_size, num_heads, head_size]
739
- torch::Tensor&
743
+ torch::stable::Tensor& key, // [num_tokens, num_heads, head_size]
744
+ torch::stable::Tensor& value, // [num_tokens, num_heads, head_size]
745
+ torch::stable::Tensor&
746
+ key_cache, // [num_blocks, block_size, num_heads, head_size]
747
+ torch::stable::Tensor&
748
value_cache, // [num_blocks, block_size, num_heads, head_size]
741
- torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
749
+ torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
750
const std::string& kv_cache_dtype,
743
- torch::Tensor& k_scale, // [1] or [num_heads]
744
- torch::Tensor& v_scale) { // [1] or [num_heads]
751
+ torch::stable::Tensor& k_scale, // [1] or [num_heads]
752
+ torch::stable::Tensor& v_scale) { // [1] or [num_heads]
753
// NOTE(woosuk): In vLLM V1, key.size(0) can be different from
754
// slot_mapping.size(0) because of padding for CUDA graphs.
755
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -756,23 +764,26 @@ void reshape_and_cache_flash(
764
int num_heads = key.size(1);
765
int head_size = key.size(2);
766
759
- const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
760
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
767
+ const torch::stable::accelerator::DeviceGuard device_guard(
768
+ key.get_device_index());
769
+ const cudaStream_t stream = get_current_cuda_stream();
770
771
if (kv_cache_dtype == "nvfp4") {
772
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
773
// NVFP4 dispatch is compiled separately for SM100+.
774
extern void reshape_and_cache_nvfp4_dispatch(
766
- torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
767
- torch::Tensor & value_cache, torch::Tensor & slot_mapping,
768
- torch::Tensor & k_scale, torch::Tensor & v_scale);
775
+ torch::stable::Tensor & key, torch::stable::Tensor & value,
776
+ torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache,
777
+ torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale,
778
+ torch::stable::Tensor & v_scale);
779
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
780
slot_mapping, k_scale, v_scale);
781
return;
782
#else
773
- TORCH_CHECK(false,
774
- "NVFP4 KV cache requires SM100+ (Blackwell). "
775
- "Please rebuild vllm with a Blackwell-compatible CUDA target.");
783
+ STD_TORCH_CHECK(
784
+ false,
785
+ "NVFP4 KV cache requires SM100+ (Blackwell). "
786
+ "Please rebuild vllm with a Blackwell-compatible CUDA target.");
787
#endif
788
}
789
@@ -784,53 +795,53 @@ void reshape_and_cache_flash(
795
int64_t block_stride = key_cache.stride(0);
796
int64_t page_stride = key_cache.stride(1);
797
int64_t head_stride = key_cache.stride(2);
787
- TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
798
+ STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
799
789
- TORCH_CHECK(k_scale.sizes() == v_scale.sizes(),
790
- "k_scale and v_scale must have the same shape");
791
- TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
792
- "k_scale and v_scale must be of shape [1] or [num_heads]");
800
+ STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()),
801
+ "k_scale and v_scale must have the same shape");
802
+ STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
803
+ "k_scale and v_scale must be of shape [1] or [num_heads]");
804
int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0;
805
806
dim3 grid(num_tokens);
807
dim3 block(std::min(num_heads * head_size, 512));
808
798
- DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
809
+ DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype,
810
CALL_RESHAPE_AND_CACHE_FLASH);
811
}
812
813
// KV_T is the data type of key and value tensors.
814
// CACHE_T is the stored data type of kv-cache.
815
// KV_DTYPE is the real data type of kv-cache.
805
-#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
806
- vllm::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
807
- <<<grid, block, 0, stream>>>( \
808
- reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
809
- reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
810
- reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
811
- slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
812
- kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
816
+#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
817
+ vllm::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
818
+ <<<grid, block, 0, stream>>>( \
819
+ reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
820
+ reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
821
+ reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
822
+ slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride, \
823
+ kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
824
reinterpret_cast<const float*>(scale.data_ptr()));
825
826
// KV_T is the data type of key and value tensors.
827
// CACHE_T is the stored data type of kv-cache.
817
-#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \
818
- vllm::concat_and_cache_ds_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
819
- <<<grid, block, 0, stream>>>( \
820
- reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
821
- reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
822
- reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
823
- slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
824
- kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
828
+#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \
829
+ vllm::concat_and_cache_ds_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
830
+ <<<grid, block, 0, stream>>>( \
831
+ reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
832
+ reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
833
+ reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
834
+ slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride, \
835
+ kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
836
reinterpret_cast<const float*>(scale.data_ptr()));
837
838
void concat_and_cache_mla(
828
- torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
829
- torch::Tensor& k_pe, // [num_tokens, pe_dim]
830
- torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank +
831
- // pe_dim)]
832
- torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
833
- const std::string& kv_cache_dtype, torch::Tensor& scale) {
839
+ torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank]
840
+ torch::stable::Tensor& k_pe, // [num_tokens, pe_dim]
841
+ torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank
842
+ // + pe_dim)]
843
+ torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
844
+ const std::string& kv_cache_dtype, torch::stable::Tensor& scale) {
845
// NOTE(woosuk): In vLLM V1, key.size(0) can be different from
846
// slot_mapping.size(0) because of padding for CUDA graphs.
847
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -847,16 +858,17 @@ void concat_and_cache_mla(
858
int block_size = kv_cache.size(1);
859
860
if (kv_cache_dtype == "fp8_ds_mla") {
850
- TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for fp8_ds_mla");
851
- TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla");
852
- TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.itemsize(),
853
- "kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
854
- TORCH_CHECK(kv_c.itemsize() == 2,
855
- "kv_c.itemsize() must be 2 for fp8_ds_mla");
856
- TORCH_CHECK(k_pe.itemsize() == 2,
857
- "k_pe.itemsize() must be 2 for fp8_ds_mla");
861
+ STD_TORCH_CHECK(kv_lora_rank == 512,
862
+ "kv_lora_rank must be 512 for fp8_ds_mla");
863
+ STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla");
864
+ STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(),
865
+ "kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
866
+ STD_TORCH_CHECK(kv_c.element_size() == 2,
867
+ "kv_c.element_size() must be 2 for fp8_ds_mla");
868
+ STD_TORCH_CHECK(k_pe.element_size() == 2,
869
+ "k_pe.element_size() must be 2 for fp8_ds_mla");
870
} else {
859
- TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
871
+ STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
872
}
873
874
int kv_c_stride = kv_c.stride(0);
@@ -864,8 +876,9 @@ void concat_and_cache_mla(
876
int block_stride = kv_cache.stride(0);
877
int entry_stride = kv_cache.stride(1);
878
867
- const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c));
868
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
879
+ const torch::stable::accelerator::DeviceGuard device_guard(
880
+ kv_c.get_device_index());
881
+ const cudaStream_t stream = get_current_cuda_stream();
882
883
if (kv_cache_dtype == "fp8_ds_mla") {
884
dim3 grid(num_tokens);
@@ -875,12 +888,12 @@ void concat_and_cache_mla(
888
// The RoPE part (last 64 elements) is handled by another 1 warp (32
889
// threads). So in total, we use 3 warps (96 threads) per block.
890
dim3 block(96);
878
- DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
891
+ DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
892
CALL_CONCAT_AND_CACHE_DS_MLA);
893
} else {
894
dim3 grid(num_tokens);
895
dim3 block(std::min(kv_lora_rank, 512));
883
- DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
896
+ DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
897
CALL_CONCAT_AND_CACHE_MLA);
898
}
899
}
@@ -908,55 +921,62 @@ __global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache,
921
reinterpret_cast<Tout*>(dst_cache.data_ptr()), scale, block_stride);
922
923
// Only for testing.
911
-void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
912
- const double scale, const std::string& kv_cache_dtype) {
913
- torch::Device src_device = src_cache.device();
914
- torch::Device dst_device = dst_cache.device();
915
- TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU")
916
- TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU")
917
- TORCH_CHECK(src_device.index() == dst_device.index(),
918
- "src and dst must be on the same GPU");
919
- at::cuda::OptionalCUDAGuard device_guard(src_device);
924
+void convert_fp8(torch::stable::Tensor& dst_cache,
925
+ torch::stable::Tensor& src_cache, const double scale,
926
+ const std::string& kv_cache_dtype) {
927
+ torch::stable::Device src_device = src_cache.device();
928
+ torch::stable::Device dst_device = dst_cache.device();
929
+ STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU")
930
+ STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU")
931
+ STD_TORCH_CHECK(src_device.index() == dst_device.index(),
932
+ "src and dst must be on the same GPU");
933
+ torch::stable::accelerator::DeviceGuard device_guard(src_device.index());
934
935
int64_t num_blocks = src_cache.size(0);
936
int64_t block_stride = src_cache.stride(0);
937
938
dim3 grid(num_blocks);
939
dim3 block(std::min(block_stride, int64_t(512)));
926
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
940
+ const cudaStream_t stream = get_current_cuda_stream();
941
942
if (kv_cache_dtype == "auto") {
929
- if (src_cache.dtype() == at::ScalarType::Float) {
943
+ if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) {
944
CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kAuto);
931
- } else if (src_cache.dtype() == at::ScalarType::Half) {
945
+ } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
946
CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);
933
- } else if (src_cache.dtype() == at::ScalarType::BFloat16) {
947
+ } else if (src_cache.scalar_type() ==
948
+ torch::headeronly::ScalarType::BFloat16) {
949
CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);
935
- } else if (dst_cache.dtype() == at::ScalarType::Float) {
950
+ } else if (dst_cache.scalar_type() ==
951
+ torch::headeronly::ScalarType::Float) {
952
CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
937
- } else if (dst_cache.dtype() == at::ScalarType::Half) {
953
+ } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
954
CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
939
- } else if (dst_cache.dtype() == at::ScalarType::BFloat16) {
955
+ } else if (dst_cache.scalar_type() ==
956
+ torch::headeronly::ScalarType::BFloat16) {
957
CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kAuto);
958
}
959
} else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
943
- if (src_cache.dtype() == at::ScalarType::Float) {
960
+ if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) {
961
CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kFp8E4M3);
945
- } else if (src_cache.dtype() == at::ScalarType::Half) {
962
+ } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
963
CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
947
- } else if (src_cache.dtype() == at::ScalarType::BFloat16) {
964
+ } else if (src_cache.scalar_type() ==
965
+ torch::headeronly::ScalarType::BFloat16) {
966
CALL_CONVERT_FP8(uint8_t, __nv_bfloat16,
967
vllm::Fp8KVCacheDataType::kFp8E4M3);
950
- } else if (dst_cache.dtype() == at::ScalarType::Float) {
968
+ } else if (dst_cache.scalar_type() ==
969
+ torch::headeronly::ScalarType::Float) {
970
CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
952
- } else if (dst_cache.dtype() == at::ScalarType::Half) {
971
+ } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) {
972
CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);
954
- } else if (dst_cache.dtype() == at::ScalarType::BFloat16) {
973
+ } else if (dst_cache.scalar_type() ==
974
+ torch::headeronly::ScalarType::BFloat16) {
975
CALL_CONVERT_FP8(__nv_bfloat16, uint8_t,
976
vllm::Fp8KVCacheDataType::kFp8E4M3);
977
}
978
} else {
959
- TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype);
979
+ STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype);
980
}
981
}
982
@@ -1053,8 +1073,9 @@ __global__ void gather_and_maybe_dequant_cache(
1073
<<<grid, block, 0, stream>>>( \
1074
reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
1075
reinterpret_cast<SCALAR_T*>(dst.data_ptr()), \
1056
- block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
1057
- token_to_seq.data_ptr<int32_t>(), num_tokens, block_size, \
1076
+ block_table.const_data_ptr<int32_t>(), \
1077
+ cu_seq_lens.const_data_ptr<int32_t>(), \
1078
+ token_to_seq.const_data_ptr<int32_t>(), num_tokens, block_size, \
1079
block_table_stride, cache_block_stride, cache_entry_stride, \
1080
dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
1081
seq_starts_ptr);
@@ -1072,42 +1093,47 @@ __global__ void gather_and_maybe_dequant_cache(
1093
// - Optionally, seq_starts (if provided) offsets the starting block index by
1094
// (seq_starts[bid] / page_size)
1095
void gather_and_maybe_dequant_cache(
1075
- torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1076
- torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1077
- torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1078
- torch::Tensor const& cu_seq_lens, // [BATCH+1]
1079
- torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
1096
+ torch::stable::Tensor const&
1097
+ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1098
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1099
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1100
+ torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
1101
+ torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
1102
int64_t num_tokens, const std::string& kv_cache_dtype,
1081
- torch::Tensor const& scale,
1082
- std::optional<torch::Tensor> seq_starts = std::nullopt) {
1083
- at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
1084
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1103
+ torch::stable::Tensor const& scale,
1104
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
1105
+ torch::stable::accelerator::DeviceGuard device_guard(
1106
+ src_cache.get_device_index());
1107
+ const cudaStream_t stream = get_current_cuda_stream();
1108
1109
int32_t block_size = src_cache.size(1);
1110
int32_t head_dim = dst.size(-1);
1111
1089
- TORCH_CHECK(block_table.dtype() == torch::kInt32,
1090
- "block_table must be int32");
1091
- TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32,
1092
- "cu_seq_lens must be int32");
1112
+ STD_TORCH_CHECK(
1113
+ block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1114
+ "block_table must be int32");
1115
+ STD_TORCH_CHECK(
1116
+ cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
1117
+ "cu_seq_lens must be int32");
1118
if (seq_starts.has_value()) {
1094
- TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
1095
- "seq_starts must be int32");
1119
+ STD_TORCH_CHECK(
1120
+ seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1121
+ "seq_starts must be int32");
1122
}
1097
- TORCH_CHECK(
1123
+ STD_TORCH_CHECK(
1124
head_dim == 320 || head_dim == 576,
1125
"gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 "
1126
"for better performance")
1127
1102
- TORCH_CHECK(src_cache.device() == dst.device(),
1103
- "src_cache and dst must be on the same device");
1104
- TORCH_CHECK(src_cache.device() == block_table.device(),
1105
- "src_cache and block_table must be on the same device");
1106
- TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1107
- "src_cache and cu_seq_lens must be on the same device");
1128
+ STD_TORCH_CHECK(src_cache.device() == dst.device(),
1129
+ "src_cache and dst must be on the same device");
1130
+ STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1131
+ "src_cache and block_table must be on the same device");
1132
+ STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1133
+ "src_cache and cu_seq_lens must be on the same device");
1134
if (seq_starts.has_value()) {
1109
- TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1110
- "src_cache and seq_starts must be on the same device");
1135
+ STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1136
+ "src_cache and seq_starts must be on the same device");
1137
}
1138
1139
int64_t block_table_stride = block_table.stride(0);
@@ -1120,13 +1146,14 @@ void gather_and_maybe_dequant_cache(
1146
dim3 block(thread_block_size);
1147
1148
const int32_t* seq_starts_ptr =
1123
- seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
1149
+ seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1150
+ : nullptr;
1151
1152
if (head_dim == 576) {
1126
- DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
1153
+ DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1154
CALL_GATHER_CACHE_576);
1155
} else {
1129
- DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
1156
+ DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1157
CALL_GATHER_CACHE_320);
1158
}
1159
}
@@ -1267,13 +1294,14 @@ __global__ void cp_gather_cache(
1294
} // namespace vllm
1295
1296
// Macro to dispatch the kernel based on the data type.
1270
-#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \
1271
- vllm::cp_gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \
1272
- reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \
1273
- reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \
1274
- block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
1275
- block_size, entry_size, block_table_stride, cache_block_stride, \
1276
- cache_entry_stride, dst_entry_stride, seq_starts_ptr);
1297
+#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \
1298
+ vllm::cp_gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>( \
1299
+ reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()), \
1300
+ reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()), \
1301
+ block_table.const_data_ptr<int32_t>(), \
1302
+ cu_seq_lens.const_data_ptr<int32_t>(), block_size, entry_size, \
1303
+ block_table_stride, cache_block_stride, cache_entry_stride, \
1304
+ dst_entry_stride, seq_starts_ptr);
1305
1306
// Gather sequences from the cache into the destination tensor.
1307
// - cu_seq_lens contains the cumulative sequence lengths for each batch
@@ -1281,36 +1309,41 @@ __global__ void cp_gather_cache(
1309
// - Optionally, seq_starts (if provided) offsets the starting slot index by
1310
// seq_starts[bid]
1311
void cp_gather_cache(
1284
- torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1285
- torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1286
- torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1287
- torch::Tensor const& cu_seq_lens, // [BATCH+1]
1312
+ torch::stable::Tensor const&
1313
+ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
1314
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
1315
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1316
+ torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
1317
int64_t batch_size,
1289
- std::optional<torch::Tensor> seq_starts = std::nullopt) {
1290
- at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
1291
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1318
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
1319
+ torch::stable::accelerator::DeviceGuard device_guard(
1320
+ src_cache.get_device_index());
1321
+ const cudaStream_t stream = get_current_cuda_stream();
1322
1323
int32_t block_size = src_cache.size(1);
1294
- int32_t entry_size = src_cache.flatten(2, -1).size(2);
1295
-
1296
- TORCH_CHECK(block_table.dtype() == torch::kInt32,
1297
- "block_table must be int32");
1298
- TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32,
1299
- "cu_seq_lens must be int32");
1324
+ int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2);
1325
+
1326
+ STD_TORCH_CHECK(
1327
+ block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1328
+ "block_table must be int32");
1329
+ STD_TORCH_CHECK(
1330
+ cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
1331
+ "cu_seq_lens must be int32");
1332
if (seq_starts.has_value()) {
1301
- TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
1302
- "seq_starts must be int32");
1333
+ STD_TORCH_CHECK(
1334
+ seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1335
+ "seq_starts must be int32");
1336
}
1337
1305
- TORCH_CHECK(src_cache.device() == dst.device(),
1306
- "src_cache and dst must be on the same device");
1307
- TORCH_CHECK(src_cache.device() == block_table.device(),
1308
- "src_cache and block_table must be on the same device");
1309
- TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1310
- "src_cache and cu_seq_lens must be on the same device");
1338
+ STD_TORCH_CHECK(src_cache.device() == dst.device(),
1339
+ "src_cache and dst must be on the same device");
1340
+ STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1341
+ "src_cache and block_table must be on the same device");
1342
+ STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
1343
+ "src_cache and cu_seq_lens must be on the same device");
1344
if (seq_starts.has_value()) {
1312
- TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1313
- "src_cache and seq_starts must be on the same device");
1345
+ STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
1346
+ "src_cache and seq_starts must be on the same device");
1347
}
1348
1349
int64_t block_table_stride = block_table.stride(0);
@@ -1323,12 +1356,13 @@ void cp_gather_cache(
1356
dim3 grid(batch_size, num_splits);
1357
dim3 block(1024);
1358
1326
- TORCH_CHECK(src_cache.dtype() == dst.dtype(),
1327
- "src_cache and dst must have the same dtype");
1359
+ STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(),
1360
+ "src_cache and dst must have the same dtype");
1361
1362
const int dtype_bits = src_cache.element_size() * 8;
1363
const int32_t* seq_starts_ptr =
1331
- seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
1364
+ seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
1365
+ : nullptr;
1366
1367
if (dtype_bits == 32) {
1368
CALL_CP_GATHER_CACHE(uint32_t);
@@ -1337,46 +1371,51 @@ void cp_gather_cache(
1371
} else if (dtype_bits == 8) {
1372
CALL_CP_GATHER_CACHE(uint8_t);
1373
} else {
1340
- TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
1374
+ STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
1375
}
1376
}
1377
1378
void cp_gather_and_upconvert_fp8_kv_cache(
1345
- torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
1346
- torch::Tensor const& dst, // [TOT_TOKENS, 576]
1347
- torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1348
- torch::Tensor const& seq_lens, // [BATCH]
1349
- torch::Tensor const& workspace_starts, // [BATCH]
1379
+ torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
1380
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
1381
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1382
+ torch::stable::Tensor const& seq_lens, // [BATCH]
1383
+ torch::stable::Tensor const& workspace_starts, // [BATCH]
1384
int64_t batch_size) {
1351
- at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
1352
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1385
+ torch::stable::accelerator::DeviceGuard device_guard(
1386
+ src_cache.get_device_index());
1387
+ const cudaStream_t stream = get_current_cuda_stream();
1388
1389
int32_t block_size = src_cache.size(1);
1390
int32_t head_dim = dst.size(1);
1391
1357
- TORCH_CHECK(block_table.dtype() == torch::kInt32,
1358
- "block_table must be int32");
1359
- TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32");
1360
- TORCH_CHECK(workspace_starts.dtype() == torch::kInt32,
1361
- "workspace_starts must be int32");
1362
-
1363
- TORCH_CHECK(src_cache.device() == dst.device(),
1364
- "src_cache and dst must be on the same device");
1365
- TORCH_CHECK(src_cache.device() == block_table.device(),
1366
- "src_cache and block_table must be on the same device");
1367
- TORCH_CHECK(src_cache.device() == seq_lens.device(),
1368
- "src_cache and seq_lens must be on the same device");
1369
- TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1370
- "src_cache and workspace_starts must be on the same device");
1392
+ STD_TORCH_CHECK(
1393
+ block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1394
+ "block_table must be int32");
1395
+ STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
1396
+ "seq_lens must be int32");
1397
+ STD_TORCH_CHECK(
1398
+ workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
1399
+ "workspace_starts must be int32");
1400
+
1401
+ STD_TORCH_CHECK(src_cache.device() == dst.device(),
1402
+ "src_cache and dst must be on the same device");
1403
+ STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1404
+ "src_cache and block_table must be on the same device");
1405
+ STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),
1406
+ "src_cache and seq_lens must be on the same device");
1407
+ STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1408
+ "src_cache and workspace_starts must be on the same device");
1409
auto dtype = src_cache.scalar_type();
1372
- TORCH_CHECK(
1373
- dtype == at::ScalarType::Byte || // uint8
1374
- dtype == at::ScalarType::Float8_e4m3fn || // fp8 e4m3
1375
- dtype == at::ScalarType::Float8_e5m2, // fp8 e5m2
1410
+ STD_TORCH_CHECK(
1411
+ dtype == torch::headeronly::ScalarType::Byte || // uint8
1412
+ dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3
1413
+ dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2
1414
"src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
1377
- src_cache.dtype());
1378
- TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");
1379
- TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1415
+ src_cache.scalar_type());
1416
+ STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1417
+ "dst must be bfloat16");
1418
+ STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1419
1420
int64_t block_table_stride = block_table.stride(0);
1421
int64_t cache_block_stride = src_cache.stride(0);
@@ -1384,8 +1423,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1423
int64_t dst_entry_stride = dst.stride(0);
1424
1425
const uint8_t* src_ptr = nullptr;
1387
- if (dtype == at::ScalarType::Byte) {
1388
- src_ptr = src_cache.data_ptr<uint8_t>();
1426
+ if (dtype == torch::headeronly::ScalarType::Byte) {
1427
+ src_ptr = src_cache.const_data_ptr<uint8_t>();
1428
} else {
1429
// float8_e4m3fn or float8_e5m2
1430
src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
@@ -1399,26 +1438,27 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1438
vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
1439
stream>>>(
1440
src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1402
- block_table.data_ptr<int32_t>(), workspace_starts.data_ptr<int32_t>(),
1441
+ block_table.const_data_ptr<int32_t>(),
1442
+ workspace_starts.const_data_ptr<int32_t>(),
1443
static_cast<int32_t>(batch_size), block_size, total_tokens,
1444
block_table_stride, cache_block_stride, cache_entry_stride,
1445
dst_entry_stride);
1446
}
1447
1448
// Macro to dispatch the kernel based on the data type.
1409
-#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
1410
- vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
1411
- <<<grid, block, 0, stream>>>( \
1412
- reinterpret_cast<KV_T*>(k.data_ptr()), \
1413
- reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
1414
- slot_mapping.data_ptr<int64_t>(), head_dim, quant_block_size, \
1449
+#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
1450
+ vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
1451
+ <<<grid, block, 0, stream>>>( \
1452
+ reinterpret_cast<KV_T*>(k.data_ptr()), \
1453
+ reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
1454
+ slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
1455
cache_block_size, cache_stride, use_ue8m0);
1456
1457
void indexer_k_quant_and_cache(
1418
- torch::Tensor& k, // [num_tokens, head_dim]
1419
- torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1420
- torch::Tensor& slot_mapping, // [num_tokens]
1421
- int64_t quant_block_size, // quantization block size
1458
+ torch::stable::Tensor& k, // [num_tokens, head_dim]
1459
+ torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1460
+ torch::stable::Tensor& slot_mapping, // [num_tokens]
1461
+ int64_t quant_block_size, // quantization block size
1462
const std::string& scale_fmt) {
1463
int num_tokens = k.size(0);
1464
int head_dim = k.size(1);
@@ -1426,65 +1466,70 @@ void indexer_k_quant_and_cache(
1466
int cache_stride = kv_cache.size(2);
1467
bool use_ue8m0 = scale_fmt == "ue8m0";
1468
1429
- TORCH_CHECK(k.device() == kv_cache.device(),
1430
- "k and kv_cache must be on the same device");
1431
- TORCH_CHECK(k.device() == slot_mapping.device(),
1432
- "k and slot_mapping must be on the same device");
1433
- TORCH_CHECK(head_dim % quant_block_size == 0,
1434
- "head_dim must be divisible by quant_block_size");
1469
+ STD_TORCH_CHECK(k.device() == kv_cache.device(),
1470
+ "k and kv_cache must be on the same device");
1471
+ STD_TORCH_CHECK(k.device() == slot_mapping.device(),
1472
+ "k and slot_mapping must be on the same device");
1473
+ STD_TORCH_CHECK(head_dim % quant_block_size == 0,
1474
+ "head_dim must be divisible by quant_block_size");
1475
1476
constexpr int vec_size = 4;
1477
dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) /
1478
(quant_block_size * vec_size));
1479
dim3 block(32, vec_size);
1440
- const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
1441
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1480
+ const torch::stable::accelerator::DeviceGuard device_guard(
1481
+ k.get_device_index());
1482
+ const cudaStream_t stream = get_current_cuda_stream();
1483
1484
static const std::string kv_cache_dtype = "fp8_e4m3";
1444
- DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
1485
+ DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype,
1486
CALL_INDEXER_K_QUANT_AND_CACHE);
1487
}
1488
1489
// Macro to dispatch the kernel based on the data amount.
1449
-#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \
1450
- vllm::cp_gather_indexer_k_quant_cache_kernel<BLOCK_Y_SIZE> \
1451
- <<<dim3((num_tokens + BLOCK_Y_SIZE - 1) / BLOCK_Y_SIZE, \
1452
- (head_dim + 8 * vec_size - 1) / (8 * vec_size)), \
1453
- dim3(8, BLOCK_Y_SIZE), 0, stream>>>( \
1454
- reinterpret_cast<char*>(kv_cache.data_ptr()), \
1455
- reinterpret_cast<char*>(dst_k.data_ptr()), \
1456
- reinterpret_cast<char*>(dst_scale.data_ptr()), \
1457
- block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
1458
- batch_size, dst_k.stride(0), dst_k.size(1), kv_cache.stride(0), \
1459
- kv_cache.stride(1), kv_cache.size(1), block_table.size(1), \
1460
- num_tokens, quant_block_size);
1490
+#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \
1491
+ vllm::cp_gather_indexer_k_quant_cache_kernel<BLOCK_Y_SIZE> \
1492
+ <<<dim3((num_tokens + BLOCK_Y_SIZE - 1) / BLOCK_Y_SIZE, \
1493
+ (head_dim + 8 * vec_size - 1) / (8 * vec_size)), \
1494
+ dim3(8, BLOCK_Y_SIZE), 0, stream>>>( \
1495
+ reinterpret_cast<char*>(kv_cache.data_ptr()), \
1496
+ reinterpret_cast<char*>(dst_k.data_ptr()), \
1497
+ reinterpret_cast<char*>(dst_scale.data_ptr()), \
1498
+ block_table.const_data_ptr<int32_t>(), \
1499
+ cu_seq_lens.const_data_ptr<int32_t>(), batch_size, dst_k.stride(0), \
1500
+ dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \
1501
+ kv_cache.size(1), block_table.size(1), num_tokens, \
1502
+ quant_block_size);
1503
1504
void cp_gather_indexer_k_quant_cache(
1463
- const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
1464
- torch::Tensor& dst_k, // [num_tokens, head_dim]
1465
- torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
1466
- const torch::Tensor& block_table, // [batch_size, num_blocks]
1467
- const torch::Tensor& cu_seq_lens // [batch_size + 1]
1505
+ const torch::stable::Tensor&
1506
+ kv_cache, // [num_blocks, block_size, cache_stride]
1507
+ torch::stable::Tensor& dst_k, // [num_tokens, head_dim]
1508
+ torch::stable::Tensor&
1509
+ dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
1510
+ const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
1511
+ const torch::stable::Tensor& cu_seq_lens // [batch_size + 1]
1512
) {
1513
int batch_size = block_table.size(0);
1514
int num_tokens = dst_k.size(0);
1515
int head_dim = dst_k.size(1);
1516
int quant_block_size = head_dim * 4 / dst_scale.size(1);
1517
1474
- TORCH_CHECK(kv_cache.device() == dst_k.device(),
1475
- "kv_cache and dst_k must be on the same device");
1476
- TORCH_CHECK(kv_cache.device() == dst_scale.device(),
1477
- "kv_cache and dst_scale must be on the same device");
1478
- TORCH_CHECK(kv_cache.device() == block_table.device(),
1479
- "kv_cache and block_table must be on the same device");
1480
- TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(),
1481
- "kv_cache and cu_seq_lens must be on the same device");
1482
- TORCH_CHECK(head_dim % quant_block_size == 0,
1483
- "head_dim must be divisible by quant_block_size");
1518
+ STD_TORCH_CHECK(kv_cache.device() == dst_k.device(),
1519
+ "kv_cache and dst_k must be on the same device");
1520
+ STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(),
1521
+ "kv_cache and dst_scale must be on the same device");
1522
+ STD_TORCH_CHECK(kv_cache.device() == block_table.device(),
1523
+ "kv_cache and block_table must be on the same device");
1524
+ STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(),
1525
+ "kv_cache and cu_seq_lens must be on the same device");
1526
+ STD_TORCH_CHECK(head_dim % quant_block_size == 0,
1527
+ "head_dim must be divisible by quant_block_size");
1528
1529
constexpr int vec_size = 16;
1486
- const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_cache));
1487
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1530
+ const torch::stable::accelerator::DeviceGuard device_guard(
1531
+ kv_cache.get_device_index());
1532
+ const cudaStream_t stream = get_current_cuda_stream();
1533
1534
if (num_tokens < 32) {
1535
CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1);
@@ -1503,27 +1548,30 @@ void cp_gather_indexer_k_quant_cache(
1548
1549
// Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA.
1550
// Replaces torch.cat((ql_nope, q_pe), dim=-1).
1506
-void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1507
- torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
1508
- torch::Tensor& q_out // [num_tokens, num_heads, nope_dim +
1509
- // rope_dim]
1551
+void concat_mla_q(
1552
+ torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1553
+ torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
1554
+ torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim +
1555
+ // rope_dim]
1556
) {
1557
const int num_tokens = ql_nope.size(0);
1558
const int num_heads = ql_nope.size(1);
1559
const int nope_dim = ql_nope.size(2);
1560
const int rope_dim = q_pe.size(2);
1561
1516
- TORCH_CHECK(nope_dim % 512 == 0, "nope_dim must be a multiple of 512, got ",
1517
- nope_dim);
1518
- TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim);
1519
- TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim);
1562
+ STD_TORCH_CHECK(nope_dim % 512 == 0,
1563
+ "nope_dim must be a multiple of 512, got ", nope_dim);
1564
+ STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim);
1565
+ STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim);
1566
1521
- TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
1522
- TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
1523
- TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
1524
- TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
1525
- ql_nope.scalar_type() == at::ScalarType::BFloat16,
1526
- "ql_nope must be float16 or bfloat16 dtype");
1567
+ STD_TORCH_CHECK(ql_nope.stride(2) == 1,
1568
+ "ql_nope must have stride 1 in dim 2");
1569
+ STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
1570
+ STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
1571
+ STD_TORCH_CHECK(
1572
+ ql_nope.scalar_type() == torch::headeronly::ScalarType::Half ||
1573
+ ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1574
+ "ql_nope must be float16 or bfloat16 dtype");
1575
1576
if (num_tokens == 0) return;
1577
@@ -1532,13 +1580,14 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
1580
const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block;
1581
const int block_size = warps_per_block * 32;
1582
1535
- const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
1536
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1583
+ const torch::stable::accelerator::DeviceGuard device_guard(
1584
+ ql_nope.get_device_index());
1585
+ const cudaStream_t stream = get_current_cuda_stream();
1586
1538
- VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
1587
+ VLLM_STABLE_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
1588
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
1540
- q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
1541
- q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
1589
+ q_out.mutable_data_ptr<scalar_t>(), ql_nope.const_data_ptr<scalar_t>(),
1590
+ q_pe.const_data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
1591
q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0),
1592
q_pe.stride(1));
1593
});
csrc/libtorch_stable/cache_kernels_fused.cu
renamed
+102
-89
@@ -1,15 +1,13 @@
1
-#include <torch/all.h>
2
-#include <ATen/cuda/CUDAContext.h>
3
-#include <c10/cuda/CUDAGuard.h>
4
-
5
-#include "cuda_compat.h"
1
+#include "torch_utils.h"
2
#include "dispatch_utils.h"
3
8
-#include "quantization/w8a8/fp8/common.cuh"
4
+#include "../cuda_compat.h"
5
+
6
+#include "../quantization/w8a8/fp8/common.cuh"
7
#ifdef USE_ROCM
10
- #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
8
+ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
9
#else
12
- #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
10
+ #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
11
#endif
12
13
#ifdef USE_ROCM
@@ -164,43 +162,52 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
162
163
} // namespace vllm
164
167
-#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
168
- do { \
169
- VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
170
- using qk_t = scalar_t; \
171
- VLLM_DISPATCH_FLOATING_TYPES( \
172
- rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
173
- [&] { \
174
- using cos_sin_t = scalar_t; \
175
- if (rope_is_neox) { \
176
- vllm::concat_and_cache_mla_rope_fused_kernel< \
177
- qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
178
- <<<grid, block, 0, stream>>>( \
179
- positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
180
- k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
181
- rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
182
- q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
183
- kv_c_stride, num_q_heads, \
184
- reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
185
- slot_mapping.data_ptr<int64_t>(), block_stride, \
186
- entry_stride, kv_lora_rank, block_size, \
187
- kv_cache_quant_scale.data_ptr<float>()); \
188
- } else { \
189
- vllm::concat_and_cache_mla_rope_fused_kernel< \
190
- qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
191
- <<<grid, block, 0, stream>>>( \
192
- positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
193
- k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
194
- rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
195
- q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
196
- kv_c_stride, num_q_heads, \
197
- reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
198
- slot_mapping.data_ptr<int64_t>(), block_stride, \
199
- entry_stride, kv_lora_rank, block_size, \
200
- kv_cache_quant_scale.data_ptr<float>()); \
201
- } \
202
- }); \
203
- }); \
165
+#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
166
+ do { \
167
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
168
+ q_pe.scalar_type(), "qk_scalar_type", [&] { \
169
+ using qk_t = scalar_t; \
170
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
171
+ rope_cos_sin_cache.scalar_type(), \
172
+ "rope_cos_sin_cache_scalar_type", [&] { \
173
+ using cos_sin_t = scalar_t; \
174
+ if (rope_is_neox) { \
175
+ vllm::concat_and_cache_mla_rope_fused_kernel< \
176
+ qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
177
+ <<<grid, block, 0, stream>>>( \
178
+ positions.const_data_ptr<int64_t>(), \
179
+ q_pe.mutable_data_ptr<qk_t>(), \
180
+ k_pe.mutable_data_ptr<qk_t>(), \
181
+ kv_c.const_data_ptr<qk_t>(), \
182
+ rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \
183
+ rot_dim, q_pe_stride_token, q_pe_stride_head, \
184
+ k_pe_stride, kv_c_stride, num_q_heads, \
185
+ reinterpret_cast<CACHE_T*>( \
186
+ kv_cache.mutable_data_ptr()), \
187
+ slot_mapping.const_data_ptr<int64_t>(), \
188
+ block_stride, entry_stride, kv_lora_rank, \
189
+ block_size, \
190
+ kv_cache_quant_scale.const_data_ptr<float>()); \
191
+ } else { \
192
+ vllm::concat_and_cache_mla_rope_fused_kernel< \
193
+ qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
194
+ <<<grid, block, 0, stream>>>( \
195
+ positions.const_data_ptr<int64_t>(), \
196
+ q_pe.mutable_data_ptr<qk_t>(), \
197
+ k_pe.mutable_data_ptr<qk_t>(), \
198
+ kv_c.const_data_ptr<qk_t>(), \
199
+ rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \
200
+ rot_dim, q_pe_stride_token, q_pe_stride_head, \
201
+ k_pe_stride, kv_c_stride, num_q_heads, \
202
+ reinterpret_cast<CACHE_T*>( \
203
+ kv_cache.mutable_data_ptr()), \
204
+ slot_mapping.const_data_ptr<int64_t>(), \
205
+ block_stride, entry_stride, kv_lora_rank, \
206
+ block_size, \
207
+ kv_cache_quant_scale.const_data_ptr<float>()); \
208
+ } \
209
+ }); \
210
+ }); \
211
} while (false)
212
213
// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -208,64 +215,69 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
215
// Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
216
// concat_and_cache_mla.
217
void concat_and_cache_mla_rope_fused(
211
- torch::Tensor& positions, // [num_tokens]
212
- torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
213
- torch::Tensor& k_pe, // [num_tokens, rot_dim]
214
- torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
215
- torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
218
+ torch::stable::Tensor& positions, // [num_tokens]
219
+ torch::stable::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
220
+ torch::stable::Tensor& k_pe, // [num_tokens, rot_dim]
221
+ torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank]
222
+ torch::stable::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
223
bool rope_is_neox,
217
- torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
218
- torch::Tensor&
224
+ torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
225
+ torch::stable::Tensor&
226
kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
220
- const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
227
+ const std::string& kv_cache_dtype,
228
+ torch::stable::Tensor& kv_cache_quant_scale) {
229
// NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
230
// slot_mapping.size(0) because of padding for CUDA graphs.
223
- // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
224
- // both include padding.
225
- // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
226
- // since key includes padding for CUDA graphs, while slot_mapping does not.
227
- // In this case, slot_mapping.size(0) represents the actual number of tokens
231
+ // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0)
232
+ // because both include padding.
233
+ // In vLLM V1, however, key.size(0) can be larger than
234
+ // slot_mapping.size(0) since key includes padding for CUDA graphs,
235
+ // while slot_mapping does not. In this case,
236
+ // slot_mapping.size(0) represents the actual number of tokens
237
// before padding.
229
- // For compatibility with both cases, we use slot_mapping.size(0) as the
230
- // number of tokens.
231
- int num_tokens = slot_mapping.size(0);
232
- int num_padded_tokens = q_pe.size(0);
233
- TORCH_CHECK_GE(num_padded_tokens, num_tokens);
238
+ // For compatibility with both cases, we use slot_mapping.size(0) as
239
+ // the number of tokens.
240
+ const int64_t num_tokens = slot_mapping.size(0);
241
+ const int64_t num_padded_tokens = q_pe.size(0);
242
+ STD_TORCH_CHECK(num_padded_tokens >= num_tokens);
243
244
const int num_q_heads = q_pe.size(1);
245
const int rot_dim = q_pe.size(2);
246
const int kv_lora_rank = kv_c.size(1);
247
239
- TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
240
- TORCH_CHECK_EQ(positions.dim(), 1);
241
- TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
248
+ STD_TORCH_CHECK(positions.size(0) == num_padded_tokens);
249
+ STD_TORCH_CHECK(positions.dim() == 1);
250
+ STD_TORCH_CHECK(positions.scalar_type() ==
251
+ torch::headeronly::ScalarType::Long);
252
243
- TORCH_CHECK_EQ(q_pe.dim(), 3);
244
- TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
245
- TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
246
- TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
253
+ STD_TORCH_CHECK(q_pe.dim() == 3);
254
+ STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens);
255
+ STD_TORCH_CHECK(q_pe.size(1) == num_q_heads);
256
+ STD_TORCH_CHECK(q_pe.size(2) == rot_dim);
257
248
- TORCH_CHECK_EQ(k_pe.dim(), 2);
249
- TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
250
- TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
251
- TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
258
+ STD_TORCH_CHECK(k_pe.dim() == 2);
259
+ STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens);
260
+ STD_TORCH_CHECK(k_pe.size(1) == rot_dim);
261
+ STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type());
262
253
- TORCH_CHECK_EQ(kv_c.dim(), 2);
254
- TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
255
- TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
256
- TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
257
- TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
263
+ STD_TORCH_CHECK(kv_c.dim() == 2);
264
+ STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens);
265
+ STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank);
266
+ STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type());
267
259
- TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
268
+ STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim);
269
+ STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type());
270
261
- TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
262
- TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
271
+ STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens);
272
+ STD_TORCH_CHECK(slot_mapping.scalar_type() ==
273
+ torch::headeronly::ScalarType::Long);
274
264
- TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
265
- TORCH_CHECK_EQ(kv_cache.dim(), 3);
275
+ STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim);
276
+ STD_TORCH_CHECK(kv_cache.dim() == 3);
277
267
- TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
268
- TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
278
+ STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1);
279
+ STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() ==
280
+ torch::headeronly::ScalarType::Float);
281
282
int64_t q_pe_stride_token = q_pe.stride(0);
283
int64_t q_pe_stride_head = q_pe.stride(1);
@@ -286,9 +298,10 @@ void concat_and_cache_mla_rope_fused(
298
dim3 grid(num_tokens, 1, 1);
299
dim3 block(thread_block_size, 1, 1);
300
289
- const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
290
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
301
+ const torch::stable::accelerator::DeviceGuard device_guard(
302
+ positions.get_device_index());
303
+ const cudaStream_t stream = get_current_cuda_stream();
304
292
- DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
305
+ DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
306
CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
307
}
csrc/libtorch_stable/concat_mla_q.cuh
renamed
csrc/libtorch_stable/cuda_vec_utils.cuh
renamed
csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
renamed
+36
-34
@@ -17,11 +17,8 @@
17
#define NVFP4_ENABLE_ELTS16 1
18
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
19
20
-#include <ATen/cuda/CUDAContext.h>
21
-#include <c10/cuda/CUDAGuard.h>
22
-#include <torch/all.h>
23
-
24
-#include "dispatch_utils.h"
20
+#include "libtorch_stable/dispatch_utils.h"
21
+#include "libtorch_stable/torch_utils.h"
22
23
namespace vllm {
24
@@ -184,12 +181,13 @@ __global__ void reshape_and_cache_nvfp4_kernel(
181
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
182
// Each KV side contains both data and scale:
183
// page = [K_data | K_scale | V_data | V_scale]
187
-void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
188
- torch::Tensor& key_cache,
189
- torch::Tensor& value_cache,
190
- torch::Tensor& slot_mapping,
191
- torch::Tensor& k_scale,
192
- torch::Tensor& v_scale) {
184
+void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
185
+ torch::stable::Tensor& value,
186
+ torch::stable::Tensor& key_cache,
187
+ torch::stable::Tensor& value_cache,
188
+ torch::stable::Tensor& slot_mapping,
189
+ torch::stable::Tensor& k_scale,
190
+ torch::stable::Tensor& v_scale) {
191
int num_tokens = slot_mapping.size(0);
192
int num_heads = key.size(1);
193
int head_size = key.size(2);
@@ -200,17 +198,18 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
198
// key_cache is kv_cache[:, 0] with shape
199
// [num_blocks, block_size, num_heads, full_dim] in logical order.
200
// Strides encode the physical layout (HND or NHD).
203
- TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
204
- TORCH_CHECK(key_cache.size(3) == full_dim,
205
- "key_cache last dim must be data_dim + scale_dim, got ",
206
- key_cache.size(3), " expected ", full_dim);
201
+ STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
202
+ STD_TORCH_CHECK(key_cache.size(3) == full_dim,
203
+ "key_cache last dim must be data_dim + scale_dim, got ",
204
+ key_cache.size(3), " expected ", full_dim);
205
206
int block_size = key_cache.size(1);
207
210
- TORCH_CHECK(head_size % 16 == 0,
211
- "head_size must be divisible by 16 for NVFP4 KV cache");
212
- TORCH_CHECK(block_size % 4 == 0,
213
- "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
208
+ STD_TORCH_CHECK(head_size % 16 == 0,
209
+ "head_size must be divisible by 16 for NVFP4 KV cache");
210
+ STD_TORCH_CHECK(block_size % 4 == 0,
211
+ "block_size must be divisible by 4 for NVFP4 KV cache "
212
+ "swizzle");
213
214
// Detect physical layout from strides (based on full_dim).
215
// HND: head stride > block_offset stride.
@@ -230,8 +229,9 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
229
// Scale follows data within each KV side.
230
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
231
233
- uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
234
- uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
232
+ uint8_t* key_scale_ptr = key_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
233
+ uint8_t* value_scale_ptr =
234
+ value_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
235
236
// Scale strides: same page stride, inner strides from layout.
237
int64_t scale_block_stride = data_block_stride;
@@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
244
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
245
}
246
247
- const float* k_scale_ptr = k_scale.data_ptr<float>();
248
- const float* v_scale_ptr = v_scale.data_ptr<float>();
247
+ const float* k_scale_ptr = k_scale.const_data_ptr<float>();
248
+ const float* v_scale_ptr = v_scale.const_data_ptr<float>();
249
250
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
251
int total_groups = num_heads * groups_per_head;
@@ -256,20 +256,22 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
256
dim3 grid(num_tokens);
257
dim3 block(num_threads);
258
259
- const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
260
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
259
+ const torch::stable::accelerator::DeviceGuard device_guard(
260
+ key.get_device_index());
261
+ const cudaStream_t stream = get_current_cuda_stream();
262
262
- AT_DISPATCH_REDUCED_FLOATING_TYPES(
263
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
264
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
265
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
266
<<<grid, block, 0, stream>>>(
266
- key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
267
- key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
268
- key_scale_ptr, value_scale_ptr,
269
- slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
270
- key.stride(0), value.stride(0), num_heads, head_size,
271
- block_size, data_block_stride, data_head_stride,
272
- data_block_offset_stride, scale_block_stride, scale_head_stride,
273
- scale_block_offset_stride);
267
+ key.const_data_ptr<scalar_t>(),
268
+ value.const_data_ptr<scalar_t>(),
269
+ key_cache.mutable_data_ptr<uint8_t>(),
270
+ value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
271
+ value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
272
+ k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
273
+ num_heads, head_size, block_size, data_block_stride,
274
+ data_head_stride, data_block_offset_stride, scale_block_stride,
275
+ scale_head_stride, scale_block_offset_stride);
276
});
277
}
csrc/libtorch_stable/ops.h
+129
@@ -355,3 +355,132 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
355
int64_t tokens);
356
357
int64_t ggml_moe_get_block_size(int64_t type);
358
+
359
+void paged_attention_v1(
360
+ torch::stable::Tensor& out, torch::stable::Tensor& query,
361
+ torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
362
+ int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
363
+ torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
364
+ const std::optional<torch::stable::Tensor>& alibi_slopes,
365
+ const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
366
+ torch::stable::Tensor& v_scale, const int64_t tp_rank,
367
+ const int64_t blocksparse_local_blocks,
368
+ const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
369
+ const int64_t blocksparse_head_sliding_step);
370
+
371
+void paged_attention_v2(
372
+ torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
373
+ torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
374
+ torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
375
+ torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
376
+ torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
377
+ int64_t block_size, int64_t max_seq_len,
378
+ const std::optional<torch::stable::Tensor>& alibi_slopes,
379
+ const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
380
+ torch::stable::Tensor& v_scale, const int64_t tp_rank,
381
+ const int64_t blocksparse_local_blocks,
382
+ const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
383
+ const int64_t blocksparse_head_sliding_step);
384
+
385
+// Cache ops (shared CUDA/ROCm)
386
+void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
387
+ int64_t block_size_in_bytes,
388
+ const torch::stable::Tensor& block_mapping);
389
+
390
+// Batch swap: submit all block copies in a single driver call.
391
+void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
392
+ const torch::stable::Tensor& dst_ptrs,
393
+ const torch::stable::Tensor& sizes,
394
+ bool is_src_access_order_any);
395
+
396
+void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
397
+ torch::stable::Tensor& key_cache,
398
+ torch::stable::Tensor& value_cache,
399
+ torch::stable::Tensor& slot_mapping,
400
+ const std::string& kv_cache_dtype,
401
+ torch::stable::Tensor& k_scale,
402
+ torch::stable::Tensor& v_scale);
403
+
404
+void reshape_and_cache_flash(
405
+ torch::stable::Tensor& key, torch::stable::Tensor& value,
406
+ torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
407
+ torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
408
+ torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
409
+
410
+void concat_and_cache_mla(torch::stable::Tensor& kv_c,
411
+ torch::stable::Tensor& k_pe,
412
+ torch::stable::Tensor& kv_cache,
413
+ torch::stable::Tensor& slot_mapping,
414
+ const std::string& kv_cache_dtype,
415
+ torch::stable::Tensor& scale);
416
+
417
+// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
418
+void concat_and_cache_mla_rope_fused(
419
+ torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
420
+ torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
421
+ torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
422
+ torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
423
+ const std::string& kv_cache_dtype,
424
+ torch::stable::Tensor& kv_cache_quant_scale);
425
+
426
+// Just for unittest
427
+void convert_fp8(torch::stable::Tensor& dst_cache,
428
+ torch::stable::Tensor& src_cache, const double scale,
429
+ const std::string& kv_cache_dtype);
430
+
431
+void gather_and_maybe_dequant_cache(
432
+ torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
433
+ // ENTRIES...]
434
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
435
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
436
+ torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
437
+ torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
438
+ int64_t num_tokens, const std::string& kv_cache_dtype,
439
+ torch::stable::Tensor const& scale,
440
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
441
+
442
+// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
443
+void cp_gather_cache(
444
+ torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
445
+ // ENTRIES...]
446
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
447
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
448
+ torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
449
+ int64_t batch_size,
450
+ std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
451
+
452
+// Gather and upconvert FP8 KV cache to BF16 workspace
453
+void cp_gather_and_upconvert_fp8_kv_cache(
454
+ torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
455
+ // 656]
456
+ torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
457
+ torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
458
+ torch::stable::Tensor const& seq_lens, // [BATCH]
459
+ torch::stable::Tensor const& workspace_starts, // [BATCH]
460
+ int64_t batch_size);
461
+
462
+// Indexer K quantization and cache function
463
+void indexer_k_quant_and_cache(
464
+ torch::stable::Tensor& k, // [num_tokens, head_dim]
465
+ torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
466
+ // cache_stride]
467
+ torch::stable::Tensor& slot_mapping, // [num_tokens]
468
+ int64_t quant_block_size, // quantization block size
469
+ const std::string& scale_fmt);
470
+
471
+// Concatenate query nope and rope for MLA/DSA attention
472
+void concat_mla_q(
473
+ torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
474
+ torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
475
+ torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim +
476
+ // rope_dim]
477
+
478
+// Extract function to gather quantized K cache
479
+void cp_gather_indexer_k_quant_cache(
480
+ const torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
481
+ // cache_stride]
482
+ torch::stable::Tensor& dst_k, // [num_tokens, head_dim]
483
+ torch::stable::Tensor& dst_scale, // [num_tokens, head_dim /
484
+ // quant_block_size * 4]
485
+ const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
486
+ const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1]
csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+1
-1
@@ -17,7 +17,7 @@
17
#include <torch/csrc/stable/tensor.h>
18
#include "libtorch_stable/torch_utils.h"
19
#include "libtorch_stable/dispatch_utils.h"
20
-#include "cuda_vec_utils.cuh"
20
+#include "../../cuda_vec_utils.cuh"
21
22
#include <cuda_runtime_api.h>
23
#include <cuda_runtime.h>
csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu
+1
-1
@@ -27,7 +27,7 @@
27
#include <torch/csrc/stable/tensor.h>
28
#include "libtorch_stable/torch_utils.h"
29
#include "libtorch_stable/dispatch_utils.h"
30
-#include "cuda_vec_utils.cuh"
30
+#include "../../cuda_vec_utils.cuh"
31
#include "cuda_utils.h"
32
33
#include "nvfp4_utils.cuh"
csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
+1
-1
@@ -17,7 +17,7 @@
17
#include <torch/csrc/stable/tensor.h>
18
#include "libtorch_stable/torch_utils.h"
19
#include "libtorch_stable/dispatch_utils.h"
20
-#include "cuda_vec_utils.cuh"
20
+#include "../../cuda_vec_utils.cuh"
21
22
#include <cuda_runtime_api.h>
23
#include <cuda_runtime.h>
csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu
+1
-1
@@ -23,7 +23,7 @@
23
24
#include "libtorch_stable/torch_utils.h"
25
#include "libtorch_stable/dispatch_utils.h"
26
-#include "cuda_vec_utils.cuh"
26
+#include "../../cuda_vec_utils.cuh"
27
28
#include "cuda_utils.h"
29
#include "launch_bounds_utils.h"
csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh
+1
-1
@@ -20,7 +20,7 @@
20
#include <cuda_fp8.h>
21
#include <utility>
22
23
-#include "cuda_vec_utils.cuh"
23
+#include "../../cuda_vec_utils.cuh"
24
25
#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
26
CUDA_VERSION >= 12090
csrc/libtorch_stable/torch_bindings.cpp
+141
@@ -474,6 +474,33 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
474
"Tensor? initial_state_idx,"
475
"Tensor? cu_chunk_seqlen,"
476
"Tensor? last_chunk_indices) -> ()");
477
+
478
+ // Attention ops
479
+ // Compute the attention between an input query and the cached
480
+ // keys/values using PagedAttention.
481
+ ops.def(
482
+ "paged_attention_v1("
483
+ " Tensor! out, Tensor query, Tensor key_cache,"
484
+ " Tensor value_cache, int num_kv_heads, float scale,"
485
+ " Tensor block_tables, Tensor seq_lens, int block_size,"
486
+ " int max_seq_len, Tensor? alibi_slopes,"
487
+ " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
488
+ " int tp_rank, int blocksparse_local_blocks,"
489
+ " int blocksparse_vert_stride, int blocksparse_block_size,"
490
+ " int blocksparse_head_sliding_step) -> ()");
491
+
492
+ // PagedAttention V2.
493
+ ops.def(
494
+ "paged_attention_v2("
495
+ " Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
496
+ " Tensor! tmp_out, Tensor query, Tensor key_cache,"
497
+ " Tensor value_cache, int num_kv_heads, float scale,"
498
+ " Tensor block_tables, Tensor seq_lens, int block_size,"
499
+ " int max_seq_len, Tensor? alibi_slopes,"
500
+ " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
501
+ " int tp_rank, int blocksparse_local_blocks,"
502
+ " int blocksparse_vert_stride, int blocksparse_block_size,"
503
+ " int blocksparse_head_sliding_step) -> ()");
504
}
505
506
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -581,6 +608,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
608
ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
609
ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
610
ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
611
+
612
+ ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
613
+ ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
614
}
615
616
// These capability-check functions take only primitive args (no tensors), so
@@ -603,4 +633,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
633
ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
634
}
635
636
+// Cache ops
637
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
638
+ // Swap in (out) the cache blocks from src to dst.
639
+ ops.def(
640
+ "swap_blocks(Tensor src, Tensor! dst,"
641
+ " int block_size_in_bytes, Tensor block_mapping) -> ()");
642
+
643
+ // Batch swap: submit all block copies in a single driver call.
644
+ ops.def(
645
+ "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
646
+ " Tensor sizes,"
647
+ " bool is_src_access_order_any=False) -> ()");
648
+
649
+ // Reshape the key and value tensors and cache them.
650
+ ops.def(
651
+ "reshape_and_cache(Tensor key, Tensor value,"
652
+ " Tensor! key_cache, Tensor! value_cache,"
653
+ " Tensor slot_mapping,"
654
+ " str kv_cache_dtype,"
655
+ " Tensor k_scale, Tensor v_scale) -> ()");
656
+
657
+ // Reshape the key and value tensors and cache them.
658
+ ops.def(
659
+ "reshape_and_cache_flash(Tensor key, Tensor value,"
660
+ " Tensor! key_cache,"
661
+ " Tensor! value_cache,"
662
+ " Tensor slot_mapping,"
663
+ " str kv_cache_dtype,"
664
+ " Tensor k_scale, Tensor v_scale) -> ()");
665
+
666
+ // Concat kv_c and k_pe and cache them.
667
+ ops.def(
668
+ "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
669
+ " Tensor! kv_cache,"
670
+ " Tensor slot_mapping,"
671
+ " str kv_cache_dtype,"
672
+ " Tensor scale) -> ()");
673
+
674
+ // Rotate Q and K, then write to kv cache for MLA
675
+ ops.def(
676
+ "concat_and_cache_mla_rope_fused("
677
+ " Tensor positions,"
678
+ " Tensor! q_pe,"
679
+ " Tensor! k_pe,"
680
+ " Tensor kv_c,"
681
+ " Tensor cos_sin_cache,"
682
+ " bool is_neox,"
683
+ " Tensor slot_mapping,"
684
+ " Tensor! kv_cache,"
685
+ " str kv_cache_dtype,"
686
+ " Tensor kv_cache_scale) -> ()");
687
+
688
+ // Convert the key and value cache to fp8 data type.
689
+ ops.def(
690
+ "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
691
+ "str kv_cache_dtype) -> ()");
692
+
693
+ // Gather cache blocks from src_cache to dst, dequantizing from
694
+ // src_cache's dtype to dst's dtype if necessary.
695
+ ops.def(
696
+ "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
697
+ " Tensor block_table, Tensor cu_seq_lens, "
698
+ " Tensor token_to_seq, "
699
+ " int num_tokens, "
700
+ " str kv_cache_dtype, "
701
+ " Tensor scale, Tensor? seq_starts) -> ()");
702
+
703
+ ops.def(
704
+ "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
705
+ "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
706
+
707
+ ops.def(
708
+ "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
709
+ "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
710
+ "batch_size) -> ()");
711
+
712
+ ops.def(
713
+ "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
714
+ "slot_mapping, "
715
+ "int quant_block_size, str kv_cache_dtype) -> ()");
716
+
717
+ ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
718
+
719
+ ops.def(
720
+ "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
721
+ "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
722
+}
723
+
724
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
725
+ ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
726
+}
727
+
728
+STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
729
+ ops.impl("swap_blocks", TORCH_BOX(&swap_blocks));
730
+ ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
731
+ ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
732
+ ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
733
+ ops.impl("concat_and_cache_mla_rope_fused",
734
+ TORCH_BOX(&concat_and_cache_mla_rope_fused));
735
+ ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
736
+ ops.impl("gather_and_maybe_dequant_cache",
737
+ TORCH_BOX(&gather_and_maybe_dequant_cache));
738
+ ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
739
+ ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
740
+ TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
741
+ ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
742
+ ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
743
+ ops.impl("cp_gather_indexer_k_quant_cache",
744
+ TORCH_BOX(&cp_gather_indexer_k_quant_cache));
745
+}
746
+
747
REGISTER_EXTENSION(_C_stable_libtorch)
csrc/ops.h
-23
@@ -31,29 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
31
return new_tensor;
32
}
33
34
-void paged_attention_v1(
35
- torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
36
- torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
37
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
38
- int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
39
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
40
- torch::Tensor& v_scale, const int64_t tp_rank,
41
- const int64_t blocksparse_local_blocks,
42
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
43
- const int64_t blocksparse_head_sliding_step);
44
-
45
-void paged_attention_v2(
46
- torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
47
- torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
48
- torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
49
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
50
- int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
51
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
52
- torch::Tensor& v_scale, const int64_t tp_rank,
53
- const int64_t blocksparse_local_blocks,
54
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
55
- const int64_t blocksparse_head_sliding_step);
56
-
34
// rms_norm and fused_add_rms_norm declarations also exist in
35
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
36
// because the CPU build still uses these torch::Tensor declarations.
csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
+12
-9
@@ -6,6 +6,7 @@
6
#include <hip/hip_bfloat16.h>
7
8
#include "../../../../attention/attention_dtypes.h"
9
+#include <torch/headeronly/core/ScalarType.h>
10
11
namespace vllm {
12
#ifdef USE_ROCM
@@ -642,27 +643,29 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
643
vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
644
vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
645
if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
645
- if (SRC_DTYPE == at::ScalarType::Float) { \
646
+ if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
647
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
647
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
648
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
649
FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \
649
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
650
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
651
FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \
652
} else { \
652
- TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
653
+ STD_TORCH_CHECK(false, \
654
+ "Unsupported input type of kv cache: ", SRC_DTYPE); \
655
} \
656
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
655
- if (SRC_DTYPE == at::ScalarType::Float) { \
657
+ if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
658
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
657
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
659
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
660
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
659
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
661
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
662
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
663
} else { \
662
- TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
664
+ STD_TORCH_CHECK(false, \
665
+ "Unsupported input type of kv cache: ", SRC_DTYPE); \
666
} \
667
} else { \
665
- TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
668
+ STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
669
}
670
671
} // namespace fp8
csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
+17
-13
@@ -1,6 +1,7 @@
1
#pragma once
2
3
#include "../../../../attention/attention_dtypes.h"
4
+#include <torch/headeronly/core/ScalarType.h>
5
#include <assert.h>
6
#include <float.h>
7
#include <stdint.h>
@@ -546,37 +547,40 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
547
vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \
548
vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \
549
if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \
549
- if (SRC_DTYPE == at::ScalarType::Float) { \
550
+ if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
551
FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \
551
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
552
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
553
FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \
553
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
554
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
555
FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \
556
} else { \
556
- TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
557
+ STD_TORCH_CHECK(false, \
558
+ "Unsupported input type of kv cache: ", SRC_DTYPE); \
559
} \
560
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \
559
- if (SRC_DTYPE == at::ScalarType::Float) { \
561
+ if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
562
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
561
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
563
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
564
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
563
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
565
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
566
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \
567
} else { \
566
- TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
568
+ STD_TORCH_CHECK(false, \
569
+ "Unsupported input type of kv cache: ", SRC_DTYPE); \
570
} \
571
} else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \
569
- if (SRC_DTYPE == at::ScalarType::Float) { \
572
+ if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \
573
FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
571
- } else if (SRC_DTYPE == at::ScalarType::Half) { \
574
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \
575
FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
573
- } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \
576
+ } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \
577
FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \
578
} else { \
576
- TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
579
+ STD_TORCH_CHECK(false, \
580
+ "Unsupported input type of kv cache: ", SRC_DTYPE); \
581
} \
582
} else { \
579
- TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
583
+ STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
584
}
585
586
} // namespace fp8
csrc/torch_bindings.cpp
+4
-138
@@ -1,4 +1,7 @@
1
-#include "cache.h"
1
+// Provides torch::Tensor for ops.h (previously included transitively via
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"
6
#include "ops.h"
7
#include "core/registration.h"
@@ -33,35 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
36
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
37
&get_cuda_view_from_cpu_tensor);
38
36
- // Attention ops
37
- // Compute the attention between an input query and the cached
38
- // keys/values using PagedAttention.
39
- ops.def(
40
- "paged_attention_v1("
41
- " Tensor! out, Tensor query, Tensor key_cache,"
42
- " Tensor value_cache, int num_kv_heads, float scale,"
43
- " Tensor block_tables, Tensor seq_lens, int block_size,"
44
- " int max_seq_len, Tensor? alibi_slopes,"
45
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
46
- " int tp_rank, int blocksparse_local_blocks,"
47
- " int blocksparse_vert_stride, int blocksparse_block_size,"
48
- " int blocksparse_head_sliding_step) -> ()");
49
- ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
50
-
51
- // PagedAttention V2.
52
- ops.def(
53
- "paged_attention_v2("
54
- " Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
55
- " Tensor! tmp_out, Tensor query, Tensor key_cache,"
56
- " Tensor value_cache, int num_kv_heads, float scale,"
57
- " Tensor block_tables, Tensor seq_lens, int block_size,"
58
- " int max_seq_len, Tensor? alibi_slopes,"
59
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
60
- " int tp_rank, int blocksparse_local_blocks,"
61
- " int blocksparse_vert_stride, int blocksparse_block_size,"
62
- " int blocksparse_head_sliding_step) -> ()");
63
- ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
64
-
39
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
40
ops.def(
41
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -217,114 +191,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
191
#endif
192
}
193
220
-TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
221
- // Cache ops
222
- // Swap in (out) the cache blocks from src to dst.
223
- cache_ops.def(
224
- "swap_blocks(Tensor src, Tensor! dst,"
225
- " int block_size_in_bytes, Tensor block_mapping) -> ()");
226
- cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
227
-
228
- // Batch swap: submit all block copies in a single driver call.
229
- cache_ops.def(
230
- "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
231
- " Tensor sizes,"
232
- " bool is_src_access_order_any=False) -> ()");
233
- cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
234
-
235
- // Reshape the key and value tensors and cache them.
236
- cache_ops.def(
237
- "reshape_and_cache(Tensor key, Tensor value,"
238
- " Tensor! key_cache, Tensor! value_cache,"
239
- " Tensor slot_mapping,"
240
- " str kv_cache_dtype,"
241
- " Tensor k_scale, Tensor v_scale) -> ()");
242
- cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);
243
-
244
- // Reshape the key and value tensors and cache them.
245
- cache_ops.def(
246
- "reshape_and_cache_flash(Tensor key, Tensor value,"
247
- " Tensor! key_cache,"
248
- " Tensor! value_cache,"
249
- " Tensor slot_mapping,"
250
- " str kv_cache_dtype,"
251
- " Tensor k_scale, Tensor v_scale) -> ()");
252
- cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
253
- &reshape_and_cache_flash);
254
-
255
- // Concat kv_c and k_pe and cache them.
256
- cache_ops.def(
257
- "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
258
- " Tensor! kv_cache,"
259
- " Tensor slot_mapping,"
260
- " str kv_cache_dtype,"
261
- " Tensor scale) -> ()");
262
- cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
263
-
264
- // Rotate Q and K, then write to kv cache for MLA
265
- cache_ops.def(
266
- "concat_and_cache_mla_rope_fused("
267
- " Tensor positions,"
268
- " Tensor! q_pe,"
269
- " Tensor! k_pe,"
270
- " Tensor kv_c,"
271
- " Tensor cos_sin_cache,"
272
- " bool is_neox,"
273
- " Tensor slot_mapping,"
274
- " Tensor! kv_cache,"
275
- " str kv_cache_dtype,"
276
- " Tensor kv_cache_scale) -> ()");
277
- cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
278
- &concat_and_cache_mla_rope_fused);
279
-
280
- // Convert the key and value cache to fp8 data type.
281
- cache_ops.def(
282
- "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
283
- "str kv_cache_dtype) -> ()");
284
- cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
285
-
286
- // Gather cache blocks from src_cache to dst, dequantizing from
287
- // src_cache's dtype to dst's dtype if necessary.
288
- cache_ops.def(
289
- "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
290
- " Tensor block_table, Tensor cu_seq_lens, "
291
- " Tensor token_to_seq, "
292
- " int num_tokens, "
293
- " str kv_cache_dtype, "
294
- " Tensor scale, Tensor? seq_starts) -> ()");
295
- cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
296
- &gather_and_maybe_dequant_cache);
297
-
298
- cache_ops.def(
299
- "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
300
- "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
301
- cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
302
-
303
- cache_ops.def(
304
- "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
305
- "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
306
- "batch_size) -> ()");
307
- cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,
308
- &cp_gather_and_upconvert_fp8_kv_cache);
309
-
310
- cache_ops.def(
311
- "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
312
- "slot_mapping, "
313
- "int quant_block_size, str kv_cache_dtype) -> ()");
314
- cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
315
- &indexer_k_quant_and_cache);
316
-
317
- cache_ops.def(
318
- "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
319
- cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
320
-
321
- cache_ops.def(
322
- "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
323
- "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
324
- cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
325
- &cp_gather_indexer_k_quant_cache);
326
-}
327
-
194
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
195
// Cuda utils
196